From c8e7929273335d7f7902a9350e3917bd0f3a47da Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Tue, 23 Sep 2025 17:12:39 -0400 Subject: [PATCH] fix: shorten db scope for dialectic and embedding --- src/routers/peers.py | 12 ++-- src/utils/embedding_store.py | 120 +++++++++++++++++------------------ 2 files changed, 66 insertions(+), 66 deletions(-) diff --git a/src/routers/peers.py b/src/routers/peers.py index eabde2ee..c9b393b5 100644 --- a/src/routers/peers.py +++ b/src/routers/peers.py @@ -15,7 +15,7 @@ from mirascope.llm import Stream from sqlalchemy.ext.asyncio import AsyncSession from src import crud, schemas -from src.dependencies import db +from src.dependencies import db, tracked_db from src.dialectic import chat as dialectic_chat from src.exceptions import AuthenticationException, ResourceNotFoundException from src.security import JWTParams, require_auth @@ -160,12 +160,14 @@ async def chat( options: schemas.DialecticOptions = Body( ..., description="Dialectic Endpoint Parameters" ), - db: AsyncSession = db, ): # Get or create the peer to ensure it exists - await crud.get_or_create_peers( - db, workspace_name=workspace_id, peers=[schemas.PeerCreate(name=peer_id)] - ) + async with tracked_db("peers.chat.get_or_create_peer") as peer_db: + await crud.get_or_create_peers( + peer_db, + workspace_name=workspace_id, + peers=[schemas.PeerCreate(name=peer_id)], + ) if not options.stream: response = await dialectic_chat( diff --git a/src/utils/embedding_store.py b/src/utils/embedding_store.py index a45c32c3..a18025d3 100644 --- a/src/utils/embedding_store.py +++ b/src/utils/embedding_store.py @@ -59,74 +59,72 @@ class EmbeddingStore: fallback_level: Reasoning level for the observations if not provided similarity_threshold: Threshold for considering observations similar """ - async with tracked_db("ed_embedding_store.save_unified_observations") as db: - # Extract conclusions for deduplication and embedding - conclusions: list[str] = [obs.conclusion for obs in observations] + # Extract conclusions for deduplication and embedding + conclusions: list[str] = [obs.conclusion for obs in observations] - # Remove duplicates before saving - unique_conclusions: list[str] = await self.remove_duplicates( - conclusions, similarity_threshold=similarity_threshold + # Remove duplicates before saving + unique_conclusions: list[str] = await self.remove_duplicates( + conclusions, similarity_threshold=similarity_threshold + ) + if settings.LANGFUSE_PUBLIC_KEY: + langfuse_context.update_current_observation( + input={"observations": [obs.model_dump() for obs in observations]}, + output={"unique_conclusions": unique_conclusions}, ) - if settings.LANGFUSE_PUBLIC_KEY: - langfuse_context.update_current_observation( - input={"observations": [obs.model_dump() for obs in observations]}, - output={"unique_conclusions": unique_conclusions}, + + if not unique_conclusions: + logger.debug("No unique observations to save after deduplication") + return + + # Create mapping from conclusion back to original observation + conclusion_to_observation: dict[str, UnifiedObservation] = { + obs.conclusion: obs for obs in observations + } + + # Filter unified observations to only unique ones + unique_observations: list[UnifiedObservation] = [ + conclusion_to_observation[conclusion] for conclusion in unique_conclusions + ] + + # Batch embed all unique conclusions (not premises) + embeddings: list[list[float]] = [] + batch_size: int = 2048 # OpenAI batch limit + + for i in range(0, len(unique_conclusions), batch_size): + batch = unique_conclusions[i : i + batch_size] + response: CreateEmbeddingResponse = ( + await embedding_client.client.embeddings.create( + input=batch, model="text-embedding-3-small" ) + ) + embeddings.extend([data.embedding for data in response.data]) - if not unique_conclusions: - logger.debug("No unique observations to save after deduplication") - return + # Batch create document objects + document_objects: list[models.Document] = [] + for obs, embedding in zip(unique_observations, embeddings, strict=True): + # Use the observation's own level or fall back to parameter level + obs_level = obs.level or fallback_level - # Create mapping from conclusion back to original observation - conclusion_to_observation: dict[str, UnifiedObservation] = { - obs.conclusion: obs for obs in observations + # Build metadata including premises + metadata: dict[str, Any] = { + "level": obs_level, + "message_id": message_id, + "session_name": session_name, + "premises": obs.premises, # Store premises in metadata + "created_at": format_datetime_utc(message_created_at), } - # Filter unified observations to only unique ones - unique_observations: list[UnifiedObservation] = [ - conclusion_to_observation[conclusion] - for conclusion in unique_conclusions - ] - - # Batch embed all unique conclusions (not premises) - embeddings: list[list[float]] = [] - batch_size: int = 2048 # OpenAI batch limit - - for i in range(0, len(unique_conclusions), batch_size): - batch = unique_conclusions[i : i + batch_size] - response: CreateEmbeddingResponse = ( - await embedding_client.client.embeddings.create( - input=batch, model="text-embedding-3-small" - ) - ) - embeddings.extend([data.embedding for data in response.data]) - - # Batch create document objects - document_objects: list[models.Document] = [] - for obs, embedding in zip(unique_observations, embeddings, strict=True): - # Use the observation's own level or fall back to parameter level - obs_level = obs.level or fallback_level - - # Build metadata including premises - metadata: dict[str, Any] = { - "level": obs_level, - "message_id": message_id, - "session_name": session_name, - "premises": obs.premises, # Store premises in metadata - "created_at": format_datetime_utc(message_created_at), - } - - doc = models.Document( - workspace_name=self.workspace_name, - peer_name=self.peer_name, - collection_name=self.collection_name, - content=obs.conclusion, # Store only conclusion as content - internal_metadata=metadata, - embedding=embedding, # Embedding generated from conclusion only - created_at=message_created_at, - ) - document_objects.append(doc) - + doc = models.Document( + workspace_name=self.workspace_name, + peer_name=self.peer_name, + collection_name=self.collection_name, + content=obs.conclusion, # Store only conclusion as content + internal_metadata=metadata, + embedding=embedding, # Embedding generated from conclusion only + created_at=message_created_at, + ) + document_objects.append(doc) + async with tracked_db("ed_embedding_store.save_unified_observations") as db: # Batch insert all documents db.add_all(document_objects) await db.commit()