src/models/rag.py

# src/models/rag.py
from dataclasses import dataclass
from typing import List, Optional, Dict
from datetime import datetime

@dataclass
class Message:
    """Single message in a conversation"""
    role: str  # 'user' or 'assistant'
    content: str
    timestamp: datetime

@dataclass
class ConversationContext:
    """Conversation context with history"""
    messages: List[Message]
    max_messages: int = 10

    def add_message(self, role: str, content: str):
        """Add a message while maintaining max size"""
        self.messages.append(Message(
            role=role,
            content=content,
            timestamp=datetime.now()
        ))
        # Keep only the most recent messages
        if len(self.messages) > self.max_messages:
            self.messages = self.messages[-self.max_messages:]

@dataclass
class RAGResponse:
    """Enhanced RAG response with conversation context"""
    response: str
    context_docs: Optional[List[str]] = None
    sources: Optional[List[Dict]] = None
    scores: Optional[List[float]] = None
    conversation_context: Optional[ConversationContext] = None