feat: add conversation persistence and context management layer
- SQLAlchemy async database layer (SQLite dev, PostgreSQL prod) - Conversation and Message models with UUID primary keys - Token counting utilities using litellm - Context summarization at 80% token threshold - REST API endpoints for multi-turn conversations - 19 conversation tests, 6 token tests (176 total passing) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.5
@@ -0,0 +1,354 @@
|
||||
"""
|
||||
Conversation service - Business logic for conversation management.
|
||||
|
||||
Handles CRUD operations, context building, and summarization triggers.
|
||||
"""
|
||||
from uuid import UUID
|
||||
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from src.domains.agents.base import get_agent
|
||||
from src.domains.conversations.models import Conversation, Message
|
||||
from src.domains.conversations.summarize import generate_summary
|
||||
from src.shared.config import get_settings
|
||||
from src.shared.logging import get_logger
|
||||
from src.shared.tokens import count_tokens
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class ConversationService:
|
||||
"""
|
||||
Service for managing conversations and messages.
|
||||
|
||||
Handles:
|
||||
- CRUD operations for conversations and messages
|
||||
- Context building for agent prompts
|
||||
- Automatic summarization when approaching token limits
|
||||
"""
|
||||
|
||||
def __init__(self, session: AsyncSession):
|
||||
"""
|
||||
Initialize with database session.
|
||||
|
||||
Args:
|
||||
session: Async SQLAlchemy session
|
||||
"""
|
||||
self.session = session
|
||||
self.settings = get_settings()
|
||||
|
||||
# === Conversation CRUD ===
|
||||
|
||||
async def create(
|
||||
self,
|
||||
user_id: str,
|
||||
agent_type: str = "explore",
|
||||
working_dir: str = ".",
|
||||
title: str | None = None,
|
||||
) -> Conversation:
|
||||
"""
|
||||
Create a new conversation.
|
||||
|
||||
Args:
|
||||
user_id: Owner's user ID
|
||||
agent_type: Type of agent for this conversation
|
||||
working_dir: Working directory for agent
|
||||
title: Optional title (auto-generated from first message if None)
|
||||
|
||||
Returns:
|
||||
Created Conversation object
|
||||
"""
|
||||
conversation = Conversation(
|
||||
user_id=user_id,
|
||||
agent_type=agent_type,
|
||||
working_dir=working_dir,
|
||||
title=title,
|
||||
)
|
||||
self.session.add(conversation)
|
||||
await self.session.flush()
|
||||
logger.info(f"Created conversation {conversation.id} for user {user_id}")
|
||||
return conversation
|
||||
|
||||
async def get(self, conversation_id: UUID) -> Conversation | None:
|
||||
"""Get conversation by ID without messages."""
|
||||
result = await self.session.execute(
|
||||
select(Conversation).where(Conversation.id == conversation_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def get_with_messages(self, conversation_id: UUID) -> Conversation | None:
|
||||
"""Get conversation by ID with messages loaded."""
|
||||
result = await self.session.execute(
|
||||
select(Conversation)
|
||||
.options(selectinload(Conversation.messages))
|
||||
.where(Conversation.id == conversation_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def list_by_user(
|
||||
self,
|
||||
user_id: str,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
) -> tuple[list[Conversation], int]:
|
||||
"""
|
||||
List conversations for a user.
|
||||
|
||||
Args:
|
||||
user_id: User ID to filter by
|
||||
limit: Maximum results to return
|
||||
offset: Offset for pagination
|
||||
|
||||
Returns:
|
||||
Tuple of (conversations, total_count)
|
||||
"""
|
||||
# Get total count
|
||||
count_result = await self.session.execute(
|
||||
select(func.count(Conversation.id))
|
||||
.where(Conversation.user_id == user_id)
|
||||
)
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# Get conversations
|
||||
result = await self.session.execute(
|
||||
select(Conversation)
|
||||
.where(Conversation.user_id == user_id)
|
||||
.order_by(Conversation.updated_at.desc())
|
||||
.limit(limit)
|
||||
.offset(offset)
|
||||
)
|
||||
conversations = list(result.scalars().all())
|
||||
|
||||
return conversations, total
|
||||
|
||||
async def delete(self, conversation_id: UUID) -> bool:
|
||||
"""Delete a conversation and all its messages."""
|
||||
conversation = await self.get(conversation_id)
|
||||
if conversation:
|
||||
await self.session.delete(conversation)
|
||||
logger.info(f"Deleted conversation {conversation_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
# === Message Operations ===
|
||||
|
||||
async def add_message(
|
||||
self,
|
||||
conversation_id: UUID,
|
||||
role: str,
|
||||
content: str,
|
||||
) -> Message:
|
||||
"""
|
||||
Add a message to a conversation.
|
||||
|
||||
Args:
|
||||
conversation_id: Conversation to add to
|
||||
role: Message role (user, assistant, system, summary)
|
||||
content: Message content
|
||||
|
||||
Returns:
|
||||
Created Message object
|
||||
"""
|
||||
# Count tokens
|
||||
token_count = count_tokens(content)
|
||||
|
||||
message = Message(
|
||||
conversation_id=conversation_id,
|
||||
role=role,
|
||||
content=content,
|
||||
token_count=token_count,
|
||||
)
|
||||
self.session.add(message)
|
||||
|
||||
# Update conversation total tokens
|
||||
conversation = await self.get(conversation_id)
|
||||
if conversation:
|
||||
conversation.total_tokens += token_count
|
||||
|
||||
# Auto-generate title from first user message
|
||||
if conversation.title is None and role == "user":
|
||||
conversation.title = content[:100] + ("..." if len(content) > 100 else "")
|
||||
|
||||
await self.session.flush()
|
||||
return message
|
||||
|
||||
# === Context Building ===
|
||||
|
||||
def build_context_prompt(
|
||||
self,
|
||||
messages: list[Message],
|
||||
current_message: str,
|
||||
) -> str:
|
||||
"""
|
||||
Build a prompt with conversation context.
|
||||
|
||||
Includes summary (if exists) and recent messages.
|
||||
|
||||
Args:
|
||||
messages: All conversation messages
|
||||
current_message: The current user message
|
||||
|
||||
Returns:
|
||||
Formatted prompt with context
|
||||
"""
|
||||
parts = []
|
||||
|
||||
# Find most recent summary
|
||||
summaries = [m for m in messages if m.is_summary]
|
||||
if summaries:
|
||||
latest_summary = summaries[-1]
|
||||
parts.append(
|
||||
f"<conversation_summary>\n{latest_summary.content}\n</conversation_summary>"
|
||||
)
|
||||
|
||||
# Get recent non-summary messages
|
||||
recent = [m for m in messages if not m.is_summary]
|
||||
keep_count = self.settings.keep_recent_messages
|
||||
recent = recent[-keep_count:] if len(recent) > keep_count else recent
|
||||
|
||||
if recent:
|
||||
parts.append("<recent_conversation>")
|
||||
for msg in recent:
|
||||
role_label = msg.role.upper()
|
||||
parts.append(f"{role_label}: {msg.content}")
|
||||
parts.append("</recent_conversation>")
|
||||
|
||||
# Add current message
|
||||
parts.append(f"<current_request>\n{current_message}\n</current_request>")
|
||||
|
||||
return "\n\n".join(parts)
|
||||
|
||||
# === Agent Integration ===
|
||||
|
||||
async def get_agent_response(
|
||||
self,
|
||||
conversation_id: UUID,
|
||||
user_message: str,
|
||||
) -> str:
|
||||
"""
|
||||
Get agent response with conversation context.
|
||||
|
||||
Args:
|
||||
conversation_id: Conversation ID
|
||||
user_message: Current user message
|
||||
|
||||
Returns:
|
||||
Agent's response text
|
||||
"""
|
||||
conversation = await self.get_with_messages(conversation_id)
|
||||
if not conversation:
|
||||
raise ValueError(f"Conversation {conversation_id} not found")
|
||||
|
||||
agent = get_agent(conversation.agent_type)
|
||||
if not agent:
|
||||
raise ValueError(f"Unknown agent type: {conversation.agent_type}")
|
||||
|
||||
# Build context prompt
|
||||
context_prompt = self.build_context_prompt(
|
||||
conversation.messages,
|
||||
user_message,
|
||||
)
|
||||
|
||||
# Run agent
|
||||
response = await agent.run(
|
||||
context_prompt,
|
||||
working_dir=conversation.working_dir,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
# === Summarization ===
|
||||
|
||||
async def should_summarize(self, conversation_id: UUID) -> bool:
|
||||
"""
|
||||
Check if conversation needs summarization.
|
||||
|
||||
Args:
|
||||
conversation_id: Conversation to check
|
||||
|
||||
Returns:
|
||||
True if summarization should be triggered
|
||||
"""
|
||||
conversation = await self.get(conversation_id)
|
||||
if not conversation:
|
||||
return False
|
||||
|
||||
threshold = self.settings.max_context_tokens * self.settings.summarization_threshold
|
||||
return conversation.total_tokens > threshold
|
||||
|
||||
async def summarize_if_needed(self, conversation_id: UUID) -> bool:
|
||||
"""
|
||||
Summarize old messages if approaching token limit.
|
||||
|
||||
Args:
|
||||
conversation_id: Conversation to check and potentially summarize
|
||||
|
||||
Returns:
|
||||
True if summarization was performed
|
||||
"""
|
||||
if not await self.should_summarize(conversation_id):
|
||||
return False
|
||||
|
||||
conversation = await self.get_with_messages(conversation_id)
|
||||
if not conversation:
|
||||
return False
|
||||
|
||||
messages = conversation.messages
|
||||
keep_count = self.settings.keep_recent_messages
|
||||
|
||||
# Don't summarize if not enough messages
|
||||
if len(messages) <= keep_count + 1:
|
||||
return False
|
||||
|
||||
# Get messages to summarize (exclude recent and existing summaries)
|
||||
non_summary_msgs = [m for m in messages if not m.is_summary]
|
||||
to_summarize = non_summary_msgs[:-keep_count]
|
||||
|
||||
if not to_summarize:
|
||||
return False
|
||||
|
||||
logger.info(
|
||||
f"Summarizing {len(to_summarize)} messages in conversation {conversation_id}"
|
||||
)
|
||||
|
||||
# Generate summary
|
||||
summary_text = await generate_summary(
|
||||
to_summarize,
|
||||
working_dir=conversation.working_dir,
|
||||
)
|
||||
|
||||
# Get ID of last summarized message
|
||||
last_summarized_id = to_summarize[-1].id
|
||||
|
||||
# Calculate tokens being removed
|
||||
removed_tokens = sum(m.token_count for m in to_summarize)
|
||||
summary_tokens = count_tokens(summary_text)
|
||||
|
||||
# Add summary message
|
||||
summary_message = Message(
|
||||
conversation_id=conversation_id,
|
||||
role="summary",
|
||||
content=summary_text,
|
||||
token_count=summary_tokens,
|
||||
is_summary=True,
|
||||
summarizes_up_to=last_summarized_id,
|
||||
)
|
||||
self.session.add(summary_message)
|
||||
|
||||
# Mark old messages as summarized (soft delete by excluding from context)
|
||||
for msg in to_summarize:
|
||||
msg.is_summary = True # Reuse flag to mark as "summarized away"
|
||||
|
||||
# Update conversation token count
|
||||
conversation.total_tokens = conversation.total_tokens - removed_tokens + summary_tokens
|
||||
|
||||
await self.session.flush()
|
||||
|
||||
logger.info(
|
||||
f"Summarization complete: removed {removed_tokens} tokens, "
|
||||
f"added {summary_tokens} token summary"
|
||||
)
|
||||
|
||||
return True
|
||||
Reference in New Issue
Block a user