app.py

import os
import chromadb
import gradio as gr
from pathlib import Path
from typing import List, Dict, Any
from dataclasses import dataclass
from sentence_transformers import SentenceTransformer
from google import genai
from google.genai import types

# --- 1. CONFIGURATION AND INITIALIZATION ---

API_KEY = os.getenv("GEMINI_API_KEY")
if not API_KEY:
    raise ValueError("Missing GEMINI_API_KEY in environment variables or Hugging Face secrets.")

client = genai.Client(api_key=API_KEY)
generation_config = types.GenerateContentConfig(
    temperature=0.1,
)
# --- 2. RAG CORE CLASSES (Final Corrected Versions) ---

@dataclass
class Document:
    text: str
    metadata: Dict[str, Any]
    doc_id: str

class VectorStore:
    def __init__(self, collection_name: str = "rag_collection", model_name: str = "all-MiniLM-L6-v2"):
        self.persist_dir = "./chroma_db"
        Path(self.persist_dir).mkdir(exist_ok=True)
        
        # Load the Sentence Transformer model
        self.embedding_model = SentenceTransformer(model_name)
        
        # Initialize Chroma Client
        self.client = chromadb.PersistentClient(path=self.persist_dir) 
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self._create_embedding_function()
        )
        print("VectorStore initialized and ready.")


    def _create_embedding_function(self):
        # Inner class to wrap the Sentence Transformer model for ChromaDB
        class SentenceTransformerEmbeddingFunction(chromadb.api.types.EmbeddingFunction):
            def __init__(self, model):
                self.embedding_model = model
            
            def __call__(self, texts: List[str]) -> List[List[float]]:
                # This function is called by ChromaDB to generate embeddings
                return self.embedding_model.encode(texts, convert_to_tensor=False).tolist()
        
        # Pass the model instance to the inner class
        return SentenceTransformerEmbeddingFunction(self.embedding_model)

    def add_documents(self, documents: List[Document]):
        texts = [doc.text for doc in documents]
        metadatas = [doc.metadata for doc in documents]
        ids = [doc.doc_id for doc in documents]
        
        self.collection.upsert(
            documents=texts,
            metadatas=metadatas,
            ids=ids
        )
        print(f"Added/Updated **{len(documents)}** documents in the vector store.")

    def search(self, query: str, k: int = 4) -> List[Document]:
        try:
            # FIX: Removed 'ids' from the include parameter. 'ids' is returned by default.
            results = self.collection.query(
                query_texts=[query],
                n_results=k,
                include=['documents', 'metadatas'] 
            )
        except Exception as e:
            print(f"ChromaDB search failed: {e}")
            return []

        retrieved_docs = []
        # Ensure the required keys exist and have data
        if results and 'documents' in results and results['documents'] and results['ids']:
            # results['ids'][0] contains the IDs corresponding to the retrieved documents
            for text, metadata, doc_id in zip(results['documents'][0], results['metadatas'][0], results['ids'][0]):
                retrieved_docs.append(Document(text=text, metadata=metadata, doc_id=doc_id))
        
        return retrieved_docs

class RAGPipeline:
    def __init__(self, client: genai.Client, vector_store: VectorStore, config: types.GenerateContentConfig):
        self.client = client
        self.vector_store = vector_store
        self.config = config
        self.model = "gemini-2.5-flash"
        self.system_prompt = (
            "You are an expert Question Answering system. "
            "Your task is to answer the user's question only based on the provided CONTEXT. "
            "If the CONTEXT does not contain the answer, state that you cannot find the answer in the provided documents. "
            "Always cite the source document ID (e.g., [chunk1]) after each piece of information you use. "
            "Format your response using Markdown."
        )

    def _create_prompt(self, query: str, context: List[Document]) -> str:
        context_text = "\n---\n".join([
            f"Document ID: {doc.doc_id} (Source: {doc.metadata.get('source', 'Unknown')})\nContent: {doc.text}"
            for doc in context
        ])
        
        prompt = (
            f"{self.system_prompt}\n\n"
            f"CONTEXT:\n{context_text}\n\n"
            f"USER QUESTION: {query}"
        )
        return prompt

    def generate_response(self, query: str, k: int = 4) -> str:
            # --- FIX APPLIED: Removed the check for self.client.api_key which caused the AttributeError ---
            # The key is checked during the initial Client setup, and the try/except block 
            # below will handle runtime errors if the key is invalid or permissions are missing.
            
            # 1. Retrieval
            retrieved_docs = self.vector_store.search(query, k=k)
            
            if not retrieved_docs:
                return "Retrieval failed. No relevant documents found to answer the query."
    
            # 2. Augmentation & Prompt Creation
            final_prompt = self._create_prompt(query, retrieved_docs)
    
            # 3. Generation
            try:
                # Use keyword argument 'text=final_prompt'
                contents = [types.Content(parts=[types.Part.from_text(text=final_prompt)], role="user")]
                
                response = self.client.models.generate_content(
                    model=self.model,
                    contents=contents,
                    config=self.config
                )
    
                # Display the context used below the final answer for transparency
                context_display = "\n\n---\n\n"
                context_display += "**Context Retrieved:**\n"
                for doc in retrieved_docs:
                    context_display += f"* **ID {doc.doc_id}** (Source: {doc.metadata.get('source', 'Unknown')}): {doc.text[:80]}...\n"
                
                return response.text + context_display
    
            except Exception as e:
                # This catch-all handles errors from the API call itself (e.g., invalid API key, quota issues)
                return f"**Gemini Generation Error:** The API call failed. Ensure your GEMINI_API_KEY is valid and has not expired. Details: {e}"


