File size: 2,879 Bytes
7e2f74d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | import time
from typing import List, Dict, Any, Optional
from pydantic import BaseModel, Field
class MessageRecord(BaseModel):
role: str # 'user' | 'assistant'
content: str
timestamp: float = Field(default_factory=time.time)
retrieved_context: Optional[List[Dict[str, Any]]] = None
agent_decisions: Optional[Dict[str, Any]] = None
class ConversationSession(BaseModel):
session_id: str
repo_id: str
history: List[MessageRecord] = Field(default_factory=list)
summary: str = ""
class ConversationManager:
"""
Manages multi-session conversation tracking, caching history in-memory.
Persists context, questions, answers, and internal timeline decisions.
"""
def __init__(self):
self._sessions: Dict[str, ConversationSession] = {}
def get_or_create_session(self, session_id: str, repo_id: str) -> ConversationSession:
if session_id not in self._sessions:
self._sessions[session_id] = ConversationSession(session_id=session_id, repo_id=repo_id)
return self._sessions[session_id]
def add_message(
self,
session_id: str,
repo_id: str,
role: str,
content: str,
retrieved_context: Optional[List[Dict[str, Any]]] = None,
agent_decisions: Optional[Dict[str, Any]] = None
) -> MessageRecord:
session = self.get_or_create_session(session_id, repo_id)
record = MessageRecord(
role=role,
content=content,
timestamp=time.time(),
retrieved_context=retrieved_context,
agent_decisions=agent_decisions
)
session.history.append(record)
# Keep summary updated with the last user prompt summary or simple description
if role == "user" and not session.summary:
# First query acts as session title/summary
session.summary = content[:40] + ("..." if len(content) > 40 else "")
return record
def update_summary(self, session_id: str, summary: str):
if session_id in self._sessions:
self._sessions[session_id].summary = summary
def get_session(self, session_id: str) -> Optional[ConversationSession]:
return self._sessions.get(session_id)
def list_sessions_for_repo(self, repo_id: str) -> List[Dict[str, Any]]:
return [
{
"session_id": s.session_id,
"repo_id": s.repo_id,
"summary": s.summary,
"message_count": len(s.history),
"last_updated": s.history[-1].timestamp if s.history else time.time()
}
for s in self._sessions.values() if s.repo_id == repo_id
]
def clear_sessions(self):
self._sessions.clear()
# Global conversation manager singleton
conversation_manager = ConversationManager()
|