Fix docs retreival/search errors

agents-api/agents_api/queries/docs/list_docs.py

@@ -10,6 +10,7 @@
 from beartype import beartype
 from fastapi import HTTPException
 from sqlglot import parse_one
+import ast
 
 from ...autogen.openapi_model import Doc
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
@@ -58,6 +59,13 @@ def transform_list_docs(d: dict) -> dict:
     content = d["content"][0] if len(d["content"]) == 1 else d["content"]
 
     embeddings = d["embeddings"][0] if len(d["embeddings"]) == 1 else d["embeddings"]
+
+    try:
+        # Embeddings are retreived as a string, so we need to evaluate it
+        embeddings = ast.literal_eval(embeddings)
+    except Exception as e:
+        raise ValueError(f"Error evaluating embeddings: {e}")
+
     if embeddings and all((e is None) for e in embeddings):
         embeddings = None
 

agents-api/agents_api/queries/docs/search_docs_by_embedding.py

@@ -7,6 +7,7 @@
 
 from ...autogen.openapi_model import DocReference
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
+from .utils import transform_to_doc_reference
 
 # Raw query for vector search
 search_docs_by_embedding_query = """
@@ -33,14 +34,7 @@
 )
 @wrap_in_class(
     DocReference,
-    transform=lambda d: {
-        "owner": {
-            "id": d["owner_id"],
-            "role": d["owner_type"],
-        },
-        "metadata": d.get("metadata", {}),
-        **d,
-    },
+    transform=transform_to_doc_reference,
 )
 @pg_query
 @beartype

agents-api/agents_api/queries/docs/search_docs_by_text.py

@@ -7,6 +7,7 @@
 
 from ...autogen.openapi_model import DocReference
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
+from .utils import transform_to_doc_reference
 
 # Raw query for text search
 search_docs_text_query = """
@@ -33,14 +34,7 @@
 )
 @wrap_in_class(
     DocReference,
-    transform=lambda d: {
-        "owner": {
-            "id": d["owner_id"],
-            "role": d["owner_type"],
-        },
-        "metadata": d.get("metadata", {}),
-        **d,
-    },
+    transform=transform_to_doc_reference,
 )
 @pg_query
 @beartype

agents-api/agents_api/queries/docs/search_docs_hybrid.py

@@ -7,6 +7,7 @@
 
 from ...autogen.openapi_model import DocReference
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
+from .utils import transform_to_doc_reference
 
 # Raw query for hybrid search
 search_docs_hybrid_query = """
@@ -36,14 +37,7 @@
 )
 @wrap_in_class(
     DocReference,
-    transform=lambda d: {
-        "owner": {
-            "id": d["owner_id"],
-            "role": d["owner_type"],
-        },
-        "metadata": d.get("metadata", {}),
-        **d,
-    },
+    transform=transform_to_doc_reference,
 )
 @pg_query
 @beartype

agents-api/agents_api/queries/docs/utils.py

@@ -0,0 +1,36 @@
+import ast
+
+
+def transform_to_doc_reference(d: dict) -> dict:
+    id = d.pop("doc_id")
+    content = d.pop("content")
+    index = d.pop("index")
+
+    embedding = d.pop("embedding")
+
+    try:
+        # Embeddings are retreived as a string, so we need to evaluate it
+        embedding = ast.literal_eval(embedding)
+    except Exception as e:
+        raise ValueError(f"Error evaluating embeddings: {e}")
+
+    owner = {
+        "id": d.pop("owner_id"),
+        "role": d.pop("owner_type"),
+    }
+    snippet = {
+        "content": content,
+        "index": index,
+        "embedding": embedding,
+    }
+    metadata = d.pop("metadata")
+
+    transformed_data = {
+        "id": id,
+        "owner": owner,
+        "snippet": snippet,
+        "metadata": metadata,
+        **d,
+    }
+
+    return transformed_data