From 57e453f51260f1458e1b0e2c0c86d8af16f3474a Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Thu, 19 Dec 2024 10:13:50 +0530 Subject: [PATCH] feat(memory-store,agents-api): Move is_leaf handling to postgres Signed-off-by: Diwank Singh Tomer --- .../agents_api/queries/agents/create_agent.py | 2 -- .../queries/agents/create_or_update_agent.py | 2 -- .../agents_api/queries/agents/delete_agent.py | 2 -- .../agents_api/queries/agents/get_agent.py | 2 -- .../agents_api/queries/agents/list_agents.py | 3 +- .../agents_api/queries/agents/patch_agent.py | 2 -- .../agents_api/queries/agents/update_agent.py | 2 -- .../queries/entries/create_entries.py | 6 +--- .../queries/entries/delete_entries.py | 4 +-- .../agents_api/queries/entries/get_history.py | 4 +-- .../queries/entries/list_entries.py | 3 +- agents-api/tests/test_entry_queries.py | 2 +- agents-api/tests/test_session_queries.py | 1 - .../migrations/000016_entry_relations.up.sql | 34 +++++++++++-------- 14 files changed, 25 insertions(+), 44 deletions(-) diff --git a/agents-api/agents_api/queries/agents/create_agent.py b/agents-api/agents_api/queries/agents/create_agent.py index 2d8df7978..76c96f46b 100644 --- a/agents-api/agents_api/queries/agents/create_agent.py +++ b/agents-api/agents_api/queries/agents/create_agent.py @@ -3,7 +3,6 @@ It includes functions to construct and execute SQL queries for inserting new agent records. """ -from typing import Any, TypeVar from uuid import UUID from beartype import beartype @@ -15,7 +14,6 @@ from ..utils import ( generate_canonical_name, pg_query, - rewrap_exceptions, wrap_in_class, ) diff --git a/agents-api/agents_api/queries/agents/create_or_update_agent.py b/agents-api/agents_api/queries/agents/create_or_update_agent.py index e96b30c77..ef3a0abe5 100644 --- a/agents-api/agents_api/queries/agents/create_or_update_agent.py +++ b/agents-api/agents_api/queries/agents/create_or_update_agent.py @@ -3,7 +3,6 @@ It constructs and executes SQL queries to insert a new agent or update an existing agent's details based on agent ID and developer ID. """ -from typing import Any, TypeVar from uuid import UUID from beartype import beartype @@ -14,7 +13,6 @@ from ..utils import ( generate_canonical_name, pg_query, - rewrap_exceptions, wrap_in_class, ) diff --git a/agents-api/agents_api/queries/agents/delete_agent.py b/agents-api/agents_api/queries/agents/delete_agent.py index 6738374db..3527f3611 100644 --- a/agents-api/agents_api/queries/agents/delete_agent.py +++ b/agents-api/agents_api/queries/agents/delete_agent.py @@ -3,7 +3,6 @@ It constructs and executes SQL queries to remove agent records and associated data. """ -from typing import Any, TypeVar from uuid import UUID from beartype import beartype @@ -13,7 +12,6 @@ from ...common.utils.datetime import utcnow from ..utils import ( pg_query, - rewrap_exceptions, wrap_in_class, ) diff --git a/agents-api/agents_api/queries/agents/get_agent.py b/agents-api/agents_api/queries/agents/get_agent.py index 916572db1..a731300fa 100644 --- a/agents-api/agents_api/queries/agents/get_agent.py +++ b/agents-api/agents_api/queries/agents/get_agent.py @@ -3,7 +3,6 @@ It constructs and executes SQL queries to fetch agent details based on agent ID and developer ID. """ -from typing import Any, TypeVar from uuid import UUID from beartype import beartype @@ -12,7 +11,6 @@ from ...autogen.openapi_model import Agent from ..utils import ( pg_query, - rewrap_exceptions, wrap_in_class, ) diff --git a/agents-api/agents_api/queries/agents/list_agents.py b/agents-api/agents_api/queries/agents/list_agents.py index ce12b32b3..87a0c942d 100644 --- a/agents-api/agents_api/queries/agents/list_agents.py +++ b/agents-api/agents_api/queries/agents/list_agents.py @@ -3,7 +3,7 @@ It constructs and executes SQL queries to fetch a list of agents based on developer ID with pagination. """ -from typing import Any, Literal, TypeVar +from typing import Any, Literal from uuid import UUID from beartype import beartype @@ -12,7 +12,6 @@ from ...autogen.openapi_model import Agent from ..utils import ( pg_query, - rewrap_exceptions, wrap_in_class, ) diff --git a/agents-api/agents_api/queries/agents/patch_agent.py b/agents-api/agents_api/queries/agents/patch_agent.py index 2325ab33f..69a5a6ca5 100644 --- a/agents-api/agents_api/queries/agents/patch_agent.py +++ b/agents-api/agents_api/queries/agents/patch_agent.py @@ -3,7 +3,6 @@ It constructs and executes SQL queries to update specific fields of an agent based on agent ID and developer ID. """ -from typing import Any, TypeVar from uuid import UUID from beartype import beartype @@ -13,7 +12,6 @@ from ...metrics.counters import increase_counter from ..utils import ( pg_query, - rewrap_exceptions, wrap_in_class, ) diff --git a/agents-api/agents_api/queries/agents/update_agent.py b/agents-api/agents_api/queries/agents/update_agent.py index 79b520cb8..f28e28264 100644 --- a/agents-api/agents_api/queries/agents/update_agent.py +++ b/agents-api/agents_api/queries/agents/update_agent.py @@ -3,7 +3,6 @@ It constructs and executes SQL queries to replace an agent's details based on agent ID and developer ID. """ -from typing import Any, TypeVar from uuid import UUID from beartype import beartype @@ -13,7 +12,6 @@ from ...metrics.counters import increase_counter from ..utils import ( pg_query, - rewrap_exceptions, wrap_in_class, ) diff --git a/agents-api/agents_api/queries/entries/create_entries.py b/agents-api/agents_api/queries/entries/create_entries.py index 72de8db90..fb61b7c7e 100644 --- a/agents-api/agents_api/queries/entries/create_entries.py +++ b/agents-api/agents_api/queries/entries/create_entries.py @@ -1,16 +1,14 @@ from typing import Literal from uuid import UUID -import asyncpg from beartype import beartype -from fastapi import HTTPException from uuid_extensions import uuid7 from ...autogen.openapi_model import CreateEntryRequest, Entry, Relation from ...common.utils.datetime import utcnow from ...common.utils.messages import content_to_json from ...metrics.counters import increase_counter -from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class +from ..utils import pg_query, wrap_in_class # Query for checking if the session exists session_exists_query = """ @@ -47,7 +45,6 @@ head, relation, tail, - is_leaf ) VALUES ($1, $2, $3, $4, $5) RETURNING *; """ @@ -169,7 +166,6 @@ async def add_entry_relations( item.get("head"), # $2 item.get("relation"), # $3 item.get("tail"), # $4 - item.get("is_leaf", False), # $5 ] ) diff --git a/agents-api/agents_api/queries/entries/delete_entries.py b/agents-api/agents_api/queries/entries/delete_entries.py index 4539ae4df..628ef9011 100644 --- a/agents-api/agents_api/queries/entries/delete_entries.py +++ b/agents-api/agents_api/queries/entries/delete_entries.py @@ -1,15 +1,13 @@ from typing import Literal from uuid import UUID -import asyncpg from beartype import beartype -from fastapi import HTTPException from sqlglot import parse_one from ...autogen.openapi_model import ResourceDeletedResponse from ...common.utils.datetime import utcnow from ...metrics.counters import increase_counter -from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class +from ..utils import pg_query, wrap_in_class # Define the raw SQL query for deleting entries with a developer check delete_entry_query = parse_one(""" diff --git a/agents-api/agents_api/queries/entries/get_history.py b/agents-api/agents_api/queries/entries/get_history.py index 7ad940c0a..b0b767c08 100644 --- a/agents-api/agents_api/queries/entries/get_history.py +++ b/agents-api/agents_api/queries/entries/get_history.py @@ -1,12 +1,10 @@ from uuid import UUID -import asyncpg from beartype import beartype -from fastapi import HTTPException from sqlglot import parse_one from ...autogen.openapi_model import History -from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class +from ..utils import pg_query, wrap_in_class # Define the raw SQL query for getting history with a developer check history_query = parse_one(""" diff --git a/agents-api/agents_api/queries/entries/list_entries.py b/agents-api/agents_api/queries/entries/list_entries.py index 4920e39c1..a6c355f53 100644 --- a/agents-api/agents_api/queries/entries/list_entries.py +++ b/agents-api/agents_api/queries/entries/list_entries.py @@ -1,13 +1,12 @@ from typing import Literal from uuid import UUID -import asyncpg from beartype import beartype from fastapi import HTTPException from ...autogen.openapi_model import Entry from ...metrics.counters import increase_counter -from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class +from ..utils import pg_query, wrap_in_class # Query for checking if the session exists session_exists_query = """ diff --git a/agents-api/tests/test_entry_queries.py b/agents-api/tests/test_entry_queries.py index 60a387591..f5b9d8d56 100644 --- a/agents-api/tests/test_entry_queries.py +++ b/agents-api/tests/test_entry_queries.py @@ -10,7 +10,7 @@ from agents_api.autogen.openapi_model import CreateEntryRequest from agents_api.clients.pg import create_db_pool from agents_api.queries.entries import create_entries, list_entries -from tests.fixtures import pg_dsn, test_developer, test_session # , test_session +from tests.fixtures import pg_dsn, test_developer # , test_session MODEL = "gpt-4o-mini" diff --git a/agents-api/tests/test_session_queries.py b/agents-api/tests/test_session_queries.py index 8e512379f..4e04468bf 100644 --- a/agents-api/tests/test_session_queries.py +++ b/agents-api/tests/test_session_queries.py @@ -29,7 +29,6 @@ from tests.fixtures import ( pg_dsn, test_agent, - test_developer, test_developer_id, test_session, test_user, diff --git a/memory-store/migrations/000016_entry_relations.up.sql b/memory-store/migrations/000016_entry_relations.up.sql index c61c7cd24..bcdb7fb72 100644 --- a/memory-store/migrations/000016_entry_relations.up.sql +++ b/memory-store/migrations/000016_entry_relations.up.sql @@ -31,25 +31,29 @@ CREATE INDEX idx_entry_relations_components ON entry_relations (session_id, head CREATE INDEX idx_entry_relations_leaf ON entry_relations (session_id, relation, is_leaf); -CREATE -OR REPLACE FUNCTION enforce_leaf_nodes () RETURNS TRIGGER AS $$ +CREATE OR REPLACE FUNCTION auto_update_leaf_status() RETURNS TRIGGER AS $$ BEGIN - IF NEW.is_leaf THEN - -- Ensure no other relations point to this leaf node as a head - IF EXISTS ( - SELECT 1 FROM entry_relations - WHERE tail = NEW.head AND session_id = NEW.session_id - ) THEN - RAISE EXCEPTION 'Cannot assign relations to a leaf node.'; - END IF; - END IF; + -- Set is_leaf = false for any existing rows that will now have this new relation as a child + UPDATE entry_relations + SET is_leaf = false + WHERE session_id = NEW.session_id + AND tail = NEW.head; + + -- Set is_leaf for the new row based on whether it has any children + NEW.is_leaf := NOT EXISTS ( + SELECT 1 + FROM entry_relations + WHERE session_id = NEW.session_id + AND head = NEW.tail + ); + RETURN NEW; END; $$ LANGUAGE plpgsql; -CREATE TRIGGER trg_enforce_leaf_nodes BEFORE INSERT -OR -UPDATE ON entry_relations FOR EACH ROW -EXECUTE FUNCTION enforce_leaf_nodes (); +CREATE TRIGGER trg_auto_update_leaf_status +BEFORE INSERT OR UPDATE ON entry_relations +FOR EACH ROW +EXECUTE FUNCTION auto_update_leaf_status(); COMMIT; \ No newline at end of file