diff --git a/src/agent.py b/src/agent.py index 21ed2698..58939b30 100644 --- a/src/agent.py +++ b/src/agent.py @@ -152,7 +152,6 @@ async def chat( user_id: str, session_id: str, queries: str | list[str], - db: AsyncSession, stream: bool = False, ) -> schemas.DialecticResponse | MessageStreamManager: """ @@ -178,72 +177,71 @@ async def chat( # Setup phase - create resources we'll need for all operations - # 1. Create embedding store - collection = await crud.get_or_create_user_protected_collection(db, app_id, user_id) + # 1. Fetch latest user message & chat history + async with tracked_db("chat.load_history") as db_history: + stmt = ( + select(models.Message) + .where(models.Message.app_id == app_id) + .where(models.Message.user_id == user_id) + .where(models.Message.session_id == session_id) + .where(models.Message.is_user) + .order_by(models.Message.id.desc()) + .limit(1) + ) + latest_messages = await db_history.execute(stmt) + latest_message = latest_messages.scalar_one_or_none() + latest_message_id = latest_message.public_id if latest_message else None + logger.debug(f"Latest user message ID: {latest_message_id}") - embedding_store = CollectionEmbeddingStore( - db=db, - app_id=app_id, - user_id=user_id, - collection_id=collection.public_id, # type: ignore + chat_history, _, _ = await history.get_summarized_history( + db_history, session_id, summary_type=history.SummaryType.SHORT + ) + if not chat_history: + logger.warning(f"No chat history found for session {session_id}") + chat_history = f"someone asked this about the user's message: {final_query}" + logger.debug(f"IDs: {app_id}, {user_id}, {session_id}") + message_count = len(chat_history.split("\n")) + logger.debug(f"Retrieved chat history: {message_count} messages") + + # Run short-term inference and long-term facts in parallel + async def fetch_long_term(): + async with tracked_db("chat.get_collection_and_facts") as db_embed: + collection = await crud.get_or_create_user_protected_collection( + db_embed, app_id, user_id + ) + embedding_store = CollectionEmbeddingStore( + db=db_embed, + app_id=app_id, + user_id=user_id, + collection_id=collection.public_id, # type: ignore + ) + facts = await get_long_term_facts(final_query, embedding_store) + return facts + + long_term_task = asyncio.create_task(fetch_long_term()) + short_term_task = asyncio.create_task(run_tom_inference(chat_history, session_id)) + + facts, tom_inference = await asyncio.gather( + long_term_task, short_term_task ) - logger.debug( - f"Created embedding store with collection_id: {collection.public_id if collection else None}" - ) - - stmt = ( - select(models.Message) - .where(models.Message.app_id == app_id) - .where(models.Message.user_id == user_id) - .where(models.Message.session_id == session_id) - .where(models.Message.is_user) - .order_by(models.Message.id.desc()) - .limit(1) - ) - - latest_messages = await db.execute(stmt) - latest_message = latest_messages.scalar_one_or_none() - latest_message_id = latest_message.public_id if latest_message else None - logger.debug(f"Latest user message ID: {latest_message_id}") - - # Get chat history for the session - chat_history, _, _ = await history.get_summarized_history( - db, session_id, summary_type=history.SummaryType.SHORT - ) - if not chat_history: - logger.warning(f"No chat history found for session {session_id}") - chat_history = f"someone asked this about the user's message: {final_query}" - logger.debug(f"IDs: {app_id}, {user_id}, {session_id}") - message_count = len(chat_history.split("\n")) - logger.debug(f"Retrieved chat history: {message_count} messages") - - # Run both long-term and short-term context retrieval concurrently - logger.debug("Starting parallel tasks for context retrieval") - long_term_task = get_long_term_facts(final_query, embedding_store) - short_term_task = run_tom_inference(chat_history, session_id) - - # Wait for both tasks to complete - facts, tom_inference = await asyncio.gather(long_term_task, short_term_task) - logger.debug(f"Retrieved {len(facts)} facts from long-term memory") - logger.debug(f"TOM inference completed with {len(tom_inference)} characters") + logger.info(f"Retrieved {len(facts)} facts from long-term memory") + logger.info(f"TOM inference completed with {len(tom_inference)} characters") # Generate a fresh user representation logger.debug("Generating user representation") - user_representation = await generate_user_representation( - app_id=app_id, - user_id=user_id, - session_id=session_id, - chat_history=chat_history, - tom_inference=tom_inference, - facts=facts, - embedding_store=embedding_store, - db=db, - message_id=latest_message_id, - with_inference=False, - ) - logger.debug( - f"User representation generated: {len(user_representation)} characters" - ) + async with tracked_db("chat.generate_user_representation") as db_rep: + user_representation = await generate_user_representation( + app_id=app_id, + user_id=user_id, + session_id=session_id, + chat_history=chat_history, + tom_inference=tom_inference, + facts=facts, + db=db_rep, + message_id=latest_message_id, + with_inference=False, + ) + logger.info(f"User representation generated: {len(user_representation)} characters") # Create a Dialectic chain with the fresh user representation chain = Dialectic( @@ -253,7 +251,7 @@ async def chat( ) generation_time = asyncio.get_event_loop().time() - start_time - logger.debug(f"User representation generation completed in {generation_time:.2f}s") + logger.info(f"User representation generation completed in {generation_time:.2f}s") langfuse_context.update_current_trace( session_id=session_id, @@ -430,7 +428,6 @@ async def generate_user_representation( chat_history: str, tom_inference: str, facts: list[str], - embedding_store: CollectionEmbeddingStore, db: AsyncSession, message_id: Optional[str] = None, with_inference: bool = False, @@ -477,7 +474,6 @@ async def generate_user_representation( chat_history=chat_history, session_id=session_id, facts=facts, - embedding_store=embedding_store, user_representation=latest_representation, tom_inference=tom_inference, )