Spaces:
Running
Running
| # src/utils/conversation_manager.py | |
| from typing import List, Dict, Optional | |
| import tiktoken | |
| from datetime import datetime | |
| class ConversationManager: | |
| def __init__( | |
| self, | |
| max_tokens: int = 4000, | |
| max_messages: int = 10, | |
| model: str = "gpt-3.5-turbo" | |
| ): | |
| """ | |
| Initialize conversation manager | |
| Args: | |
| max_tokens (int): Maximum tokens to keep in context | |
| max_messages (int): Maximum number of messages to keep | |
| model (str): Model name for token counting | |
| """ | |
| self.max_tokens = max_tokens | |
| self.max_messages = max_messages | |
| self.encoding = tiktoken.encoding_for_model(model) | |
| def format_messages(self, messages: List[Dict]) -> str: | |
| """Format messages into a conversation string""" | |
| formatted = [] | |
| for msg in messages: | |
| role = msg.get('role', 'unknown') | |
| content = msg.get('content', '') | |
| formatted.append(f"{role.capitalize()}: {content}") | |
| return "\n".join(formatted) | |
| def count_tokens(self, text: str) -> int: | |
| """Count tokens in text""" | |
| return len(self.encoding.encode(text)) | |
| def get_relevant_history( | |
| self, | |
| messages: List[Dict], | |
| current_query: str, | |
| max_tokens: Optional[int] = None | |
| ) -> List[Dict]: | |
| """ | |
| Get relevant conversation history within token limit | |
| Args: | |
| messages (List[Dict]): Full message history | |
| current_query (str): Current user query | |
| max_tokens (Optional[int]): Override default max tokens | |
| Returns: | |
| List[Dict]: Relevant message history | |
| """ | |
| max_tokens = max_tokens or self.max_tokens | |
| current_tokens = self.count_tokens(current_query) | |
| # Keep track of tokens and messages | |
| history = [] | |
| total_tokens = current_tokens | |
| # Process messages from most recent to oldest | |
| for msg in reversed(messages[-self.max_messages:]): | |
| msg_text = f"{msg['role']}: {msg['content']}\n" | |
| msg_tokens = self.count_tokens(msg_text) | |
| # Check if adding this message would exceed token limit | |
| if total_tokens + msg_tokens > max_tokens: | |
| break | |
| total_tokens += msg_tokens | |
| history.append(msg) | |
| # Reverse back to chronological order | |
| return list(reversed(history)) | |
| def generate_prompt_with_history( | |
| self, | |
| current_query: str, | |
| history: List[Dict], | |
| context_docs: List[str] | |
| ) -> str: | |
| """ | |
| Generate a prompt that includes conversation history and context | |
| Args: | |
| current_query (str): Current user query | |
| history (List[Dict]): Relevant conversation history | |
| context_docs (List[str]): Retrieved context documents | |
| Returns: | |
| str: Formatted prompt | |
| """ | |
| # Format conversation history | |
| conversation_context = self.format_messages(history) | |
| # Format context documents | |
| context_str = "\n\n".join(context_docs) | |
| prompt = f""" | |
| Previous Conversation: | |
| {conversation_context} | |
| Relevant Context: | |
| {context_str} | |
| Current Query: {current_query} | |
| Based on the previous conversation and the provided context, please provide a comprehensive and accurate response that maintains continuity with the conversation history.""" | |
| return prompt |