# --- 3. DATA INDEXING (Runs on app startup) ---

def initialize_rag_system() -> RAGPipeline:
    """Initializes the VectorStore and indexes the data."""
    print("--- Initializing RAG System ---")

    # Define Sample Knowledge Base
    knowledge_base = {
        "doc1": "The **Gemini 2.5 Pro** model is highly capable for complex reasoning tasks and can handle up to 2 million tokens of context, making it suitable for long-form content analysis. It was released in 2024.",
        "doc2": "Google's **Gemini 2.5 Flash** model is optimized for speed and efficiency, making it ideal for high-volume, low-latency tasks like summarization and chat. It also has a 2 million token context window.",
        "doc3": "The **Retrieval-Augmented Generation (RAG)** architecture improves LLM responses by fetching relevant, external documents and providing them as context to the model, reducing hallucinations and improving grounding.",
        "doc4": "The new **google-genai** SDK provides a unified interface for interacting with the entire Gemini family of models, including features like function calling and streaming. This SDK is the recommended way to access the latest features.",
        "doc5": "ChromaDB is used in this pipeline as the **vector store** to efficiently manage and search the document embeddings. It leverages **Sentence Transformers** for local embedding generation."
    }

    # Chunking and Document Creation
    rag_documents: List[Document] = []
    chunk_id_counter = 0

    for doc_id, text in knowledge_base.items():
        chunk_id_counter += 1
        rag_documents.append(
            Document(
                text=text,
                metadata={"source": doc_id, "chunk_index": chunk_id_counter},
                doc_id=f"chunk{chunk_id_counter}"
            )
        )
    print(f"Created **{len(rag_documents)}** individual documents/chunks.")
    
    # Initialize Vector Store and Index
    vector_store = VectorStore(collection_name="rag_knowledge_base")
    vector_store.add_documents(rag_documents)
    print("--- Indexing Complete ---")

    # Initialize RAG Pipeline
    rag_pipeline = RAGPipeline(
        client=client, 
        vector_store=vector_store, 
        config=generation_config
    )
    return rag_pipeline

# Initialize the RAG system globally when the app starts
try:
    rag_system = initialize_rag_system()
except Exception as e:
    # This might catch errors during model loading if no network is available, etc.
    print(f"FATAL ERROR during RAG system initialization: {e}")
    rag_system = None


# --- 4. GRADIO INTERFACE ---

def query_rag(query: str, k_chunks: int) -> str:
    """Interface function to run the RAG pipeline."""
    if rag_system is None:
        return "**FATAL ERROR:** The RAG system failed to initialize. Check the console logs for initialization errors (e.g., API key, model downloads)."
    
    if not query:
        return "Please enter a question to query the knowledge base."
    
    return rag_system.generate_response(query=query, k=k_chunks)

# Define the Gradio Interface
iface = gr.Interface(
    fn=query_rag,
    inputs=[
        gr.Textbox(
            label="Your Question", 
            placeholder="e.g., What is the purpose of the RAG architecture?", 
            lines=2
        ),
        gr.Slider(
            minimum=1, 
            maximum=5, 
            step=1, 
            value=3, 
            label="Number of Context Chunks (k)"
        )
    ],
    outputs=gr.Markdown(label="Grounded Answer (Powered by Gemini)"),
    title="RAG Demo with Gemini 2.5 Flash and ChromaDB",
    description=(
        "Ask a question about Gemini 2.5 models or RAG architecture. "
        "The system retrieves relevant context from its small knowledge base and uses **Gemini 2.5 Flash** "
        "to generate a **grounded, cited answer**."
    ),
    examples=[
        ["What makes the Gemini 2.5 Pro model suitable for long content?", 2],
        ["What is the role of ChromaDB in this setup and what embedding model is used?", 3],
        ["Who was the first person to walk on the moon?", 1]
    ]
)

# Launch the app
if __name__ == "__main__":
    iface.launch()