feat: Semantic splitter (#63)

* feat: Semantic Spliter + minor improvements * chore: testing * fix: BaseDocumentChunk output fix * feat: Combined chunk title with chunk content * chore: Updated semantic router version * chore: Updated semantic router version * Small tweaks and bug fixes * Update README * chore: Merging * chore: Testing * chore: Minor improvements * Add support for querying code interpreter (#66) * Add support for queryig code interpreter * Fix formatting * Ensure the sandbox close is called on exceptions * Update service/code_interpreter.py Co-authored-by: Tomas Valenta <[email protected]> * Update service/code_interpreter.py Co-authored-by: Tomas Valenta <[email protected]> * Update service/router.py Co-authored-by: Tomas Valenta <[email protected]> * Update service/code_interpreter.py Co-authored-by: Tomas Valenta <[email protected]> * Add system prompt * Format code * Bump dependencies * Minor tweaks --------- Co-authored-by: Tomas Valenta <[email protected]> * Minor tweaks --------- Co-authored-by: Ismail Pelaseyed <[email protected]> Co-authored-by: Tomas Valenta <[email protected]>
superagent-ai · Mar 2, 2024 · ae3b113 · ae3b113
1 parent db5c57a
commit ae3b113
Show file tree

Hide file tree

Showing 22 changed files with 803 additions and 272 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -0,0 +1,6 @@
+{
+    "cSpell.words": [
+        "tiktoken",
+        "Upserted"
+    ]
+}
diff --git a/README.md b/README.md
@@ -66,23 +66,41 @@ Super-Rag comes with a built in REST API powered by FastApi.
 
 // Payload
 {
-    "files": [{
-        "url": "https://arxiv.org/pdf/2210.03629.pdf"
-    }],
+    "files": [
+        {
+            "name": "My file", // Optional
+            "url": "https://path-to-my-file.pdf"
+        }
+    ],
+    "document_processor": { // Optional
+        "encoder": {
+            "dimensions": 384,
+            "model_name": "embed-multilingual-light-v3.0",
+            "provider": "cohere"
+        },
+        "unstructured": {
+            "hi_res_model_name": "detectron2_onnx",
+            "partition_strategy": "auto",
+            "process_tables": false
+        },
+        "splitter": {
+            "max_tokens": 400,
+            "min_tokens": 30,
+            "name": "semantic",
+            "prefix_summary": true,
+            "prefix_title": true,
+            "rolling_window_size": 1
+        }
+    },
     "vector_database": {
         "type": "qdrant",
         "config": {
             "api_key": "YOUR API KEY",
             "host": "THE QDRANT HOST"
         }
     },
-    "encoder": {
-        "type": "openai",
-        "name": "text-embedding-3-small",
-        "dimensions": 1536  // encoder depends on the provider and model
-    },
-    "index_name": "YOUR INDEX",
-    "webhook_url": "https://webhook.site/0e217d1c-49f1-424a-9992-497db09f7793"
+    "index_name": "my_index",
+    "webhook_url": "https://my-webhook-url"
 }
 ```
 
@@ -103,12 +121,13 @@ Super-Rag comes with a built in REST API powered by FastApi.
     "index_name": "YOUR INDEX",
     "interpreter_mode": true,
     "encoder": {
-        "type": "cohere",
-        "name": "embed-multilingual-light-v3.0",
+        "provider": "openai",
+        "name": "text-embedding-3-small",
         "dimensions": 384
     },
-    "exclude_fields": ["metadata"],
-    "session_id": "test"
+    "exclude_fields": ["metadata"], // Exclude specific fields
+    "interpreter_mode": False, // Set to True if you wish to run computation Q&A with a code interpreter
+    "session_id": "my_session_id" // keeps micro-vm sessions and enables caching 
 }
 ```
 

diff --git a/api/delete.py b/api/delete.py
@@ -1,7 +1,6 @@
 from fastapi import APIRouter
 
 from models.delete import RequestPayload, ResponsePayload
-from service.embedding import get_encoder
 from vectordbs import get_vector_service
 from vectordbs.base import BaseVectorDatabase
 
@@ -10,12 +9,12 @@
 
 @router.delete("/delete", response_model=ResponsePayload)
 async def delete(payload: RequestPayload):
-    encoder = get_encoder(encoder_config=payload.encoder)
+    encoder = payload.encoder.get_encoder()
     vector_service: BaseVectorDatabase = get_vector_service(
         index_name=payload.index_name,
         credentials=payload.vector_database,
         encoder=encoder,
-        dimensions=encoder.dimensions,
+        dimensions=payload.encoder.dimensions,
     )
 
     for file in payload.files:

diff --git a/api/ingest.py b/api/ingest.py
@@ -5,7 +5,7 @@
 from fastapi import APIRouter
 
 from models.ingest import RequestPayload
-from service.embedding import EmbeddingService, get_encoder
+from service.embedding import EmbeddingService
 from service.ingest import handle_google_drive, handle_urls
 from utils.summarise import SUMMARY_SUFFIX
 
@@ -14,25 +14,33 @@
 
 @router.post("/ingest")
 async def ingest(payload: RequestPayload) -> Dict:
-    encoder = get_encoder(encoder_config=payload.encoder)
+    encoder = payload.document_processor.encoder.get_encoder()
     embedding_service = EmbeddingService(
+        encoder=encoder,
         index_name=payload.index_name,
         vector_credentials=payload.vector_database,
-        dimensions=payload.encoder.dimensions,
+        dimensions=payload.document_processor.encoder.dimensions,
     )
+    chunks = []
+    summary_documents = []
     if payload.files:
-        chunks, summary_documents = await handle_urls(embedding_service, payload.files)
+        chunks, summary_documents = await handle_urls(
+            embedding_service=embedding_service,
+            files=payload.files,
+            config=payload.document_processor,
+        )
+
     elif payload.google_drive:
         chunks, summary_documents = await handle_google_drive(
             embedding_service, payload.google_drive
-        )
+        )  # type: ignore TODO: Fix typing
 
     await asyncio.gather(
-        embedding_service.generate_and_upsert_embeddings(
-            documents=chunks, encoder=encoder, index_name=payload.index_name
+        embedding_service.embed_and_upsert(
+            chunks=chunks, encoder=encoder, index_name=payload.index_name
         ),
-        embedding_service.generate_and_upsert_embeddings(
-            documents=summary_documents,
+        embedding_service.embed_and_upsert(
+            chunks=summary_documents,
             encoder=encoder,
             index_name=f"{payload.index_name}{SUMMARY_SUFFIX}",
         ),

diff --git a/api/query.py b/api/query.py
@@ -1,6 +1,6 @@
 from fastapi import APIRouter
 
-from models.query import RequestPayload, ResponseData, ResponsePayload
+from models.query import RequestPayload, ResponsePayload
 from service.router import query as _query
 
 router = APIRouter()
@@ -9,5 +9,9 @@
 @router.post("/query", response_model=ResponsePayload)
 async def query(payload: RequestPayload):
     chunks = await _query(payload=payload)
-    response_data = [ResponseData(**chunk.model_dump()) for chunk in chunks]
-    return {"success": True, "data": response_data}
+    # NOTE: Filter out fields before given to LLM
+    response_payload = ResponsePayload(success=True, data=chunks)
+    response_data = response_payload.model_dump(
+        exclude=set(payload.exclude_fields) if payload.exclude_fields else None
+    )
+    return response_data
diff --git a/dev/embedding.ipynb b/dev/embedding.ipynb
@@ -40,7 +40,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "elements = await embedding_service._download_and_extract_elements(file, strategy=\"auto\")\n"
+    "elements = await embedding_service._partition_file(file, strategy=\"auto\")\n"
    ]
   },
   {