From 2e054746a3b8c28b2daa839d2751be0953ee3bb0 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:13:54 +0300 Subject: [PATCH 1/9] Make `tool_calls` nullable in `entries` table --- memory-store/migrations/000015_entries.up.sql | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/memory-store/migrations/000015_entries.up.sql b/memory-store/migrations/000015_entries.up.sql index 10e7693a4..5b9302f05 100644 --- a/memory-store/migrations/000015_entries.up.sql +++ b/memory-store/migrations/000015_entries.up.sql @@ -37,15 +37,14 @@ CREATE TABLE IF NOT EXISTS entries ( name TEXT, content JSONB[] NOT NULL, tool_call_id TEXT DEFAULT NULL, - tool_calls JSONB[] NOT NULL DEFAULT '{}'::JSONB[], + tool_calls JSONB[] DEFAULT NULL, model TEXT NOT NULL, token_count INTEGER DEFAULT NULL, tokenizer TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, timestamp DOUBLE PRECISION NOT NULL, CONSTRAINT pk_entries PRIMARY KEY (session_id, entry_id, created_at), - CONSTRAINT ct_content_is_array_of_objects CHECK (all_jsonb_elements_are_objects (content)), - CONSTRAINT ct_tool_calls_is_array_of_objects CHECK (all_jsonb_elements_are_objects (tool_calls)) + CONSTRAINT ct_content_is_array_of_objects CHECK (all_jsonb_elements_are_objects (content)) ); -- Convert to hypertable if not already From 9a932aec58ff16bd6808fdebc2e86d06cb12bc17 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:14:05 +0300 Subject: [PATCH 2/9] Fix: Storing `tool_calls` in `entries` relation properly --- agents-api/agents_api/queries/entries/create_entries.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agents-api/agents_api/queries/entries/create_entries.py b/agents-api/agents_api/queries/entries/create_entries.py index 1eb24f798..c67f4f115 100644 --- a/agents-api/agents_api/queries/entries/create_entries.py +++ b/agents-api/agents_api/queries/entries/create_entries.py @@ -127,7 +127,7 @@ async def create_entries( item.get("name"), # $6 content_to_json(item.get("content") or {}), # $7 item.get("tool_call_id"), # $8 - content_to_json(item.get("tool_calls") or {}), # $9 + item.get("tool_calls"), # $9 item.get("model"), # $10 item.get("token_count"), # $11 select_tokenizer(item.get("model"))["type"], # $12 From 2a63876f1ff8d6ab2fea75fbff4e07f7eba145f5 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:14:17 +0300 Subject: [PATCH 3/9] Fix: Remove alien attribute --- agents-api/agents_api/routers/sessions/chat.py | 1 - 1 file changed, 1 deletion(-) diff --git a/agents-api/agents_api/routers/sessions/chat.py b/agents-api/agents_api/routers/sessions/chat.py index 2fc5a859e..b5ded8522 100644 --- a/agents-api/agents_api/routers/sessions/chat.py +++ b/agents-api/agents_api/routers/sessions/chat.py @@ -219,7 +219,6 @@ async def chat( developer_id=developer.id, session_id=session_id, data=new_entries, - mark_session_as_updated=True, ) # Adaptive context handling From ff8e9bdf7beafdfb4f297549835c1a51f6aa5b57 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:14:33 +0300 Subject: [PATCH 4/9] Fix: Wrong await usage --- .../agents_api/routers/docs/search_docs.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/agents-api/agents_api/routers/docs/search_docs.py b/agents-api/agents_api/routers/docs/search_docs.py index ead9e1edb..dfbeba711 100644 --- a/agents-api/agents_api/routers/docs/search_docs.py +++ b/agents-api/agents_api/routers/docs/search_docs.py @@ -20,7 +20,7 @@ from .router import router -async def get_search_fn_and_params( +def get_search_fn_and_params( search_params, ) -> Tuple[ Any, Optional[Dict[str, Union[float, int, str, Dict[str, float], List[float]]]] @@ -31,7 +31,7 @@ async def get_search_fn_and_params( case TextOnlyDocSearchRequest( text=query, limit=k, metadata_filter=metadata_filter ): - search_fn = await search_docs_by_text + search_fn = search_docs_by_text params = dict( query=query, k=k, @@ -44,7 +44,7 @@ async def get_search_fn_and_params( confidence=confidence, metadata_filter=metadata_filter, ): - search_fn = await search_docs_by_embedding + search_fn = search_docs_by_embedding params = dict( query_embedding=query_embedding, k=k * 3 if search_params.mmr_strength > 0 else k, @@ -60,7 +60,7 @@ async def get_search_fn_and_params( alpha=alpha, metadata_filter=metadata_filter, ): - search_fn = await search_docs_hybrid + search_fn = search_docs_hybrid params = dict( query=query, query_embedding=query_embedding, @@ -94,10 +94,10 @@ async def search_user_docs( """ # MMR here - search_fn, params = await get_search_fn_and_params(search_params) + search_fn, params = get_search_fn_and_params(search_params) start = time.time() - docs: list[DocReference] = search_fn( + docs: list[DocReference] = await search_fn( developer_id=x_developer_id, owners=[("user", user_id)], **params, @@ -145,10 +145,10 @@ async def search_agent_docs( DocSearchResponse: The search results. """ - search_fn, params = await get_search_fn_and_params(search_params) + search_fn, params = get_search_fn_and_params(search_params) start = time.time() - docs: list[DocReference] = search_fn( + docs: list[DocReference] = await search_fn( developer_id=x_developer_id, owners=[("agent", agent_id)], **params, From 61c8edb1623c87fea2b1bae121bce03f397a54d4 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:14:56 +0300 Subject: [PATCH 5/9] Fix: dump `recall_options` for serialization purpose --- .../agents_api/queries/sessions/create_or_update_session.py | 2 +- agents-api/agents_api/queries/sessions/create_session.py | 2 +- agents-api/agents_api/queries/sessions/patch_session.py | 2 +- agents-api/agents_api/queries/sessions/update_session.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/agents-api/agents_api/queries/sessions/create_or_update_session.py b/agents-api/agents_api/queries/sessions/create_or_update_session.py index b6c280b01..aef7fd1cd 100644 --- a/agents-api/agents_api/queries/sessions/create_or_update_session.py +++ b/agents-api/agents_api/queries/sessions/create_or_update_session.py @@ -143,7 +143,7 @@ async def create_or_update_session( data.token_budget, # $7 data.context_overflow, # $8 data.forward_tool_calls, # $9 - data.recall_options or {}, # $10 + data.recall_options.model_dump() if data.recall_options else {}, # $10 ] # Prepare lookup parameters diff --git a/agents-api/agents_api/queries/sessions/create_session.py b/agents-api/agents_api/queries/sessions/create_session.py index edfe9e1bb..b7196459a 100644 --- a/agents-api/agents_api/queries/sessions/create_session.py +++ b/agents-api/agents_api/queries/sessions/create_session.py @@ -138,7 +138,7 @@ async def create_session( data.token_budget, # $7 data.context_overflow, # $8 data.forward_tool_calls, # $9 - data.recall_options or {}, # $10 + data.recall_options.model_dump() if data.recall_options else {}, # $10 ] # Prepare lookup parameters as a list of parameter lists diff --git a/agents-api/agents_api/queries/sessions/patch_session.py b/agents-api/agents_api/queries/sessions/patch_session.py index d7533e124..1306cde63 100644 --- a/agents-api/agents_api/queries/sessions/patch_session.py +++ b/agents-api/agents_api/queries/sessions/patch_session.py @@ -88,7 +88,7 @@ async def patch_session( data.token_budget, # $7 data.context_overflow, # $8 data.forward_tool_calls, # $9 - data.recall_options or {}, # $10 + data.recall_options.model_dump() if data.recall_options else {}, # $10 ] return [(session_query, session_params)] diff --git a/agents-api/agents_api/queries/sessions/update_session.py b/agents-api/agents_api/queries/sessions/update_session.py index e3f46c0af..c271af488 100644 --- a/agents-api/agents_api/queries/sessions/update_session.py +++ b/agents-api/agents_api/queries/sessions/update_session.py @@ -86,7 +86,7 @@ async def update_session( data.token_budget, # $7 data.context_overflow, # $8 data.forward_tool_calls, # $9 - data.recall_options or {}, # $10 + data.recall_options.model_dump() if data.recall_options else {}, # $10 ] return [ From ba4b802c4066378e6c37a99718b0cad595242aa3 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:15:05 +0300 Subject: [PATCH 6/9] Refactor transform function --- .../agents_api/queries/entries/get_history.py | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/agents-api/agents_api/queries/entries/get_history.py b/agents-api/agents_api/queries/entries/get_history.py index 6a734d4c5..75db9c110 100644 --- a/agents-api/agents_api/queries/entries/get_history.py +++ b/agents-api/agents_api/queries/entries/get_history.py @@ -51,6 +51,29 @@ """).sql(pretty=True) +def _transform(d): + transformed_data = { + "entries": [ + { + **entry, + } + for entry in json.loads(d.get("entries") or "[]") + ], + "relations": [ + { + "head": r["head"], + "relation": r["relation"], + "tail": r["tail"], + } + for r in (d.get("relations") or []) + ], + "session_id": d.get("session_id"), + "created_at": utcnow(), + } + + return transformed_data + + @rewrap_exceptions( { asyncpg.ForeignKeyViolationError: partialclass( @@ -73,19 +96,7 @@ @wrap_in_class( History, one=True, - transform=lambda d: { - "entries": json.loads(d.get("entries") or "[]"), - "relations": [ - { - "head": r["head"], - "relation": r["relation"], - "tail": r["tail"], - } - for r in (d.get("relations") or []) - ], - "session_id": d.get("session_id"), - "created_at": utcnow(), - }, + transform=_transform, ) @pg_query @beartype From cde3e08763ea68e91586840ee008343b7d39851c Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:15:21 +0300 Subject: [PATCH 7/9] Fix docs retreival/search errors --- .../agents_api/queries/docs/list_docs.py | 8 +++++ .../queries/docs/search_docs_by_embedding.py | 10 ++---- .../queries/docs/search_docs_by_text.py | 10 ++---- .../queries/docs/search_docs_hybrid.py | 10 ++---- agents-api/agents_api/queries/docs/utils.py | 36 +++++++++++++++++++ 5 files changed, 50 insertions(+), 24 deletions(-) create mode 100644 agents-api/agents_api/queries/docs/utils.py 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 From 53dcbe4e9254d186c2bc0e0424fbeb543cd8a071 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 03:15:34 +0300 Subject: [PATCH 8/9] Fix: Adjust to new function parameters --- agents-api/agents_api/queries/chat/gather_messages.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/agents-api/agents_api/queries/chat/gather_messages.py b/agents-api/agents_api/queries/chat/gather_messages.py index fb3205acf..56889f110 100644 --- a/agents-api/agents_api/queries/chat/gather_messages.py +++ b/agents-api/agents_api/queries/chat/gather_messages.py @@ -126,8 +126,8 @@ async def gather_messages( doc_references: list[DocReference] = await search_docs_hybrid( developer_id=developer.id, owners=owners, - query=query_text, - query_embedding=query_embedding, + text_query=query_text, + embedding=query_embedding, ) case "text": doc_references: list[DocReference] = await search_docs_by_text( From ed2c342ddf1357b89f132ac453a46dc82338ad88 Mon Sep 17 00:00:00 2001 From: HamadaSalhab Date: Sat, 28 Dec 2024 00:16:45 +0000 Subject: [PATCH 9/9] refactor: Lint agents-api (CI) --- agents-api/agents_api/queries/docs/list_docs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agents-api/agents_api/queries/docs/list_docs.py b/agents-api/agents_api/queries/docs/list_docs.py index 5fcc03e76..12f3b2f1f 100644 --- a/agents-api/agents_api/queries/docs/list_docs.py +++ b/agents-api/agents_api/queries/docs/list_docs.py @@ -3,6 +3,7 @@ It constructs and executes SQL queries to fetch document details based on various filters. """ +import ast from typing import Any, Literal from uuid import UUID @@ -10,7 +11,6 @@ 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