refactor: scope db sessions in chat endpoint to reduce connection hold times. Replace route-scoped db session with multiple tracked_db sessions
This commit is contained in:
Rajat Ahuja
parent
77c9c0f06a
commit
5ae3a92baf
126
src/agent.py
126
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,
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue