Add embeddings to doc get/list

julep-ai · Oct 3, 2024 · 70eebc4 · 70eebc4
1 parent 2d62857
commit 70eebc4
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 2 deletions.
diff --git a/agents-api/agents_api/autogen/Docs.py b/agents-api/agents_api/autogen/Docs.py
@@ -57,6 +57,13 @@ class Doc(BaseModel):
     """
     Contents of the document
     """
+    embeddings: Annotated[
+        list[float] | list[list[float]] | None,
+        Field(None, json_schema_extra={"readOnly": True}),
+    ]
+    """
+    Embeddings for the document
+    """
 
 
 class DocOwner(BaseModel):

diff --git a/agents-api/agents_api/models/docs/get_doc.py b/agents-api/agents_api/models/docs/get_doc.py
@@ -37,6 +37,7 @@
     one=True,
     transform=lambda d: {
         "content": [s[1] for s in sorted(d["snippet_data"], key=lambda x: x[0])],
+        "embeddings": [s[2] for s in sorted(d["snippet_data"], key=lambda x: x[0])],
         **d,
     },
 )
@@ -68,8 +69,9 @@ def get_doc(
                 doc_id,
                 index,
                 content,
+                embedding,
             },
-            snippet_data = [index, content]
+            snippet_data = [index, content, embedding]
 
         ?[
             id,

diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py
@@ -34,6 +34,7 @@
     Doc,
     transform=lambda d: {
         "content": [s[1] for s in sorted(d["snippet_data"], key=lambda x: x[0])],
+        "embeddings": [s[2] for s in sorted(d["snippet_data"], key=lambda x: x[0])],
         **d,
     },
 )
@@ -67,8 +68,9 @@ def list_docs(
                 doc_id: id,
                 index,
                 content,
+                embedding,
             }},
-            snippet_data = [index, content]
+            snippet_data = [index, content, embedding]
 
         ?[
             owner_type,

diff --git a/cozo.db b/cozo.db
diff --git a/typespec/docs/models.tsp b/typespec/docs/models.tsp
@@ -23,6 +23,10 @@ model Doc {
 
     /** Contents of the document */
     content: string | string[];
+
+    /** Embeddings for the document */
+    @visibility("read")
+    embeddings?: float32[] | float32[][];
 }
 
 /** Payload for creating a doc */