agents/rag_agent.py

from openai import AsyncOpenAI
import os
from typing import List, Dict, Any, Optional
import logging
from pydantic import BaseModel
import asyncio

from vector.qdrant_client import get_qdrant_manager
from scripts.embedder import get_embedder
from agents.system_prompt import get_system_prompt

logger = logging.getLogger(__name__)

class Citation(BaseModel):
    """
    Citation model for tracking textbook references
    """
    chapter_title: str
    section_title: str
    url_path: str
    confidence_score: float
    content_snippet: str


class RAGResponse(BaseModel):
    """
    Response model for RAG queries
    """
    content: str
    citations: List[Citation]
    query: str
    retrieved_chunks_count: int


class RAGAgent:
    """
    RAG Agent that retrieves relevant content and generates responses
    """
    def __init__(self):
        # Initialize OpenAI client for Gemini
        # Note: base_url should NOT have trailing slash for Gemini API
        base_url = os.getenv("OPENAI_BASE_URL", "https://generativelanguage.googleapis.com/v1beta/openai")
        # Remove trailing slash if present
        base_url = base_url.rstrip('/')
        api_key = os.getenv("OPENAI_API_KEY")

        if not api_key:
            raise ValueError("OPENAI_API_KEY environment variable is required")

        self.client = AsyncOpenAI(
            base_url=base_url,
            api_key=api_key
        )

        self.model = os.getenv("GEMINI_MODEL", "gemini-2.0-flash")
        self.top_k = int(os.getenv("RAG_TOP_K", "5"))
        self.min_score = float(os.getenv("RAG_MIN_SCORE", "0.1"))  # Lowered threshold
        self.max_tokens = int(os.getenv("MAX_RESPONSE_TOKENS", "2000"))

    async def retrieve_context(self, query: str) -> List[Dict[str, Any]]:
        """
        Retrieve relevant context from vector database
        """
        try:
            # Generate embedding for the query
            logger.info(f"Generating embedding for query: {query}")
            embedder = get_embedder()
            query_embedding = await embedder.generate_embedding(query)
            logger.info(f"Embedding generated, length: {len(query_embedding)}")

            # Search in Qdrant
            qdrant_manager = get_qdrant_manager()
            logger.info(f"Searching Qdrant with top_k={self.top_k}, min_score={self.min_score}")
            search_results = await qdrant_manager.search_similar(
                query_vector=query_embedding,
                top_k=self.top_k,
                min_score=self.min_score
            )
            logger.info(f"Qdrant returned {len(search_results)} results")

            # Format results
            retrieved_chunks = []
            for result in search_results:
                payload = result['payload']
                logger.info(f"Result score: {result['score']}, section: {payload.get('section_path', 'unknown')}")
                retrieved_chunks.append({
                    'id': result['id'],
                    'score': result['score'],
                    'content': payload.get('content', ''),
                    'chapter_id': payload.get('chapter_id', ''),
                    'section_path': payload.get('section_path', ''),
                    'token_count': payload.get('token_count', 0),
                    'content_type': payload.get('content_type', 'text'),
                    'chunk_index': payload.get('chunk_index', 0)
                })

            logger.info(f"Retrieved {len(retrieved_chunks)} chunks for query: {query[:50]}...")
            return retrieved_chunks

        except Exception as e:
            logger.error(f"Error retrieving context: {e}")
            import traceback
            traceback.print_exc()
            return []

    async def generate_response(self, query: str, context_chunks: List[Dict[str, Any]]) -> RAGResponse:
        """
        Generate a response based on query and retrieved context
        """
        try:
            # Build context string from retrieved chunks
            context_str = ""
            citations = []

            for chunk in context_chunks:
                if chunk.get('content'):
                    context_str += f"\n\n[Relevant Content]\n{chunk['content']}\n"

                    # Create citation
                    citation = Citation(
                        chapter_title=chunk.get('chapter_id', 'Unknown Chapter'),
                        section_title=chunk.get('section_path', 'Unknown Section'),
                        url_path=chunk.get('section_path', '#'),
                        confidence_score=chunk.get('score', 0.0),
                        content_snippet=chunk['content'][:200] + "..." if len(chunk['content']) > 200 else chunk['content']
                    )
                    citations.append(citation)

            # Build the full prompt
            system_prompt = get_system_prompt()
            user_prompt = f"""
            Query: {query}

            Retrieved Context:
            {context_str}

            Please provide a comprehensive answer based on the textbook content.
            """

            # Call the LLM
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                max_tokens=self.max_tokens,
                temperature=0.3  # Lower temperature for more consistent, factual responses
            )

            response_content = response.choices[0].message.content

            return RAGResponse(
                content=response_content,
                citations=citations,
                query=query,
                retrieved_chunks_count=len(context_chunks)
            )

        except Exception as e:
            logger.error(f"Error generating response: {e}")
            import traceback
            traceback.print_exc()
            # Return a fallback response when generation fails
            return RAGResponse(
                content="I'm sorry, but I couldn't generate a response at this time. Please try again later.",
                citations=[],
                query=query,
                retrieved_chunks_count=0
            )

    def _is_casual_message(self, query: str) -> bool:
        """
        Check if the query is a casual/greeting message that doesn't need RAG
        """
        casual_patterns = [
            'hi', 'hello', 'hey', 'hii', 'hiii', 'yo',
            'good morning', 'good afternoon', 'good evening', 'good night',
            'thanks', 'thank you', 'thankyou', 'thx',
            'bye', 'goodbye', 'see you', 'ok', 'okay',
            'how are you', 'whats up', "what's up", 'wassup',
            'nice', 'great', 'awesome', 'cool',
            'yes', 'no', 'yeah', 'yep', 'nope',
            'sorry', 'help me', 'can you help'
        ]
        query_lower = query.lower().strip().rstrip('!?.،')
        return len(query_lower) < 20 and any(
            query_lower == p or query_lower.startswith(p) for p in casual_patterns
        )

    async def generate_casual_response(self, query: str) -> RAGResponse:
        """
        Generate a conversational response using LLM without RAG context
        """
        try:
            system_prompt = get_system_prompt()
            user_prompt = f"""User message: {query}

This is a casual/greeting message. Respond naturally and friendly as a helpful robotics tutor. Keep it brief and conversational."""

            response = await self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                max_tokens=300,
                temperature=0.7
            )

            return RAGResponse(
                content=response.choices[0].message.content,
                citations=[],
                query=query,
                retrieved_chunks_count=0
            )
        except Exception as e:
            logger.error(f"Error generating casual response: {e}")
            return RAGResponse(
                content="Hello! I'm your robotics textbook assistant. How can I help you learn about ROS 2 and Physical AI today?",
                citations=[],
                query=query,
                retrieved_chunks_count=0
            )

    async def query(self, query: str) -> RAGResponse:
        """
        Main method to process a query through RAG
        """
        logger.info(f"Processing RAG query: {query[:50]}...")

        # Check for casual messages - use LLM directly without RAG
        if self._is_casual_message(query):
            logger.info("Detected casual message, generating conversational response")
            return await self.generate_casual_response(query)

        # Retrieve relevant context for knowledge questions
        context_chunks = await self.retrieve_context(query)

        if not context_chunks:
            # No relevant content - still use LLM to respond helpfully
            return await self.generate_casual_response(query + "\n(Note: No relevant textbook content found for this query)")

        # Generate response based on context
        response = await self.generate_response(query, context_chunks)

        return response

    async def query_with_streaming(self, query: str):
        """
        Generator method for streaming responses via SSE
        """
        logger.info(f"Processing streaming RAG query: {query[:50]}...")

        # Retrieve relevant context
        context_chunks = await self.retrieve_context(query)

        if not context_chunks:
            # No relevant content found
            yield {
                "content_chunk": "I couldn't find relevant content in the textbook to answer your question. Please try rephrasing your question or check if the topic is covered in the textbook.",
                "is_complete": True,
                "citations": []
            }
            return

        # Build context string from retrieved chunks
        context_str = ""
        citations = []

        for chunk in context_chunks:
            if chunk.get('content'):
                context_str += f"\n\n[Relevant Content]\n{chunk['content']}\n"

                # Create citation
                citation = Citation(
                    chapter_title=chunk.get('chapter_id', 'Unknown Chapter'),
                    section_title=chunk.get('section_path', 'Unknown Section'),
                    url_path=chunk.get('section_path', '#'),
                    confidence_score=chunk.get('score', 0.0),
                    content_snippet=chunk['content'][:200] + "..." if len(chunk['content']) > 200 else chunk['content']
                )
                citations.append(citation)

        # Build the full prompt
        system_prompt = get_system_prompt()
        user_prompt = f"""
        Query: {query}

        Retrieved Context:
        {context_str}

        Please provide a comprehensive answer based on the textbook content.
        """

        try:
            # Stream the response from the LLM
            stream = await self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                max_tokens=self.max_tokens,
                temperature=0.3,
                stream=True
            )

            full_response = ""
            async for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    content_chunk = chunk.choices[0].delta.content
                    full_response += content_chunk

                    yield {
                        "content_chunk": content_chunk,
                        "is_complete": False,
                        "citations": []
                    }

            # Send final chunk with citations
            yield {
                "content_chunk": "",
                "is_complete": True,
                "citations": [c.dict() for c in citations]
            }

        except Exception as e:
            logger.error(f"Error in streaming response: {e}")
            yield {
                "content_chunk": "I'm sorry, but I encountered an error while processing your request. Please try again later.",
                "is_complete": True,
                "citations": []
            }


# Lazy-initialized RAG agent instance
_rag_agent: Optional[RAGAgent] = None

def get_rag_agent() -> RAGAgent:
    """Get or create the RAG agent instance (lazy initialization)"""
    global _rag_agent
    if _rag_agent is None:
        _rag_agent = RAGAgent()
    return _rag_agent