96 lines
3.2 KiB
Python
Executable File
96 lines
3.2 KiB
Python
Executable File
"""
|
|
Conversation Management
|
|
Handles message history and context window management.
|
|
"""
|
|
from typing import List, Dict, Optional
|
|
from datetime import datetime
|
|
import uuid
|
|
from loguru import logger
|
|
|
|
|
|
class ConversationManager:
|
|
"""
|
|
Manages conversation history and context.
|
|
|
|
Features:
|
|
- Track multiple conversations
|
|
- Automatic context window management
|
|
- Message summarization when context grows too large
|
|
"""
|
|
|
|
def __init__(self, max_messages: int = 50, max_tokens: int = 8000):
|
|
self.conversations: Dict[str, List[Dict]] = {}
|
|
self.max_messages = max_messages
|
|
self.max_tokens = max_tokens
|
|
|
|
def create_conversation(self) -> str:
|
|
"""Create a new conversation and return its ID."""
|
|
conv_id = str(uuid.uuid4())
|
|
self.conversations[conv_id] = []
|
|
logger.debug(f"Created conversation: {conv_id}")
|
|
return conv_id
|
|
|
|
def get_conversation(self, conv_id: str) -> List[Dict]:
|
|
"""Get messages for a conversation."""
|
|
return self.conversations.get(conv_id, [])
|
|
|
|
def add_message(
|
|
self,
|
|
conv_id: str,
|
|
role: str,
|
|
content: str,
|
|
metadata: Optional[Dict] = None
|
|
) -> None:
|
|
"""Add a message to a conversation."""
|
|
if conv_id not in self.conversations:
|
|
self.conversations[conv_id] = []
|
|
|
|
message = {
|
|
"role": role,
|
|
"content": content,
|
|
"timestamp": datetime.now().isoformat(),
|
|
}
|
|
|
|
if metadata:
|
|
message["metadata"] = metadata
|
|
|
|
self.conversations[conv_id].append(message)
|
|
|
|
# Trim if needed
|
|
self._trim_conversation(conv_id)
|
|
|
|
def _trim_conversation(self, conv_id: str) -> None:
|
|
"""Trim conversation if it exceeds limits."""
|
|
messages = self.conversations.get(conv_id, [])
|
|
|
|
if len(messages) > self.max_messages:
|
|
# Keep system messages and last N messages
|
|
system_messages = [m for m in messages if m["role"] == "system"]
|
|
other_messages = [m for m in messages if m["role"] != "system"]
|
|
|
|
# Keep last N-1 messages (plus system)
|
|
keep_count = self.max_messages - len(system_messages) - 1
|
|
trimmed = system_messages + other_messages[-keep_count:]
|
|
|
|
self.conversations[conv_id] = trimmed
|
|
logger.debug(f"Trimmed conversation {conv_id} to {len(trimmed)} messages")
|
|
|
|
def delete_conversation(self, conv_id: str) -> None:
|
|
"""Delete a conversation."""
|
|
if conv_id in self.conversations:
|
|
del self.conversations[conv_id]
|
|
logger.debug(f"Deleted conversation: {conv_id}")
|
|
|
|
def list_conversations(self) -> List[str]:
|
|
"""List all conversation IDs."""
|
|
return list(self.conversations.keys())
|
|
|
|
def estimate_tokens(self, messages: List[Dict]) -> int:
|
|
"""Estimate token count for messages."""
|
|
# Rough estimate: ~4 characters per token
|
|
total_chars = sum(
|
|
len(m.get("content", "")) + len(m.get("role", ""))
|
|
for m in messages
|
|
)
|
|
return total_chars // 4
|