Add support for semantic-router encoders (#45)

* Add support for semantic-router encoders * Fix formatting
superagent-ai · Feb 15, 2024 · 3f20dd8 · 3f20dd8
1 parent 18adba6
commit 3f20dd8
Show file tree

Hide file tree

Showing 20 changed files with 67 additions and 345 deletions.
diff --git a/README.md b/README.md
@@ -23,7 +23,11 @@ Input example:
         }
     },
     "index_name": "my_index",
-    "encoder": "my_encoder"
+    "encoder": {
+        "type": "openai",
+        "name": "text-embedding-3-small",
+        "dimensions": 1536  # encoder depends on the provider and model
+    },
     "webhook_url": "https://my-webhook-url"
 }
 ```
@@ -42,7 +46,10 @@ Input example:
         }
     },
     "index_name": "my_index",
-    "encoder": "my_encoder",
+    "encoder": {
+        "type": "openai",
+        "name": "text-embedding-3-small",
+    } 
 }
 ```
 

diff --git a/api/delete.py b/api/delete.py
@@ -10,11 +10,12 @@
 
 @router.delete("/delete", response_model=ResponsePayload)
 async def delete(payload: RequestPayload):
-    encoder = get_encoder(encoder_type=payload.encoder)
+    encoder = get_encoder(encoder_config=payload.encoder)
     vector_service: BaseVectorDatabase = get_vector_service(
         index_name=payload.index_name,
         credentials=payload.vector_database,
         encoder=encoder,
+        dimensions=encoder.dimensions,
     )
     data = await vector_service.delete(file_url=payload.file_url)
     return ResponsePayload(success=True, data=data)
diff --git a/api/ingest.py b/api/ingest.py
@@ -17,9 +17,10 @@ async def ingest(payload: RequestPayload) -> Dict:
         files=payload.files,
         index_name=payload.index_name,
         vector_credentials=payload.vector_database,
+        dimensions=payload.encoder.dimensions,
     )
     chunks = await embedding_service.generate_chunks()
-    encoder = get_encoder(encoder_type=payload.encoder)
+    encoder = get_encoder(encoder_config=payload.encoder)
     summary_documents = await embedding_service.generate_summary_documents(
         documents=chunks
     )

diff --git a/encoders/__init__.py b/encoders/__init__.py
diff --git a/encoders/base.py b/encoders/base.py
diff --git a/encoders/bm25.py b/encoders/bm25.py
diff --git a/encoders/cohere.py b/encoders/cohere.py
diff --git a/encoders/huggingface.py b/encoders/huggingface.py
diff --git a/encoders/openai.py b/encoders/openai.py
diff --git a/models/delete.py b/models/delete.py
@@ -1,14 +1,14 @@
 from pydantic import BaseModel
 
-from models.ingest import EncoderEnum
+from models.ingest import Encoder
 from models.vector_database import VectorDatabase
 
 
 class RequestPayload(BaseModel):
     index_name: str
     file_url: str
     vector_database: VectorDatabase
-    encoder: EncoderEnum
+    encoder: Encoder
 
 
 class DeleteResponse(BaseModel):