Create backend/main.py

backend/main.py

@@ -0,0 +1,308 @@
+"""
+RAG + Voice + MCP Chatbot Backend - FastAPI
+LangChain + LangGraph + Groq (UPDATED)
+"""
+
+import os
+import json
+import logging
+from typing import Optional, List, TypedDict
+
+import numpy as np
+from fastapi import FastAPI, HTTPException, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse, FileResponse
+from pydantic import BaseModel
+from dotenv import load_dotenv
+
+# LangChain / LangGraph
+from langchain_groq import ChatGroq
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+from langgraph.graph import StateGraph, END
+
+load_dotenv()
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# =========================================================
+# FastAPI App
+# =========================================================
+app = FastAPI(title="Voice RAG Chatbot", version="2.0.0")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# =========================================================
+# RAG Engine (unchanged, your implementation is good)
+# =========================================================
+class RAGEngine:
+    def __init__(self):
+        self.documents: List[dict] = []
+        self.embeddings: Optional[np.ndarray] = None
+        self.model = None
+        self._initialized = False
+
+    def _lazy_init(self):
+        if self._initialized:
+            return
+        try:
+            from sentence_transformers import SentenceTransformer
+            self.model = SentenceTransformer("all-MiniLM-L6-v2")
+            self._initialized = True
+            logger.info("RAG engine initialized")
+        except Exception as e:
+            logger.error(f"RAG init error: {e}")
+
+    def add_documents(self, texts: List[str], source: str = "upload"):
+        self._lazy_init()
+        if not self.model:
+            return False
+
+        chunks = []
+        for text in texts:
+            for i in range(0, len(text), 450):
+                chunk = text[i:i + 450].strip()
+                if chunk:
+                    chunks.append({"text": chunk, "source": source})
+
+        if not chunks:
+            return False
+
+        new_embeddings = self.model.encode([c["text"] for c in chunks])
+        self.documents.extend(chunks)
+
+        if self.embeddings is None:
+            self.embeddings = new_embeddings
+        else:
+            self.embeddings = np.vstack([self.embeddings, new_embeddings])
+
+        logger.info(f"Added {len(chunks)} chunks from {source}")
+        return True
+
+    def retrieve(self, query: str, top_k: int = 4) -> List[str]:
+        if not self.documents or self.embeddings is None:
+            return []
+
+        self._lazy_init()
+        if not self.model:
+            return []
+
+        q_emb = self.model.encode([query])
+
+        norms = np.linalg.norm(self.embeddings, axis=1, keepdims=True) + 1e-9
+        normed = self.embeddings / norms
+        q_norm = q_emb / (np.linalg.norm(q_emb) + 1e-9)
+
+        scores = normed @ q_norm.T
+        top_idx = np.argsort(scores[:, 0])[::-1][:top_k]
+        return [self.documents[i]["text"] for i in top_idx]
+
+    def clear(self):
+        self.documents = []
+        self.embeddings = None
+
+
+rag = RAGEngine()
+
+# =========================================================
+# MCP Tool Registry
+# =========================================================
+class MCPToolRegistry:
+    def __init__(self):
+        self.tools = {}
+
+    def register(self, name: str, description: str, handler):
+        self.tools[name] = {"description": description, "handler": handler}
+
+    def get_tool_descriptions(self) -> str:
+        return "\n".join(
+            f"- {name}: {info['description']}"
+            for name, info in self.tools.items()
+        )
+
+    async def run(self, name: str, args: dict) -> str:
+        if name not in self.tools:
+            return f"Tool '{name}' not found."
+        try:
+            return str(await self.tools[name]["handler"](**args))
+        except Exception as e:
+            return f"Tool error: {e}"
+
+
+mcp = MCPToolRegistry()
+
+# ---------------- Tools ----------------
+async def tool_current_time() -> str:
+    from datetime import datetime
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+
+async def tool_word_count(text: str = "") -> str:
+    return f"{len(text.split())} words, {len(text)} characters"
+
+
+async def tool_search_docs(query: str = "") -> str:
+    results = rag.retrieve(query, top_k=3)
+    if not results:
+        return "No documents found."
+    return "\n---\n".join(results[:2])
+
+
+mcp.register("get_time", "Get current time", tool_current_time)
+mcp.register("word_count", "Count words in text", tool_word_count)
+mcp.register("search_docs", "Search knowledge base", tool_search_docs)
+
+# =========================================================
+# Groq LLM via LangChain (FIXED MODEL)
+# =========================================================
+api_key = os.getenv("GROQ_API_KEY")
+if not api_key:
+    raise RuntimeError("❌ GROQ_API_KEY not set")
+
+llm = ChatGroq(
+    model="llama-3.1-8b-instant",  # ✅ NOT deprecated
+    groq_api_key=api_key,
+    temperature=0.7,
+)
+
+# =========================================================
+# LangGraph Agent
+# =========================================================
+class AgentState(TypedDict):
+    messages: List
+
+
+def agent_node(state: AgentState):
+    response = llm.invoke(state["messages"])
+    return {"messages": state["messages"] + [response]}
+
+
+graph = StateGraph(AgentState)
+graph.add_node("agent", agent_node)
+graph.set_entry_point("agent")
+graph.add_edge("agent", END)
+
+app_graph = graph.compile()
+
+# =========================================================
+# Conversation Store
+# =========================================================
+conversations: dict = {}
+
+# =========================================================
+# Schemas
+# =========================================================
+class ChatRequest(BaseModel):
+    message: str
+    session_id: str = "default"
+    use_rag: bool = True
+
+
+class ChatResponse(BaseModel):
+    reply: str
+    sources: List[str] = []
+    session_id: str
+
+
+# =========================================================
+# Routes
+# =========================================================
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    return FileResponse("frontend/index.html")
+
+
+@app.post("/chat", response_model=ChatResponse)
+async def chat(req: ChatRequest):
+    session_id = req.session_id
+    if session_id not in conversations:
+        conversations[session_id] = []
+
+    history = conversations[session_id]
+
+    # -------- RAG --------
+    rag_context = ""
+    sources = []
+
+    if req.use_rag and rag.documents:
+        retrieved = rag.retrieve(req.message)
+        if retrieved:
+            sources = retrieved[:3]
+            rag_context = "\n\n[Knowledge Base]\n" + "\n---\n".join(retrieved[:3])
+
+    # -------- Build messages --------
+    system_text = (
+        "You are a helpful AI assistant.\n"
+        f"{rag_context}\n\n"
+        f"Tools available:\n{mcp.get_tool_descriptions()}"
+    )
+
+    lc_messages = [SystemMessage(content=system_text)]
+
+    for h in history[-10:]:
+        if h["role"] == "user":
+            lc_messages.append(HumanMessage(content=h["content"]))
+        else:
+            lc_messages.append(AIMessage(content=h["content"]))
+
+    lc_messages.append(HumanMessage(content=req.message))
+
+    # -------- LangGraph invoke --------
+    try:
+        result = app_graph.invoke({"messages": lc_messages})
+        reply = result["messages"][-1].content
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    # -------- Save history --------
+    history.append({"role": "user", "content": req.message})
+    history.append({"role": "assistant", "content": reply})
+
+    conversations[session_id] = history[-40:]
+
+    return ChatResponse(
+        reply=reply,
+        sources=sources,
+        session_id=session_id,
+    )
+
+
+@app.post("/upload")
+async def upload_document(file: UploadFile = File(...)):
+    content = await file.read()
+    filename = file.filename or "document"
+    texts = []
+
+    try:
+        if filename.endswith(".pdf"):
+            import PyPDF2, io
+            reader = PyPDF2.PdfReader(io.BytesIO(content))
+            for page in reader.pages:
+                t = page.extract_text()
+                if t:
+                    texts.append(t)
+        else:
+            texts.append(content.decode("utf-8", errors="ignore"))
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    success = rag.add_documents(texts, source=filename)
+    if not success:
+        raise HTTPException(status_code=500, detail="Indexing failed")
+
+    return {"message": f"Indexed {filename}", "chunks": len(rag.documents)}
+
+
+@app.get("/status")
+async def status():
+    return {
+        "status": "online",
+        "rag_chunks": len(rag.documents),
+        "model": "llama-3.1-8b-instant",
+    }