diff --git a/agents-api/agents_api/app.py b/agents-api/agents_api/app.py index c178df7c2..469cbb85c 100644 --- a/agents-api/agents_api/app.py +++ b/agents-api/agents_api/app.py @@ -22,15 +22,22 @@ class ObjectWithState(Protocol): state: State +pool = None + + # TODO: This currently doesn't use env.py, we should move to using them @asynccontextmanager async def lifespan(*containers: FastAPI | ObjectWithState): # INIT POSTGRES # pg_dsn = os.environ.get("PG_DSN") + global pool + if not pool: + pool = await create_db_pool(pg_dsn) + for container in containers: if hasattr(container, "state") and not getattr(container.state, "postgres_pool", None): - container.state.postgres_pool = await create_db_pool(pg_dsn) + container.state.postgres_pool = pool # INIT S3 # s3_access_key = os.environ.get("S3_ACCESS_KEY") @@ -50,13 +57,13 @@ async def lifespan(*containers: FastAPI | ObjectWithState): try: yield finally: - # CLOSE POSTGRES # - for container in containers: - if hasattr(container, "state") and getattr(container.state, "postgres_pool", None): - pool = getattr(container.state, "postgres_pool", None) - if pool: - await pool.close() - container.state.postgres_pool = None + # # CLOSE POSTGRES # + # for container in containers: + # if hasattr(container, "state") and getattr(container.state, "postgres_pool", None): + # pool = getattr(container.state, "postgres_pool", None) + # if pool: + # await pool.close() + # container.state.postgres_pool = None # CLOSE S3 # for container in containers: diff --git a/agents-api/agents_api/queries/chat/gather_messages.py b/agents-api/agents_api/queries/chat/gather_messages.py index 0a644a7bb..b589f2360 100644 --- a/agents-api/agents_api/queries/chat/gather_messages.py +++ b/agents-api/agents_api/queries/chat/gather_messages.py @@ -120,7 +120,7 @@ async def gather_messages( doc_references = await search_docs_by_embedding( developer_id=developer.id, owners=owners, - query_embedding=query_embedding, + embedding=query_embedding, connection_pool=connection_pool, ) case "hybrid": 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 fb5110a56..3addcd236 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 @@ -13,7 +13,7 @@ search_docs_by_embedding_query = """ SELECT * FROM search_by_vector( $1, -- developer_id - $2::vector(1024), -- query_embedding + $2::vector(1024), -- embedding $3::text[], -- owner_types $4::uuid[], -- owner_ids $5, -- k @@ -33,7 +33,7 @@ async def search_docs_by_embedding( *, developer_id: UUID, - query_embedding: list[float], + embedding: list[float], k: int = 10, owners: list[tuple[Literal["user", "agent"], UUID]], confidence: float = 0.5, @@ -44,7 +44,7 @@ async def search_docs_by_embedding( Parameters: developer_id (UUID): The ID of the developer. - query_embedding (List[float]): The vector to query. + embedding (List[float]): The vector to query. k (int): The number of results to return. owners (list[tuple[Literal["user", "agent"], UUID]]): List of (owner_type, owner_id) tuples. confidence (float): The confidence threshold for the search. @@ -56,11 +56,11 @@ async def search_docs_by_embedding( if k < 1: raise HTTPException(status_code=400, detail="k must be >= 1") - if not query_embedding: + if not embedding: raise HTTPException(status_code=400, detail="Empty embedding provided") - # Convert query_embedding to a string - query_embedding_str = f"[{', '.join(map(str, query_embedding))}]" + # Convert embedding to a string + embedding_str = f"[{', '.join(map(str, embedding))}]" # Extract owner types and IDs owner_types: list[str] = [owner[0] for owner in owners] @@ -70,7 +70,7 @@ async def search_docs_by_embedding( search_docs_by_embedding_query, [ developer_id, - query_embedding_str, + embedding_str, owner_types, owner_ids, k, diff --git a/agents-api/agents_api/queries/utils.py b/agents-api/agents_api/queries/utils.py index 92841cf83..82add454b 100644 --- a/agents-api/agents_api/queries/utils.py +++ b/agents-api/agents_api/queries/utils.py @@ -116,7 +116,7 @@ async def wrapper( pool = ( connection_pool if connection_pool is not None - else cast(asyncpg.Pool, app.state.postgres_pool) + else cast(asyncpg.Pool, getattr(app.state, "postgres_pool", None)) ) try: diff --git a/agents-api/agents_api/routers/docs/search_docs.py b/agents-api/agents_api/routers/docs/search_docs.py index 947025bed..42ed9b21c 100644 --- a/agents-api/agents_api/routers/docs/search_docs.py +++ b/agents-api/agents_api/routers/docs/search_docs.py @@ -40,14 +40,14 @@ def get_search_fn_and_params( } case VectorDocSearchRequest( - vector=query_embedding, + vector=embedding, limit=k, confidence=confidence, metadata_filter=metadata_filter, ): search_fn = search_docs_by_embedding params = { - "query_embedding": query_embedding, + "embedding": embedding, "k": k * 3 if search_params.mmr_strength > 0 else k, "confidence": confidence, "metadata_filter": metadata_filter, @@ -55,7 +55,7 @@ def get_search_fn_and_params( case HybridDocSearchRequest( text=query, - vector=query_embedding, + vector=embedding, lang=lang, limit=k, confidence=confidence, @@ -66,7 +66,7 @@ def get_search_fn_and_params( search_fn = search_docs_hybrid params = { "text_query": query, - "embedding": query_embedding, + "embedding": embedding, "k": k * 3 if search_params.mmr_strength > 0 else k, "confidence": confidence, "alpha": alpha, @@ -111,7 +111,7 @@ async def search_user_docs( and len(docs) > search_params.limit ): indices = maximal_marginal_relevance( - np.asarray(params["query_embedding"]), + np.asarray(params["embedding"]), [doc.snippet.embedding for doc in docs], k=search_params.limit, ) @@ -160,7 +160,7 @@ async def search_agent_docs( and len(docs) > search_params.limit ): indices = maximal_marginal_relevance( - np.asarray(params["query_embedding"]), + np.asarray(params["embedding"]), [doc.snippet.embedding for doc in docs], k=search_params.limit, ) diff --git a/memory-store/migrations/000004_agents.up.sql b/memory-store/migrations/000004_agents.up.sql index 23de2b68d..02e51e78e 100644 --- a/memory-store/migrations/000004_agents.up.sql +++ b/memory-store/migrations/000004_agents.up.sql @@ -14,7 +14,7 @@ CREATE TABLE IF NOT EXISTS agents ( ), about TEXT CONSTRAINT ct_agents_about_length CHECK ( about IS NULL - OR length(about) <= 1000 + OR length(about) <= 5000 ), instructions TEXT[] DEFAULT ARRAY[]::TEXT[], model TEXT NOT NULL, diff --git a/memory-store/migrations/000015_entries.up.sql b/memory-store/migrations/000015_entries.up.sql index 5b9302f05..3d04bafd3 100644 --- a/memory-store/migrations/000015_entries.up.sql +++ b/memory-store/migrations/000015_entries.up.sql @@ -42,7 +42,7 @@ CREATE TABLE IF NOT EXISTS entries ( token_count INTEGER DEFAULT NULL, tokenizer TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - timestamp DOUBLE PRECISION NOT NULL, + timestamp TIMESTAMPTZ 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)) );