Update app.py

app.py

@@ -1,601 +1,634 @@
-import os
-import time
-import gradio as gr
-import uvicorn
-from fastapi import FastAPI, HTTPException, Depends, File, UploadFile
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from pydantic import BaseModel
-from typing import Optional, Dict, Any
-import threading
-import logging
-from langchain.document_loaders import TextLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
-import tiktoken
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# --- Configuration ---
-CHUNK_SIZE = 800
-CHUNK_OVERLAP = 100
-MAX_TOKENS = 512
-TEMPERATURE = 0.5
-RETRIEVAL_K = 5
-
-# --- Token Counting Setup ---
-try:
-    tokenizer = tiktoken.get_encoding("cl100k_base")
-except:
-    print("Tiktoken encoder 'cl100k_base' not found. Using basic split().")
-    tokenizer = type('obj', (object,), {'encode': lambda x: x.split()})()
-
-def estimate_tokens(text):
-    """Estimates token count for a given text."""
-    return len(tokenizer.encode(text))
-
-# Custom Callback Handler to track LLM token usage
-class TokenUsageCallbackHandler(BaseCallbackHandler):
-    """Callback handler to track token usage in LLM calls."""
-    def __init__(self):
-        super().__init__()
-        self.reset_counters()
-
-    def reset_counters(self):
-        self.total_prompt_tokens = 0
-        self.total_completion_tokens = 0
-        self.total_llm_calls = 0
-
-    def on_llm_end(self, response, **kwargs):
-        """Collect token usage from the LLM response."""
-        self.total_llm_calls += 1
-        llm_output = response.llm_output
-        
-        if llm_output and 'usage_metadata' in llm_output:
-            usage = llm_output['usage_metadata']
-            prompt_tokens = usage.get('prompt_token_count', 0)
-            completion_tokens = usage.get('candidates_token_count', 0)
-            
-            self.total_prompt_tokens += prompt_tokens
-            self.total_completion_tokens += completion_tokens
-
-    def get_total_tokens(self):
-        """Returns the total prompt and completion tokens."""
-        return {
-            "total_prompt_tokens": self.total_prompt_tokens,
-            "total_completion_tokens": self.total_completion_tokens,
-            "total_llm_tokens": self.total_prompt_tokens + self.total_completion_tokens,
-            "total_llm_calls": self.total_llm_calls
-        }
-
-# --- Pydantic Models for API ---
-class InitializeRequest(BaseModel):
-    api_key: str
-    document_content: Optional[str] = None
-
-class QueryRequest(BaseModel):
-    query: str
-    api_key: str
-
-class InitializeResponse(BaseModel):
-    success: bool
-    message: str
-    chunks: Optional[int] = None
-    estimated_tokens: Optional[int] = None
-
-class QueryResponse(BaseModel):
-    success: bool
-    answer: str
-    response_time: float
-    query_tokens: int
-    llm_tokens: Dict[str, int]
-    session_stats: Dict[str, int]
-
-class StatsResponse(BaseModel):
-    total_queries: int
-    total_embedding_tokens: int
-    total_llm_tokens: int
-    total_llm_calls: int
-    initialization_complete: bool
-
-# --- Global Variables ---
-class RAGSystem:
-    def __init__(self):
-        self.vector_store = None
-        self.qa_chain = None
-        self.token_callback_handler = TokenUsageCallbackHandler()
-        self.session_stats = {
-            "total_queries": 0,
-            "total_embedding_tokens": 0,
-            "initialization_complete": False
-        }
-        self.current_api_key = None
-
-# Global RAG system instance
-rag_system = RAGSystem()
-
-def initialize_rag_system(api_key, file_content=None):
-    """Initialize the RAG system with API key and optional file content."""
-    global rag_system
-    
-    try:
-        # Set API key
-        os.environ["GOOGLE_API_KEY"] = api_key
-        rag_system.current_api_key = api_key
-        
-        # Initialize embeddings
-        embeddings = GoogleGenerativeAIEmbeddings(
-            model="models/embedding-001",
-            google_api_key=api_key
-        )
-        
-        # Initialize LLM
-        llm = ChatGoogleGenerativeAI(
-            model="gemini-1.5-flash",
-            google_api_key=api_key,
-            temperature=TEMPERATURE,
-            max_tokens=MAX_TOKENS,
-            callbacks=[rag_system.token_callback_handler],
-            verbose=False
-        )
-        
-        # Load or use default document
-        if file_content:
-            # Save uploaded file content
-            with open("uploaded_document.txt", "w", encoding="utf-8") as f:
-                f.write(file_content)
-            loader = TextLoader("uploaded_document.txt")
-        else:
-            # Check if default maize_data.txt exists
-            if os.path.exists("maize_data.txt"):
-                loader = TextLoader("maize_data.txt")
-            else:
-                return "❌ No document found. Please upload a file or ensure maize_data.txt exists."
-        
-        # Load and split documents
-        documents = loader.load()
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=CHUNK_SIZE, 
-            chunk_overlap=CHUNK_OVERLAP
-        )
-        chunks = text_splitter.split_documents(documents)
-        
-        # Estimate embedding tokens
-        initial_embedding_tokens = sum(estimate_tokens(chunk.page_content) for chunk in chunks)
-        rag_system.session_stats["total_embedding_tokens"] = initial_embedding_tokens
-        
-        # Create vector store
-        rag_system.vector_store = FAISS.from_documents(chunks, embeddings)
-        
-        # Create prompt template
-        prompt_template = PromptTemplate(
-            input_variables=["context", "question"],
-            template="""
-You are an expert in maize agriculture. Use the following context ONLY to answer the question accurately and helpfully. If the context doesn't contain the answer, say "Based on the provided context, I cannot answer this question.".
-
-Context:
-{context}
-
-Question: {question}
-
-Answer:"""
-        )
-        
-        # Set up QA chain
-        rag_system.qa_chain = RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",
-            retriever=rag_system.vector_store.as_retriever(search_kwargs={"k": RETRIEVAL_K}),
-            chain_type_kwargs={"prompt": prompt_template},
-            callbacks=[rag_system.token_callback_handler],
-            return_source_documents=True
-        )
-        
-        rag_system.session_stats["initialization_complete"] = True
-        
-        return f"✅ RAG system initialized successfully!\n📄 Document processed: {len(chunks)} chunks\n🔢 Estimated embedding tokens: ~{initial_embedding_tokens}"
-        
-    except Exception as e:
-        logger.error(f"Initialization failed: {str(e)}")
-        return f"❌ Initialization failed: {str(e)}"
-
-def process_query(query, api_key):
-    """Process a user query through the RAG system."""
-    global rag_system
-    
-    if not api_key:
-        return "❌ Please provide a Google API key first.", ""
-    
-    if not rag_system.qa_chain:
-        return "❌ RAG system not initialized. Please initialize first.", ""
-    
-    if not query.strip():
-        return "❌ Please enter a question.", ""
-    
-    try:
-        # Estimate query embedding tokens
-        query_tokens = estimate_tokens(query)
-        rag_system.session_stats["total_embedding_tokens"] += query_tokens
-        rag_system.session_stats["total_queries"] += 1
-        
-        # Process query
-        start_time = time.time()
-        result = rag_system.qa_chain({"query": query})
-        end_time = time.time()
-        
-        # Get token usage
-        llm_tokens = rag_system.token_callback_handler.get_total_tokens()
-        
-        # Format response
-        answer = result['result']
-        
-        # Create stats summary
-        stats = f"""
-📊 **Query Statistics:**
-- Response time: {end_time - start_time:.2f} seconds
-- Query tokens (estimated): ~{query_tokens}
-- LLM tokens (this query): Prompt: {llm_tokens['total_prompt_tokens']}, Completion: {llm_tokens['total_completion_tokens']}
-
-📈 **Session Statistics:**
-- Total queries: {rag_system.session_stats['total_queries']}
-- Total embedding tokens: ~{rag_system.session_stats['total_embedding_tokens']}
-- Total LLM calls: {llm_tokens['total_llm_calls']}
-- Total LLM tokens: {llm_tokens['total_llm_tokens']}
-"""
-        
-        return answer, stats
-        
-    except Exception as e:
-        logger.error(f"Error processing query: {str(e)}")
-        return f"❌ Error processing query: {str(e)}", ""
-
-def upload_file_and_initialize(api_key, file):
-    """Handle file upload and system initialization."""
-    if not api_key:
-        return "❌ Please provide a Google API key first."
-    
-    if file is None:
-        return initialize_rag_system(api_key)
-    
-    try:
-        # Read uploaded file
-        file_content = file.decode('utf-8')
-        return initialize_rag_system(api_key, file_content)
-    except Exception as e:
-        return f"❌ Error reading uploaded file: {str(e)}"
-
-def reset_session():
-    """Reset the session statistics."""
-    global rag_system
-    rag_system.token_callback_handler.reset_counters()
-    rag_system.session_stats = {
-        "total_queries": 0,
-        "total_embedding_tokens": 0,
-        "initialization_complete": False
-    }
-    return "🔄 Session statistics reset."
-
-# --- FastAPI Setup ---
-app = FastAPI(
-    title="Maize RAG Q&A System API",
-    description="API for the Maize Agriculture RAG Q&A System",
-    version="1.0.0"
-)
-
-# Optional: Add API key authentication for API endpoints
-security = HTTPBearer(auto_error=False)
-
-async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    """Extract API key from Authorization header (optional)"""
-    if credentials:
-        return credentials.credentials
-    return None
-
-# --- API Endpoints ---
-
-@app.get("/")
-async def root():
-    """Root endpoint"""
-    return {"message": "Maize RAG Q&A System API", "status": "running"}
-
-@app.get("/health")
-async def health_check():
-    """Health check endpoint"""
-    return {
-        "status": "healthy",
-        "system_initialized": rag_system.session_stats["initialization_complete"]
-    }
-
-@app.post("/initialize", response_model=InitializeResponse)
-async def initialize_system(request: InitializeRequest):
-    """Initialize the RAG system"""
-    try:
-        result = initialize_rag_system(request.api_key, request.document_content)
-        
-        if "✅" in result:
-            # Parse successful result
-            lines = result.split('\n')
-            chunks = None
-            tokens = None
-            
-            for line in lines:
-                if "chunks" in line:
-                    chunks = int(line.split(': ')[1].split(' ')[0])
-                elif "tokens" in line:
-                    tokens = int(line.split('~')[1])
-            
-            return InitializeResponse(
-                success=True,
-                message=result,
-                chunks=chunks,
-                estimated_tokens=tokens
-            )
-        else:
-            return InitializeResponse(
-                success=False,
-                message=result
-            )
-    
-    except Exception as e:
-        logger.error(f"API initialization error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/query", response_model=QueryResponse)
-async def query_system(request: QueryRequest):
-    """Query the RAG system"""
-    try:
-        if not rag_system.session_stats["initialization_complete"]:
-            raise HTTPException(status_code=400, detail="System not initialized")
-        
-        # Estimate query embedding tokens
-        query_tokens = estimate_tokens(request.query)
-        rag_system.session_stats["total_embedding_tokens"] += query_tokens
-        rag_system.session_stats["total_queries"] += 1
-        
-        # Process query
-        start_time = time.time()
-        result = rag_system.qa_chain({"query": request.query})
-        end_time = time.time()
-        
-        # Get token usage
-        llm_tokens = rag_system.token_callback_handler.get_total_tokens()
-        
-        response_time = end_time - start_time
-        
-        return QueryResponse(
-            success=True,
-            answer=result['result'],
-            response_time=response_time,
-            query_tokens=query_tokens,
-            llm_tokens=llm_tokens,
-            session_stats=rag_system.session_stats
-        )
-    
-    except Exception as e:
-        logger.error(f"API query error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/stats", response_model=StatsResponse)
-async def get_stats():
-    """Get current session statistics"""
-    llm_tokens = rag_system.token_callback_handler.get_total_tokens()
-    
-    return StatsResponse(
-        total_queries=rag_system.session_stats["total_queries"],
-        total_embedding_tokens=rag_system.session_stats["total_embedding_tokens"],
-        total_llm_tokens=llm_tokens["total_llm_tokens"],
-        total_llm_calls=llm_tokens["total_llm_calls"],
-        initialization_complete=rag_system.session_stats["initialization_complete"]
-    )
-
-@app.post("/reset")
-async def reset_system():
-    """Reset session statistics"""
-    reset_session()
-    return {"message": "Session reset successfully"}
-
-@app.post("/upload-document")
-async def upload_document(
-    file: UploadFile = File(...),
-    api_key: str = None
-):
-    """Upload a document and initialize the system"""
-    try:
-        if not api_key:
-            raise HTTPException(status_code=400, detail="API key required")
-        
-        # Read uploaded file
-        content = await file.read()
-        file_content = content.decode('utf-8')
-        
-        # Initialize system with uploaded content
-        result = initialize_rag_system(api_key, file_content)
-        
-        if "✅" in result:
-            return {"success": True, "message": result}
-        else:
-            return {"success": False, "message": result}
-    
-    except Exception as e:
-        logger.error(f"Document upload error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-# Create Gradio interface
-def create_interface():
-    with gr.Blocks(title="Maize RAG Q&A System", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🌽 Maize Agriculture RAG Q&A System
-        
-        This system uses Retrieval-Augmented Generation (RAG) to answer questions about maize agriculture.
-        Upload your own document or use the default maize dataset.
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=2):
-                api_key_input = gr.Textbox(
-                    label="🔑 Google API Key",
-                    placeholder="Enter your Google Generative AI API key",
-                    type="password",
-                    info="Get your API key from Google AI Studio"
-                )
-            
-            with gr.Column(scale=1):
-                reset_btn = gr.Button("🔄 Reset Session", variant="secondary")
-        
-        with gr.Row():
-            with gr.Column():
-                file_upload = gr.File(
-                    label="📁 Upload Document (Optional)",
-                    file_types=[".txt"],
-                    info="Upload a text file or use the default maize dataset"
-                )
-                
-                init_btn = gr.Button("🚀 Initialize RAG System", variant="primary")
-                init_output = gr.Textbox(
-                    label="📋 Initialization Status",
-                    lines=3,
-                    interactive=False
-                )
-        
-        gr.Markdown("## 💬 Ask Questions")
-        
-        with gr.Row():
-            with gr.Column(scale=3):
-                query_input = gr.Textbox(
-                    label="❓ Your Question",
-                    placeholder="Ask something about maize agriculture...",
-                    lines=2
-                )
-                
-                # Sample questions
-                sample_questions = [
-                    "What are the main pests affecting maize crops?",
-                    "How should maize be irrigated?",
-                    "What is the ideal soil type for maize?",
-                    "What are the nutritional requirements of maize?",
-                    "When is the best time to harvest maize?"
-                ]
-                
-                gr.Examples(
-                    examples=sample_questions,
-                    inputs=query_input,
-                    label="💡 Sample Questions"
-                )
-            
-            with gr.Column(scale=1):
-                submit_btn = gr.Button("🔍 Ask", variant="primary")
-        
-        with gr.Row():
-            with gr.Column(scale=2):
-                answer_output = gr.Textbox(
-                    label="🤖 Answer",
-                    lines=6,
-                    interactive=False
-                )
-            
-            with gr.Column(scale=1):
-                stats_output = gr.Markdown(
-                    label="📊 Statistics",
-                    value="Statistics will appear here after queries."
-                )
-        
-        # Event handlers
-        init_btn.click(
-            upload_file_and_initialize,
-            inputs=[api_key_input, file_upload],
-            outputs=init_output
-        )
-        
-        submit_btn.click(
-            process_query,
-            inputs=[query_input, api_key_input],
-            outputs=[answer_output, stats_output]
-        )
-        
-        query_input.submit(
-            process_query,
-            inputs=[query_input, api_key_input],
-            outputs=[answer_output, stats_output]
-        )
-        
-        reset_btn.click(
-            reset_session,
-            outputs=init_output
-        )
-        
-        gr.Markdown("""
-        ## 📝 Instructions:
-        1. **Enter your Google API Key** (required)
-        2. **Upload a document** (optional - uses default maize dataset if not provided)
-        3. **Initialize the RAG system** by clicking "Initialize RAG System"
-        4. **Ask questions** about the document content
-        5. **View statistics** to monitor token usage and costs
-        
-        ## 💰 Cost Information:
-        - **Gemini 1.5 Flash**: Input: $0.075/1M tokens, Output: $0.30/1M tokens
-        - **Embedding Model**: $0.025/1M tokens
-        
-        Token usage is estimated and displayed for cost tracking.
-        """)
-    
-    return demo
-
-# Create and launch the interface
-def run_gradio():
-    """Run Gradio interface"""
-    demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        quiet=True  # Reduce Gradio logs in combined mode
-    )
-
-def run_fastapi():
-    """Run FastAPI server"""
-    uvicorn.run(
-        app, 
-        host="0.0.0.0", 
-        port=8000,
-        log_level="info"
-    )
-
-if __name__ == "__main__":
-    import sys
-    
-    if len(sys.argv) > 1:
-        mode = sys.argv[1]
-        
-        if mode == "api":
-            # Run only FastAPI
-            print("Starting FastAPI server on port 8000...")
-            run_fastapi()
-        elif mode == "gradio":
-            # Run only Gradio
-            print("Starting Gradio interface on port 7860...")
-            run_gradio()
-        elif mode == "both":
-            # Run both servers
-            print("Starting both FastAPI (port 8000) and Gradio (port 7860)...")
-            
-            # Start FastAPI in a separate thread
-            fastapi_thread = threading.Thread(target=run_fastapi)
-            fastapi_thread.daemon = True
-            fastapi_thread.start()
-            
-            # Start Gradio in main thread
-            time.sleep(2)  # Give FastAPI time to start
-            run_gradio()
-        else:
-            print("Usage: python app.py [api|gradio|both]")
-            print("Default: gradio only")
-            run_gradio()
-    else:
-        # Default: run only Gradio (for Hugging Face Spaces compatibility)
-        print("Starting Gradio interface on port 7860...")
+import os
+import time
+import gradio as gr
+import uvicorn
+from fastapi import FastAPI, HTTPException, Depends, File, UploadFile
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel
+from typing import Optional, Dict, Any
+import threading
+import logging
+from langchain_community.document_loaders import TextLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
+import tiktoken
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# --- Configuration ---
+CHUNK_SIZE = 800
+CHUNK_OVERLAP = 100
+MAX_TOKENS = 512
+TEMPERATURE = 0.5
+RETRIEVAL_K = 5
+
+# --- Token Counting Setup ---
+try:
+    tokenizer = tiktoken.get_encoding("cl100k_base")
+except:
+    print("Tiktoken encoder 'cl100k_base' not found. Using basic split().")
+    tokenizer = type('obj', (object,), {'encode': lambda x: x.split()})()
+
+def estimate_tokens(text):
+    """Estimates token count for a given text."""
+    return len(tokenizer.encode(text))
+
+# Custom Callback Handler to track LLM token usage
+class TokenUsageCallbackHandler(BaseCallbackHandler):
+    """Callback handler to track token usage in LLM calls."""
+    def __init__(self):
+        super().__init__()
+        self.reset_counters()
+
+    def reset_counters(self):
+        self.total_prompt_tokens = 0
+        self.total_completion_tokens = 0
+        self.total_llm_calls = 0
+
+    def on_llm_end(self, response, **kwargs):
+        """Collect token usage from the LLM response."""
+        self.total_llm_calls += 1
+        llm_output = response.llm_output
+        
+        if llm_output and 'usage_metadata' in llm_output:
+            usage = llm_output['usage_metadata']
+            prompt_tokens = usage.get('prompt_token_count', 0)
+            completion_tokens = usage.get('candidates_token_count', 0)
+            
+            self.total_prompt_tokens += prompt_tokens
+            self.total_completion_tokens += completion_tokens
+
+    def get_total_tokens(self):
+        """Returns the total prompt and completion tokens."""
+        return {
+            "total_prompt_tokens": self.total_prompt_tokens,
+            "total_completion_tokens": self.total_completion_tokens,
+            "total_llm_tokens": self.total_prompt_tokens + self.total_completion_tokens,
+            "total_llm_calls": self.total_llm_calls
+        }
+
+# --- Pydantic Models for API ---
+class InitializeRequest(BaseModel):
+    api_key: str
+    document_content: Optional[str] = None
+
+class QueryRequest(BaseModel):
+    query: str
+    api_key: str
+
+class InitializeResponse(BaseModel):
+    success: bool
+    message: str
+    chunks: Optional[int] = None
+    estimated_tokens: Optional[int] = None
+
+class QueryResponse(BaseModel):
+    success: bool
+    answer: str
+    response_time: float
+    query_tokens: int
+    llm_tokens: Dict[str, int]
+    session_stats: Dict[str, int]
+
+class StatsResponse(BaseModel):
+    total_queries: int
+    total_embedding_tokens: int
+    total_llm_tokens: int
+    total_llm_calls: int
+    initialization_complete: bool
+
+# --- Global Variables ---
+class RAGSystem:
+    def __init__(self):
+        self.vector_store = None
+        self.qa_chain = None
+        self.token_callback_handler = TokenUsageCallbackHandler()
+        self.session_stats = {
+            "total_queries": 0,
+            "total_embedding_tokens": 0,
+            "initialization_complete": False
+        }
+        self.current_api_key = None
+
+# Global RAG system instance
+rag_system = RAGSystem()
+
+def initialize_rag_system(api_key, file_content=None):
+    """Initialize the RAG system with API key and optional file content."""
+    global rag_system
+    
+    try:
+        # Set API key
+        os.environ["GOOGLE_API_KEY"] = api_key
+        rag_system.current_api_key = api_key
+        
+        # Initialize embeddings
+        embeddings = GoogleGenerativeAIEmbeddings(
+            model="models/embedding-001",
+            google_api_key=api_key
+        )
+        
+        # Initialize LLM
+        llm = ChatGoogleGenerativeAI(
+            model="gemini-1.5-flash",
+            google_api_key=api_key,
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+            callbacks=[rag_system.token_callback_handler],
+            verbose=False
+        )
+        
+        # Load or use default document
+        if file_content:
+            # Save uploaded file content
+            with open("uploaded_document.txt", "w", encoding="utf-8") as f:
+                f.write(file_content)
+            loader = TextLoader("uploaded_document.txt")
+        else:
+            # Check if default maize_data.txt exists
+            if os.path.exists("maize_data.txt"):
+                loader = TextLoader("maize_data.txt")
+            else:
+                return "❌ No document found. Please upload a file or ensure maize_data.txt exists."
+        
+        # Load and split documents
+        documents = loader.load()
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=CHUNK_SIZE, 
+            chunk_overlap=CHUNK_OVERLAP
+        )
+        chunks = text_splitter.split_documents(documents)
+        
+        # Estimate embedding tokens
+        initial_embedding_tokens = sum(estimate_tokens(chunk.page_content) for chunk in chunks)
+        rag_system.session_stats["total_embedding_tokens"] = initial_embedding_tokens
+        
+        # Create vector store
+        rag_system.vector_store = FAISS.from_documents(chunks, embeddings)
+        
+        # Create prompt template
+        prompt_template = PromptTemplate(
+            input_variables=["context", "question"],
+            template="""
+You are an expert in maize agriculture. Use the following context ONLY to answer the question accurately and helpfully. If the context doesn't contain the answer, say "Based on the provided context, I cannot answer this question.".
+
+Context:
+{context}
+
+Question: {question}
+
+Answer:"""
+        )
+        
+        # Set up QA chain
+        rag_system.qa_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=rag_system.vector_store.as_retriever(search_kwargs={"k": RETRIEVAL_K}),
+            chain_type_kwargs={"prompt": prompt_template},
+            callbacks=[rag_system.token_callback_handler],
+            return_source_documents=True
+        )
+        
+        rag_system.session_stats["initialization_complete"] = True
+        
+        return f"✅ RAG system initialized successfully!\n📄 Document processed: {len(chunks)} chunks\n🔢 Estimated embedding tokens: ~{initial_embedding_tokens}"
+        
+    except Exception as e:
+        logger.error(f"Initialization failed: {str(e)}")
+        return f"❌ Initialization failed: {str(e)}"
+
+def process_query(query, api_key):
+    """Process a user query through the RAG system."""
+    global rag_system
+    
+    if not api_key:
+        return "❌ Please provide a Google API key first.", ""
+    
+    if not rag_system.qa_chain:
+        return "❌ RAG system not initialized. Please initialize first.", ""
+    
+    if not query.strip():
+        return "❌ Please enter a question.", ""
+    
+    try:
+        # Estimate query embedding tokens
+        query_tokens = estimate_tokens(query)
+        rag_system.session_stats["total_embedding_tokens"] += query_tokens
+        rag_system.session_stats["total_queries"] += 1
+        
+        # Process query
+        start_time = time.time()
+        result = rag_system.qa_chain({"query": query})
+        end_time = time.time()
+        
+        # Get token usage
+        llm_tokens = rag_system.token_callback_handler.get_total_tokens()
+        
+        # Format response
+        answer = result['result']
+        
+        # Create stats summary
+        stats = f"""
+📊 **Query Statistics:**
+- Response time: {end_time - start_time:.2f} seconds
+- Query tokens (estimated): ~{query_tokens}
+- LLM tokens (this query): Prompt: {llm_tokens['total_prompt_tokens']}, Completion: {llm_tokens['total_completion_tokens']}
+
+📈 **Session Statistics:**
+- Total queries: {rag_system.session_stats['total_queries']}
+- Total embedding tokens: ~{rag_system.session_stats['total_embedding_tokens']}
+- Total LLM calls: {llm_tokens['total_llm_calls']}
+- Total LLM tokens: {llm_tokens['total_llm_tokens']}
+"""
+        
+        return answer, stats
+        
+    except Exception as e:
+        logger.error(f"Error processing query: {str(e)}")
+        return f"❌ Error processing query: {str(e)}", ""
+
+def upload_file_and_initialize(api_key, file):
+    """Handle file upload and system initialization."""
+    if not api_key:
+        return "❌ Please provide a Google API key first."
+    
+    if file is None:
+        return initialize_rag_system(api_key)
+    
+    try:
+        # Handle different file object types based on Gradio version
+        if hasattr(file, 'name'):
+            # Newer Gradio versions - file has .name attribute
+            with open(file.name, 'r', encoding='utf-8') as f:
+                file_content = f.read()
+        elif isinstance(file, str):
+            # File path as string
+            with open(file, 'r', encoding='utf-8') as f:
+                file_content = f.read()
+        elif hasattr(file, 'read'):
+            # File-like object
+            file_content = file.read()
+            if isinstance(file_content, bytes):
+                file_content = file_content.decode('utf-8')
+        else:
+            # Fallback - try to read as bytes and decode
+            file_content = file.decode('utf-8') if isinstance(file, bytes) else str(file)
+            
+        return initialize_rag_system(api_key, file_content)
+        
+    except Exception as e:
+        logger.error(f"Error reading uploaded file: {str(e)}")
+        return f"❌ Error reading uploaded file: {str(e)}"
+
+def reset_session():
+    """Reset the session statistics."""
+    global rag_system
+    rag_system.token_callback_handler.reset_counters()
+    rag_system.session_stats = {
+        "total_queries": 0,
+        "total_embedding_tokens": 0,
+        "initialization_complete": False
+    }
+    return "🔄 Session statistics reset."
+
+# --- FastAPI Setup ---
+app = FastAPI(
+    title="Maize RAG Q&A System API",
+    description="API for the Maize Agriculture RAG Q&A System",
+    version="1.0.0"
+)
+
+# Optional: Add API key authentication for API endpoints
+security = HTTPBearer(auto_error=False)
+
+async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """Extract API key from Authorization header (optional)"""
+    if credentials:
+        return credentials.credentials
+    return None
+
+# --- API Endpoints ---
+
+@app.get("/")
+async def root():
+    """Root endpoint"""
+    return {"message": "Maize RAG Q&A System API", "status": "running"}
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "system_initialized": rag_system.session_stats["initialization_complete"]
+    }
+
+@app.post("/initialize", response_model=InitializeResponse)
+async def initialize_system(request: InitializeRequest):
+    """Initialize the RAG system"""
+    try:
+        result = initialize_rag_system(request.api_key, request.document_content)
+        
+        if "✅" in result:
+            # Parse successful result
+            lines = result.split('\n')
+            chunks = None
+            tokens = None
+            
+            for line in lines:
+                if "chunks" in line:
+                    chunks = int(line.split(': ')[1].split(' ')[0])
+                elif "tokens" in line:
+                    tokens = int(line.split('~')[1])
+            
+            return InitializeResponse(
+                success=True,
+                message=result,
+                chunks=chunks,
+                estimated_tokens=tokens
+            )
+        else:
+            return InitializeResponse(
+                success=False,
+                message=result
+            )
+    
+    except Exception as e:
+        logger.error(f"API initialization error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/query", response_model=QueryResponse)
+async def query_system(request: QueryRequest):
+    """Query the RAG system"""
+    try:
+        if not rag_system.session_stats["initialization_complete"]:
+            raise HTTPException(status_code=400, detail="System not initialized")
+        
+        # Estimate query embedding tokens
+        query_tokens = estimate_tokens(request.query)
+        rag_system.session_stats["total_embedding_tokens"] += query_tokens
+        rag_system.session_stats["total_queries"] += 1
+        
+        # Process query
+        start_time = time.time()
+        result = rag_system.qa_chain({"query": request.query})
+        end_time = time.time()
+        
+        # Get token usage
+        llm_tokens = rag_system.token_callback_handler.get_total_tokens()
+        
+        response_time = end_time - start_time
+        
+        return QueryResponse(
+            success=True,
+            answer=result['result'],
+            response_time=response_time,
+            query_tokens=query_tokens,
+            llm_tokens=llm_tokens,
+            session_stats=rag_system.session_stats
+        )
+    
+    except Exception as e:
+        logger.error(f"API query error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/stats", response_model=StatsResponse)
+async def get_stats():
+    """Get current session statistics"""
+    llm_tokens = rag_system.token_callback_handler.get_total_tokens()
+    
+    return StatsResponse(
+        total_queries=rag_system.session_stats["total_queries"],
+        total_embedding_tokens=rag_system.session_stats["total_embedding_tokens"],
+        total_llm_tokens=llm_tokens["total_llm_tokens"],
+        total_llm_calls=llm_tokens["total_llm_calls"],
+        initialization_complete=rag_system.session_stats["initialization_complete"]
+    )
+
+@app.post("/reset")
+async def reset_system():
+    """Reset session statistics"""
+    reset_session()
+    return {"message": "Session reset successfully"}
+
+@app.post("/upload-document")
+async def upload_document(
+    file: UploadFile = File(...),
+    api_key: str = None
+):
+    """Upload a document and initialize the system"""
+    try:
+        if not api_key:
+            raise HTTPException(status_code=400, detail="API key required")
+        
+        # Read uploaded file
+        content = await file.read()
+        file_content = content.decode('utf-8')
+        
+        # Initialize system with uploaded content
+        result = initialize_rag_system(api_key, file_content)
+        
+        if "✅" in result:
+            return {"success": True, "message": result}
+        else:
+            return {"success": False, "message": result}
+    
+    except Exception as e:
+        logger.error(f"Document upload error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Create Gradio interface with version compatibility
+def create_interface():
+    # Check Gradio version for compatibility
+    import gradio as gr
+    gradio_version = gr.__version__
+    
+    with gr.Blocks(title="Maize RAG Q&A System", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🌽 Maize Agriculture RAG Q&A System
+        
+        This system uses Retrieval-Augmented Generation (RAG) to answer questions about maize agriculture.
+        Upload your own document or use the default maize dataset.
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                api_key_input = gr.Textbox(
+                    label="🔑 Google API Key",
+                    placeholder="Enter your Google Generative AI API key",
+                    type="password"
+                )
+                gr.Markdown("Get your API key from Google AI Studio")
+            
+            with gr.Column(scale=1):
+                reset_btn = gr.Button("🔄 Reset Session", variant="secondary")
+        
+        with gr.Row():
+            with gr.Column():
+                file_upload = gr.File(
+                    label="📁 Upload Document (Optional)",
+                    file_types=[".txt"]
+                )
+                gr.Markdown("Upload a text file or use the default maize dataset")
+                
+                init_btn = gr.Button("🚀 Initialize RAG System", variant="primary")
+                init_output = gr.Textbox(
+                    label="📋 Initialization Status",
+                    lines=3,
+                    interactive=False
+                )
+        
+        gr.Markdown("## 💬 Ask Questions")
+        
+        with gr.Row():
+            with gr.Column(scale=3):
+                query_input = gr.Textbox(
+                    label="❓ Your Question",
+                    placeholder="Ask something about maize agriculture...",
+                    lines=2
+                )
+                
+                # Sample questions
+                sample_questions = [
+                    "What are the main pests affecting maize crops?",
+                    "How should maize be irrigated?",
+                    "What is the ideal soil type for maize?",
+                    "What are the nutritional requirements of maize?",
+                    "When is the best time to harvest maize?"
+                ]
+                
+                # Use Examples component if available, otherwise just show as markdown
+                try:
+                    gr.Examples(
+                        examples=sample_questions,
+                        inputs=query_input,
+                        label="💡 Sample Questions"
+                    )
+                except:
+                    gr.Markdown("💡 **Sample Questions:**\n" + 
+                               "\n".join([f"- {q}" for q in sample_questions]))
+            
+            with gr.Column(scale=1):
+                submit_btn = gr.Button("🔍 Ask", variant="primary")
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                answer_output = gr.Textbox(
+                    label="🤖 Answer",
+                    lines=6,
+                    interactive=False
+                )
+            
+            with gr.Column(scale=1):
+                stats_output = gr.Markdown(
+                    value="📊 Statistics will appear here after queries."
+                )
+        
+        # Event handlers
+        init_btn.click(
+            upload_file_and_initialize,
+            inputs=[api_key_input, file_upload],
+            outputs=init_output
+        )
+        
+        submit_btn.click(
+            process_query,
+            inputs=[query_input, api_key_input],
+            outputs=[answer_output, stats_output]
+        )
+        
+        query_input.submit(
+            process_query,
+            inputs=[query_input, api_key_input],
+            outputs=[answer_output, stats_output]
+        )
+        
+        reset_btn.click(
+            reset_session,
+            outputs=init_output
+        )
+        
+        gr.Markdown("""
+        ## 📝 Instructions:
+        1. **Enter your Google API Key** (required)
+        2. **Upload a document** (optional - uses default maize dataset if not provided)
+        3. **Initialize the RAG system** by clicking "Initialize RAG System"
+        4. **Ask questions** about the document content
+        5. **View statistics** to monitor token usage and costs
+        
+        ## 💰 Cost Information:
+        - **Gemini 1.5 Flash**: Input: $0.075/1M tokens, Output: $0.30/1M tokens
+        - **Embedding Model**: $0.025/1M tokens
+        
+        Token usage is estimated and displayed for cost tracking.
+        
+        ## 🔗 API Access:
+        This system also provides REST API endpoints:
+        - **API Docs**: Add `/docs` to the URL for interactive API documentation
+        - **Health Check**: `GET /health`
+        - **Initialize**: `POST /initialize`
+        - **Query**: `POST /query`
+        """)
+    
+    return demo
+
+# Create and launch the interface
+def run_gradio():
+    """Run Gradio interface"""
+    demo = create_interface()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        quiet=True  # Reduce Gradio logs in combined mode
+    )
+
+def run_fastapi():
+    """Run FastAPI server"""
+    uvicorn.run(
+        app, 
+        host="0.0.0.0", 
+        port=8000,
+        log_level="info"
+    )
+
+if __name__ == "__main__":
+    import sys
+    
+    if len(sys.argv) > 1:
+        mode = sys.argv[1]
+        
+        if mode == "api":
+            # Run only FastAPI
+            print("Starting FastAPI server on port 8000...")
+            run_fastapi()
+        elif mode == "gradio":
+            # Run only Gradio
+            print("Starting Gradio interface on port 7860...")
+            run_gradio()
+        elif mode == "both":
+            # Run both servers
+            print("Starting both FastAPI (port 8000) and Gradio (port 7860)...")
+            
+            # Start FastAPI in a separate thread
+            fastapi_thread = threading.Thread(target=run_fastapi)
+            fastapi_thread.daemon = True
+            fastapi_thread.start()
+            
+            # Start Gradio in main thread
+            time.sleep(2)  # Give FastAPI time to start
+            run_gradio()
+        else:
+            print("Usage: python app.py [api|gradio|both]")
+            print("Default: gradio only")
+            run_gradio()
+    else:
+        # Default: run only Gradio (for Hugging Face Spaces compatibility)
+        print("Starting Gradio interface on port 7860...")
         run_gradio()