src/services/rag_service.py

import google.generativeai as genai
from qdrant_client import QdrantClient
from groq import Groq
from src.core.config import settings

# --- Configuration ---
EMBEDDING_MODEL = 'models/gemini-embedding-001'
GROQ_MODEL = "llama-3.3-70b-versatile"
COLLECTION_NAME = "textbook_content"

class RAGService:
    """
    A service to handle the Retrieval-Augmented Generation pipeline.
    """
    def __init__(self):
        """
        Initializes the RAG service, configuring API clients.
        """
        try:
            # Configure and initialize Gemini client for embeddings
            genai.configure(api_key=settings.GEMINI_API_KEY)
            
            # Initialize Groq client
            self.groq_client = Groq(
                api_key=settings.GROQ_API_KEY,
            )
            
            self.qdrant_client = QdrantClient(url=settings.QDRANT_URL, api_key=settings.QDRANT_API_KEY)
            print("RAGService initialized successfully with Groq for generation.")
        except Exception as e:
            print(f"Error initializing RAGService: {e}")
            # In a real app, you'd want more robust error handling or to prevent the app from starting.
            raise

    def get_answer(self, question: str, highlighted_text: str = None) -> str:
        """
        Generates an answer to a question using a RAG pipeline.

        Args:
            question: The user's question.
            highlighted_text: Optional context from text selected by the user.

        Returns:
            A generated answer based on the retrieved context.
        """
        if not question:
            return "Please provide a question."

        print(f"Received question: {question}")
        print(f"Highlighted context: {highlighted_text}")

        # 1. Create embedding for the query using Gemini
        query_text = f"{highlighted_text}. {question}" if highlighted_text else question
        try:
            response = genai.embed_content(
                model=EMBEDDING_MODEL,
                content=query_text,
                task_type="retrieval_query"
            )
            query_embedding = response['embedding']
        except Exception as e:
            print(f"Error generating embedding with Gemini: {e}")
            return "Sorry, I couldn't process your question at the moment."

        # 2. Query Qdrant for relevant document chunks
        try:
            query_response = self.qdrant_client.query_points(
                collection_name=COLLECTION_NAME,
                query=query_embedding,
                limit=5, # Return top 5 most relevant chunks
            )
            retrieved_chunks = [result.payload['text'] for result in query_response.points]
            context_for_prompt = "\n---\
".join(retrieved_chunks)
            
            if not context_for_prompt:
                print("No relevant context found in the database.")
                context_for_prompt = "No specific context found."

        except Exception as e:
            print(f"Error querying Qdrant: {e}")
            return "Sorry, I had trouble searching my knowledge base."

        # 3. Build a prompt and generate the final answer with Groq
        prompt = f"""
        You are an expert assistant for a textbook on Physical AI and Humanoid Robotics.
        Your task is to answer the user's question based *only* on the provided context below.
        If the context does not contain the answer, state that you cannot answer based on the provided information.
        Do not use any external knowledge.

        **Provided Context:**
        ---
        {context_for_prompt}
        ---

        **User's Question:**
        {question}

        **Answer:**
        """

        try:
            completion = self.groq_client.chat.completions.create(
                model=GROQ_MODEL,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant that answers questions based on provided context."},
                    {"role": "user", "content": prompt},
                ],
            )
            print("Successfully generated answer from Groq.")
            return completion.choices[0].message.content
        except Exception as e:
            print(f"Error generating answer with Groq: {e}")
            return "Sorry, I was unable to generate an answer."

# Singleton instance of the service
rag_service = RAGService()