src/agents/rag_agent.py

# src/agents/rag_agent.py
from typing import List, Optional, Tuple, Dict
import uuid

from .excel_aware_rag import ExcelAwareRAGAgent
from ..llms.base_llm import BaseLLM
from src.embeddings.base_embedding import BaseEmbedding
from src.vectorstores.base_vectorstore import BaseVectorStore
from src.utils.conversation_manager import ConversationManager
from src.db.mongodb_store import MongoDBStore
from src.models.rag import RAGResponse
from src.utils.logger import logger
from config.config import settings


class RAGAgent(ExcelAwareRAGAgent):
    def __init__(
        self,
        llm: BaseLLM,
        embedding: BaseEmbedding,
        vector_store: BaseVectorStore,
        mongodb: MongoDBStore,
        max_history_tokens: int = 4000,
        max_history_messages: int = 10
    ):
        """
        Initialize RAG Agent

        Args:
            llm (BaseLLM): Language model instance
            embedding (BaseEmbedding): Embedding model instance
            vector_store (BaseVectorStore): Vector store instance
            mongodb (MongoDBStore): MongoDB store instance
            max_history_tokens (int): Maximum tokens in conversation history
            max_history_messages (int): Maximum messages to keep in history
        """
        super().__init__()  # Initialize ExcelAwareRAGAgent
        self.llm = llm
        self.embedding = embedding
        self.vector_store = vector_store
        self.mongodb = mongodb
        self.conversation_manager = ConversationManager(
            max_tokens=max_history_tokens,
            max_messages=max_history_messages
        )

    def _extract_markdown_section(self, docs: List[str], section_header: str) -> str:
        """Extract complete section content from markdown documents"""
        combined_text = '\n'.join(docs)

        section_start = combined_text.find(section_header)
        if section_start == -1:
            return ""

        next_section = combined_text.find(
            "\n\n**", section_start + len(section_header))
        if next_section == -1:
            section_content = combined_text[section_start:]
        else:
            section_content = combined_text[section_start:next_section]

        return self._clean_markdown_content(section_content)

    def _clean_markdown_content(self, content: str) -> str:
        """Clean and format markdown content"""
        lines = content.split('\n')
        seen_lines = set()
        cleaned_lines = []

        for line in lines:
            # Always keep headers and table formatting
            if '| :----' in line or line.startswith('**'):
                if line not in seen_lines:
                    cleaned_lines.append(line)
                    seen_lines.add(line)
                continue

            # Keep table rows and list items
            if line.strip().startswith('|') or line.strip().startswith('-'):
                cleaned_lines.append(line)
                continue

            # Remove duplicates for other content
            stripped = line.strip()
            if stripped and stripped not in seen_lines:
                cleaned_lines.append(line)
                seen_lines.add(stripped)

        return '\n'.join(cleaned_lines)

    async def generate_response(
        self,
        query: str,
        conversation_id: Optional[str],
        temperature: float,
        max_tokens: Optional[int] = None,
        context_docs: Optional[List[str]] = None
    ) -> RAGResponse:
        """Generate response with improved markdown and conversation handling"""
        try:
            # Handle introduction/welcome message queries
            is_introduction = (
                "wants support" in query and
                "This is Introduction" in query and
                ("A new user with name:" in query or "An old user with name:" in query)
            )

            if is_introduction:
                welcome_message = self._handle_contact_query(query)
                return RAGResponse(
                    response=welcome_message,
                    context_docs=[],
                    sources=[],
                    scores=None
                )

            # Get conversation history if conversation_id exists
            history = []
            if conversation_id:
                history = await self.mongodb.get_recent_messages(
                    conversation_id,
                    limit=self.conversation_manager.max_messages
                )
                history = self.conversation_manager.get_relevant_history(
                    messages=history,
                    current_query=query
                )

            # Retrieve context if not provided
            if not context_docs:
                context_docs, sources, scores = await self.retrieve_context(
                    query=query,
                    conversation_history=history
                )
            else:
                sources = None
                scores = None

            # Special handling for markdown section queries
            if "DISCUSSIONS AND ACTION ITEMS" in query.upper():
                section_content = self._extract_markdown_section(
                    context_docs,
                    "**DISCUSSIONS AND ACTION ITEMS**"
                )

                if section_content:
                    return RAGResponse(
                        response=section_content.strip(),
                        context_docs=context_docs,
                        sources=sources,
                        scores=scores
                    )

            # Check if we have any relevant context
            if not context_docs:
                return RAGResponse(
                    response="Information about this is not available, do you want to inquire about something else?",
                    context_docs=[],
                    sources=[],
                    scores=None
                )

            # Generate prompt with context and history
            augmented_prompt = self.conversation_manager.generate_prompt_with_history(
                current_query=query,
                history=history,
                context_docs=context_docs
            )

            # Generate response
            response = self.llm.generate(
                prompt=augmented_prompt,
                temperature=temperature,
                max_tokens=max_tokens
            )

            # Clean the response
            cleaned_response = self._clean_response(response)

            # Return the final response
            return RAGResponse(
                response=cleaned_response,
                context_docs=context_docs,
                sources=sources,
                scores=scores
            )

        except Exception as e:
            logger.error(f"Error in RAGAgent: {str(e)}")
            raise

    def _create_response_prompt(self, query: str, context_docs: List[str]) -> str:
        """
        Create prompt for generating response from context

        Args:
            query (str): User query
            context_docs (List[str]): Retrieved context documents

        Returns:
            str: Formatted prompt for the LLM
        """
        if not context_docs:
            return f"Query: {query}\nResponse: Information about this is not available, do you want to inquire about something else?"

        # Format context documents
        formatted_context = "\n\n".join(
            f"Context {i+1}:\n{doc.strip()}"
            for i, doc in enumerate(context_docs)
            if doc and doc.strip()
        )

        # Build the prompt with detailed instructions
        prompt = f"""You are a knowledgeable assistant. Use the following context to answer the query accurately and informatively.

    Context Information:
    {formatted_context}

    Query: {query}

    Instructions:
    1. Base your response ONLY on the information provided in the context above
    2. If the context contains numbers, statistics, or specific details, include them in your response
    3. Keep your response focused and relevant to the query
    4. Use clear and professional language
    5. If the context includes technical terms, explain them appropriately
    6. Do not make assumptions or add information not present in the context
    7. If specific sections of a report are mentioned, maintain their original structure
    8. Format the response in a clear, readable manner
    9. If the context includes chronological information, maintain the proper sequence

    Response:"""

        return prompt

    async def retrieve_context(
        self,
        query: str,
        conversation_history: Optional[List[Dict]] = None
    ) -> Tuple[List[str], List[Dict], Optional[List[float]]]:
        """
        Retrieve context with conversation history enhancement
        """
        # Enhance query with conversation history
        if conversation_history:
            recent_queries = [
                msg['query'] for msg in conversation_history[-2:]
                if msg.get('query')
            ]
            enhanced_query = " ".join([*recent_queries, query])
        else:
            enhanced_query = query

        # Debug log the enhanced query
        logger.info(f"Enhanced query: {enhanced_query}")

        # Embed the enhanced query
        query_embedding = self.embedding.embed_query(enhanced_query)

        # Debug log embedding shape
        logger.info(f"Query embedding shape: {len(query_embedding)}")

        # Retrieve similar documents
        results = self.vector_store.similarity_search(
            query_embedding,
            top_k=settings.TOP_CHUNKS
        )

        # Debug log search results
        logger.info(f"Number of search results: {len(results)}")
        for i, result in enumerate(results):
            logger.info(f"Result {i} score: {result.get('score', 'N/A')}")
            logger.info(
                f"Result {i} text preview: {result.get('text', '')[:100]}...")

        # Process results
        documents = [doc['text'] for doc in results]
        sources = [self._convert_metadata_to_strings(doc['metadata'])
                   for doc in results]
        scores = [doc['score'] for doc in results
                  if doc.get('score') is not None]

        # Return scores only if available for all documents
        if len(scores) != len(documents):
            scores = None

        return documents, sources, scores

    def _convert_metadata_to_strings(self, metadata: Dict) -> Dict:
        """Convert numeric metadata values to strings"""
        converted = {}
        for key, value in metadata.items():
            if isinstance(value, (int, float)):
                converted[key] = str(value)
            else:
                converted[key] = value
        return converted