RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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>
This commit is contained in:
Rajat Ahuja 2025-05-13 16:10:16 -05:00 committed by GitHub
parent 606aa14411
commit ecced6b9f8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 218 additions and 150 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

136
src/agent.py
View File

@ -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(

8
src/db.py
View File

@ -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"],

44
src/dependencies.py
View File

@ -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)

123
src/deriver/queue.py
View File

@ -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):

33
src/main.py
View File

@ -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)

20
src/routers/messages.py
View File

@ -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)

4
src/routers/sessions.py
View File

@ -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: