fix: Query fixes for repeated docs

julep-ai · Sep 19, 2024 · 7eddbad · 7eddbad
1 parent d52bce7
commit 7eddbad
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/agents-api/agents_api/clients/litellm.py b/agents-api/agents_api/clients/litellm.py
@@ -14,7 +14,7 @@
 @wraps(_acompletion)
 @beartype
 async def acompletion(
-    *, model: str, messages: list[dict], custom_api_key: str = None, **kwargs
+    *, model: str, messages: list[dict], custom_api_key: None | str = None, **kwargs
 ) -> ModelResponse | CustomStreamWrapper:
 
     supported_params = get_supported_openai_params(model)

diff --git a/agents-api/agents_api/models/docs/search_docs_by_embedding.py b/agents-api/agents_api/models/docs/search_docs_by_embedding.py
@@ -194,11 +194,11 @@ def search_docs_by_embedding(
     """
 
     collect_query = """
-        m[
+        n[
             doc_id,
             owner_type,
             owner_id,
-            collect(snippet),
+            unique(snippet_data),
             distance,
             title,
         ] := 
@@ -209,11 +209,30 @@ def search_docs_by_embedding(
                 snippet_data,
                 distance,
                 title,
-            }, snippet = {
-                "index": snippet_data->0,
-                "content": snippet_data->1,
             }
 
+        m[
+            doc_id,
+            owner_type,
+            owner_id,
+            collect(snippet),
+            distance,
+            title,
+        ] := 
+            n[
+                doc_id,
+                owner_type,
+                owner_id,
+                snippet_data,
+                distance,
+                title,
+            ],
+            snippet = {
+                "index": snippet_datum->0,
+                "content": snippet_datum->1
+            },
+            snippet_datum in snippet_data
+
         ?[
             id,
             owner_type,