diff --git a/agents-api/agents_api/queries/docs/list_docs.py b/agents-api/agents_api/queries/docs/list_docs.py index 67bbe83fc..5fcc03e76 100644 --- a/agents-api/agents_api/queries/docs/list_docs.py +++ b/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 diff --git a/agents-api/agents_api/queries/docs/search_docs_by_embedding.py b/agents-api/agents_api/queries/docs/search_docs_by_embedding.py index fd750bc0f..5efe286dc 100644 --- a/agents-api/agents_api/queries/docs/search_docs_by_embedding.py +++ b/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 diff --git a/agents-api/agents_api/queries/docs/search_docs_by_text.py b/agents-api/agents_api/queries/docs/search_docs_by_text.py index 787a83651..e6668293b 100644 --- a/agents-api/agents_api/queries/docs/search_docs_by_text.py +++ b/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 diff --git a/agents-api/agents_api/queries/docs/search_docs_hybrid.py b/agents-api/agents_api/queries/docs/search_docs_hybrid.py index 23eb12318..b91aa2d83 100644 --- a/agents-api/agents_api/queries/docs/search_docs_hybrid.py +++ b/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 diff --git a/agents-api/agents_api/queries/docs/utils.py b/agents-api/agents_api/queries/docs/utils.py new file mode 100644 index 000000000..8a49f4bdc --- /dev/null +++ b/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