From ecced6b9f865b127faf173b61abf5b7af627da6d Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Tue, 13 May 2025 16:10:16 -0500 Subject: [PATCH] Switch SessionLocal in Dialectic to use get_db (#105) * fix: Manually close transaction with get_db * use get_db instead of session local * fix comment * fix (db): Add application name to each transaction and switch everything to use dependency * chore: Address coderabbit comment --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- src/agent.py | 136 ++++++++++++++++++++-------------------- src/db.py | 8 ++- src/dependencies.py | 44 ++++++++++++- src/deriver/queue.py | 123 +++++++++++++++++++----------------- src/main.py | 33 +++++++--- src/routers/messages.py | 20 +++--- src/routers/sessions.py | 4 ++ 7 files changed, 218 insertions(+), 150 deletions(-) diff --git a/src/agent.py b/src/agent.py index 246ed8d5..3786c937 100644 --- a/src/agent.py +++ b/src/agent.py @@ -14,7 +14,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas -from src.db import SessionLocal +from src.dependencies import tracked_db from src.deriver.tom import get_tom_inference from src.deriver.tom.embeddings import CollectionEmbeddingStore from src.deriver.tom.long_term import get_user_representation_long_term @@ -152,6 +152,7 @@ async def chat( user_id: str, session_id: str, queries: str | list[str], + db: AsyncSession, stream: bool = False, ) -> schemas.DialecticResponse | MessageStreamManager: """ @@ -175,80 +176,77 @@ async def chat( start_time = asyncio.get_event_loop().time() - async with SessionLocal() as db: - # Setup phase - create resources we'll need for all operations + # 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. Create embedding store + collection = await crud.get_or_create_user_protected_collection(db, app_id, user_id) - embedding_store = CollectionEmbeddingStore( - db=db, - app_id=app_id, - user_id=user_id, - collection_id=collection.public_id, # type: ignore - ) - logger.debug( - f"Created embedding store with collection_id: {collection.public_id if collection else None}" - ) + embedding_store = CollectionEmbeddingStore( + db=db, + app_id=app_id, + user_id=user_id, + collection_id=collection.public_id, # type: ignore + ) + logger.debug( + f"Created embedding store with collection_id: {collection.public_id if collection else None}" + ) - # 2. Get the latest user message to attach the user representation to - stmt = ( - select(models.Message) - .join(models.Session, models.Session.public_id == models.Message.session_id) - .join(models.User, models.User.public_id == models.Session.user_id) - .join(models.App, models.App.public_id == models.User.app_id) - .where(models.App.public_id == app_id) - .where(models.User.public_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}") + # 2. Get the latest user message to attach the user representation to + stmt = ( + select(models.Message) + .join(models.Session, models.Session.public_id == models.Message.session_id) + .join(models.User, models.User.public_id == models.Session.user_id) + .join(models.App, models.App.public_id == models.User.app_id) + .where(models.App.public_id == app_id) + .where(models.User.public_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") + # 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) + # 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") + # 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") - # 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" - ) + # 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" + ) # Create a Dialectic chain with the fresh user representation chain = Dialectic( @@ -516,7 +514,7 @@ RELEVANT LONG-TERM FACTS ABOUT THE USER: logger.debug(f"Saving representation to message_id: {message_id}") save_start = asyncio.get_event_loop().time() try: - async with SessionLocal() as save_db: + async with tracked_db("agent.generate_user_representation") as save_db: try: # First check if message exists message_check_stmt = select(models.Message).where( diff --git a/src/db.py b/src/db.py index 8bfe8dfb..4960cc0a 100644 --- a/src/db.py +++ b/src/db.py @@ -1,3 +1,4 @@ +import contextvars import os from dotenv import load_dotenv @@ -7,9 +8,10 @@ from sqlalchemy.orm import declarative_base load_dotenv() -connect_args = { - "prepare_threshold": None, -} +connect_args = {"prepare_threshold": None} + +# Context variable to store request context +request_context = contextvars.ContextVar("request_context", default=None) engine = create_async_engine( os.environ["CONNECTION_URI"], diff --git a/src/dependencies.py b/src/dependencies.py index 6183dd3e..a36153f8 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -1,19 +1,61 @@ +import uuid +from contextlib import asynccontextmanager + from fastapi import Depends +from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession -from .db import SessionLocal +from .db import SessionLocal, request_context async def get_db(): """FastAPI Dependency Generator for Database""" + + context = request_context.get() or "unknown" + db: AsyncSession = SessionLocal() try: + await db.execute(text(f"SET application_name = '{context}'")) yield db except Exception: await db.rollback() raise finally: + if db.in_transaction(): + await db.rollback() await db.close() +@asynccontextmanager +async def tracked_db(operation_name=None): + """Context manager for tracked database sessions""" + # Get request ID if available, or create operation-specific one + context = request_context.get() + token = None + + if not context and operation_name: + context = f"task:{operation_name}:{str(uuid.uuid4())[:8]}" + token = request_context.set(context) + + # Create session with tracking info + db = SessionLocal() + + try: + await db.execute( + text(f"SET application_name = '{context or f'task:{operation_name}'}'") + ) + + yield db + # Explicitly end transaction if still open + if db.in_transaction(): + await db.rollback() # Or commit if needed for write operations + except Exception: + await db.rollback() + raise + finally: + await db.close() + if token: # Only reset if we set it + request_context.reset(token) + + db: AsyncSession = Depends(get_db) diff --git a/src/deriver/queue.py b/src/deriver/queue.py index 33d87e5a..c784482f 100644 --- a/src/deriver/queue.py +++ b/src/deriver/queue.py @@ -13,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import func from .. import models -from ..db import SessionLocal +from ..dependencies import tracked_db from .consumer import process_item logger = getLogger(__name__) @@ -91,7 +91,8 @@ class QueueManager: if self.owned_sessions: logger.info(f"Cleaning up {len(self.owned_sessions)} owned sessions...") try: - async with SessionLocal() as db: + # Use the tracked_db dependency for transaction safety + async with tracked_db("queue_cleanup") as db: await db.execute( delete(models.ActiveQueueSession).where( models.ActiveQueueSession.session_id.in_( @@ -153,7 +154,8 @@ class QueueManager: await asyncio.sleep(1) # Wait before trying again continue - async with SessionLocal() as db: + # Use the dependency for transaction safety + async with tracked_db("queue_polling_loop") as db: try: new_sessions = await self.get_available_sessions(db) @@ -181,7 +183,7 @@ class QueueManager: ) self.add_task(task) except IntegrityError: - await db.rollback() + # Note: rollback is handled by tracked_db dependency logger.debug( f"Failed to claim session {session_id}, already owned" ) @@ -192,7 +194,7 @@ class QueueManager: logger.error(f"Error in polling loop: {str(e)}", exc_info=True) if os.getenv("SENTRY_ENABLED", "False").lower() == "true": sentry_sdk.capture_exception(e) - await db.rollback() + # Note: rollback is handled by tracked_db dependency await asyncio.sleep(1) finally: logger.info("Polling loop stopped") @@ -205,66 +207,69 @@ class QueueManager: async def process_session(self, session_id: int): """Process all messages for a session""" logger.debug(f"Starting to process session {session_id}") - async with self.semaphore: # Hold the semaphore for the entire session duration - async with SessionLocal() as db: - try: - message_count = 0 - while not self.shutdown_event.is_set(): - message = await self.get_next_message(db, session_id) - if not message: - logger.debug(f"No more messages for session {session_id}") - break - - message_count += 1 - logger.debug( - f"Processing message {message.id} for session {session_id} (message {message_count})" - ) - try: - logger.info( - f"Processing message {message.id} from session {session_id}" - ) - await process_item(db, payload=message.payload) - logger.debug(f"Successfully processed message {message.id}") - except Exception as e: - logger.error( - f"Error processing message {message.id}: {str(e)}", - exc_info=True, - ) - if os.getenv("SENTRY_ENABLED", "False").lower() == "true": - sentry_sdk.capture_exception(e) - finally: - # Prevent malformed messages from stalling queue indefinitely - message.processed = True - await db.commit() - logger.debug(f"Marked message {message.id} as processed") - - if self.shutdown_event.is_set(): - logger.debug( - f"Shutdown requested, stopping processing for session {session_id}" - ) - break - - # Update last_updated timestamp to show this session is still being processed - await db.execute( - update(models.ActiveQueueSession) - .where(models.ActiveQueueSession.session_id == session_id) - .values(last_updated=func.now()) - ) - await db.commit() + # Use the tracked_db dependency for transaction safety + async with ( + self.semaphore, + tracked_db("queue_process_session") as db, + ): # Hold the semaphore for the entire session duration + try: + message_count = 0 + while not self.shutdown_event.is_set(): + message = await self.get_next_message(db, session_id) + if not message: + logger.debug(f"No more messages for session {session_id}") + break + message_count += 1 logger.debug( - f"Completed processing session {session_id}, processed {message_count} messages" + f"Processing message {message.id} for session {session_id} (message {message_count})" ) - finally: - # Remove session from active_sessions when done - logger.debug(f"Removing session {session_id} from active sessions") - await db.execute( - delete(models.ActiveQueueSession).where( - models.ActiveQueueSession.session_id == session_id + try: + logger.info( + f"Processing message {message.id} from session {session_id}" ) + await process_item(db, payload=message.payload) + logger.debug(f"Successfully processed message {message.id}") + except Exception as e: + logger.error( + f"Error processing message {message.id}: {str(e)}", + exc_info=True, + ) + if os.getenv("SENTRY_ENABLED", "False").lower() == "true": + sentry_sdk.capture_exception(e) + finally: + # Prevent malformed messages from stalling queue indefinitely + message.processed = True + await db.commit() + logger.debug(f"Marked message {message.id} as processed") + + if self.shutdown_event.is_set(): + logger.debug( + f"Shutdown requested, stopping processing for session {session_id}" + ) + break + + # Update last_updated timestamp to show this session is still being processed + await db.execute( + update(models.ActiveQueueSession) + .where(models.ActiveQueueSession.session_id == session_id) + .values(last_updated=func.now()) ) await db.commit() - self.untrack_session(session_id) + + logger.debug( + f"Completed processing session {session_id}, processed {message_count} messages" + ) + finally: + # Remove session from active_sessions when done + logger.debug(f"Removing session {session_id} from active sessions") + await db.execute( + delete(models.ActiveQueueSession).where( + models.ActiveQueueSession.session_id == session_id + ) + ) + await db.commit() + self.untrack_session(session_id) @sentry_sdk.trace async def get_next_message(self, db: AsyncSession, session_id: int): diff --git a/src/main.py b/src/main.py index a7369efd..08874c50 100644 --- a/src/main.py +++ b/src/main.py @@ -1,5 +1,6 @@ import logging import os +import uuid from contextlib import asynccontextmanager import sentry_sdk @@ -10,7 +11,7 @@ from fastapi_pagination import add_pagination from sentry_sdk.integrations.fastapi import FastApiIntegration from sentry_sdk.integrations.starlette import StarletteIntegration -from src.db import engine +from src.db import engine, request_context from src.exceptions import HonchoException from src.routers import ( apps, @@ -66,17 +67,17 @@ async def setup_admin_jwt(): # Sentry Setup SENTRY_ENABLED = os.getenv("SENTRY_ENABLED", "False").lower() == "true" -if SENTRY_ENABLED: - - def before_send(event, hint): - if 'exc_info' in hint: - exc_type, exc_value, _ = hint['exc_info'] +if SENTRY_ENABLED: + + def before_send(event, hint): + if "exc_info" in hint: + exc_type, exc_value, _ = hint["exc_info"] # Filter out HonchoExceptions from being sent to Sentry if isinstance(exc_value, HonchoException): return None - + return event - + sentry_sdk.init( dsn=os.getenv("SENTRY_DSN"), traces_sample_rate=0.4, @@ -167,3 +168,19 @@ async def global_exception_handler(request: Request, exc: Exception): status_code=500, content={"detail": "An unexpected error occurred"}, ) + + +@app.middleware("http") +async def track_request(request: Request, call_next): + # Create a request ID that includes endpoint information + endpoint = request.url.path.replace("/", "_") + request_id = f"{endpoint}:{str(uuid.uuid4())[:8]}" + + # Store in request state and context var + request.state.request_id = request_id + token = request_context.set(f"api:{request_id}") + + try: + return await call_next(request) + finally: + request_context.reset(token) diff --git a/src/routers/messages.py b/src/routers/messages.py index 602508ed..c4d4f3a6 100644 --- a/src/routers/messages.py +++ b/src/routers/messages.py @@ -8,8 +8,7 @@ from fastapi_pagination.ext.sqlalchemy import paginate from sqlalchemy.sql import insert from src import crud, schemas -from src.db import SessionLocal -from src.dependencies import db +from src.dependencies import db, tracked_db from src.exceptions import ResourceNotFoundException from src.models import QueueItem from src.security import require_auth @@ -34,7 +33,9 @@ async def enqueue(payload: dict | list[dict]): Args: payload: Single message payload or list of message payloads """ - async with SessionLocal() as db: + + # Use the get_db dependency to ensure proper transaction handling + async with tracked_db("message_enqueue") as db_session: try: if isinstance(payload, list): if not payload: # Empty list check @@ -46,7 +47,7 @@ async def enqueue(payload: dict | list[dict]): # Check session once since all messages are for same session try: session = await crud.get_session( - db, + db_session, app_id=payload[0]["app_id"], user_id=payload[0]["user_id"], session_id=payload[0]["session_id"], @@ -89,8 +90,8 @@ async def enqueue(payload: dict | list[dict]): # Use insert to maintain order stmt = insert(QueueItem).returning(QueueItem) - await db.execute(stmt, queue_records) - await db.commit() + await db_session.execute(stmt, queue_records) + await db_session.commit() logger.info(f"Successfully enqueued batch of {len(payload)} messages") return else: @@ -101,7 +102,7 @@ async def enqueue(payload: dict | list[dict]): try: session = await crud.get_session( - db, + db_session, app_id=payload["app_id"], user_id=payload["user_id"], session_id=payload["session_id"], @@ -134,8 +135,8 @@ async def enqueue(payload: dict | list[dict]): .values(payload=processed_payload, session_id=session.id) .returning(QueueItem) ) - await db.execute(stmt) - await db.commit() + await db_session.execute(stmt) + await db_session.commit() logger.info( f"Successfully enqueued message for session {payload['session_id']}" ) @@ -146,7 +147,6 @@ async def enqueue(payload: dict | list[dict]): import sentry_sdk sentry_sdk.capture_exception(e) - await db.rollback() @router.post("", response_model=schemas.Message) diff --git a/src/routers/sessions.py b/src/routers/sessions.py index f8087bdc..53e64f1f 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -206,7 +206,9 @@ async def chat( options: schemas.DialecticOptions = Body( ..., description="Dialectic Endpoint Parameters" ), + db=db, ): + """Chat with the Dialectic API""" if not options.stream: return await agent.chat( @@ -214,6 +216,7 @@ async def chat( user_id=user_id, session_id=session_id, queries=options.queries, + db=db, ) else: @@ -225,6 +228,7 @@ async def chat( session_id=session_id, queries=options.queries, stream=True, + db=db, ) if type(stream) is AsyncMessageStreamManager: async with stream as stream_manager: