Update RAG.py

RAG.py

@@ -1,310 +1,310 @@
-from langgraph.graph import StateGraph, START, END
-from typing import TypedDict, Annotated
-from langchain_groq import ChatGroq
-from langchain_openai import OpenAIEmbeddings, ChatOpenAI
-from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
-from langgraph.graph.message import add_messages
-from langchain_core.tools import tool
-from dotenv import load_dotenv
-from langgraph.checkpoint.memory import MemorySaver
-import os
-from langchain_community.vectorstores import FAISS
-from langchain_community.tools.tavily_search import TavilySearchResults
-
-load_dotenv()
-
-
-#===========================================
-# Load FAISS DB & Reload Logic [FEATURE ADDED]
-#===========================================
-
-FAISS_DB_PATH = "vectorstore/db_faiss"
-embeddings = OpenAIEmbeddings(model='text-embedding-3-small')
-
-# Global variable for the database
-db = None
-
-def reload_vector_store():
-    """
-    Reloads the FAISS index from disk. 
-    Call this function after a new file is ingested.
-    """
-    global db
-    if os.path.exists(FAISS_DB_PATH):
-        print(f"Loading FAISS from {FAISS_DB_PATH}...")
-        try:
-            db = FAISS.load_local(
-                FAISS_DB_PATH, 
-                embeddings, 
-                allow_dangerous_deserialization=True
-            )
-            print("Vector store loaded successfully.")
-        except Exception as e:
-            print(f"Error loading vector store: {e}")
-            db = None
-    else:
-        print("Warning: No Vector DB found. Please run ingestion first.")
-        db = None 
-
-# Initial Load
-reload_vector_store()
-
-
-#===========================================
-# Class Schema
-#===========================================
-
-class Ragbot_State(TypedDict):
-    query       :   str
-    context     :   list[str]
-    metadata    :   list[dict]
-    RAG         :   bool
-    web_search  :   bool
-    model_name  :   str
-    web_context :   str
-    response    :   Annotated[list[BaseMessage], add_messages]
-
-#===========================================
-# LLM'S
-#===========================================
-
-
-llm_kimi2       =   ChatGroq(model='moonshotai/kimi-k2-instruct-0905', streaming=True, temperature=0.4)
-llm_gpt         =   ChatOpenAI(model='gpt-4.1-nano', streaming=True, temperature=0.2)
-llm_gpt_oss     =   ChatGroq(model='openai/gpt-oss-120b', streaming=True, temperature=0.3)
-llm_lamma4      =   ChatGroq(model='meta-llama/llama-4-scout-17b-16e-instruct', streaming=True, temperature=0.5)
-llm_qwen3       =   ChatGroq(model='qwen/qwen3-32b', streaming=True, temperature=0.5)
-
-def get_llm(model_name: str):
-    if model_name == "kimi2":
-        return llm_kimi2
-    elif model_name == "gpt":
-        return llm_gpt
-    elif model_name == "gpt_oss":
-        return llm_gpt_oss
-    elif model_name == "lamma4":
-        return llm_lamma4
-    elif model_name == "qwen3":
-        return llm_qwen3
-    else:
-        return llm_gpt   # fallback if no match
-
-#===========================================
-# Search tool
-#===========================================
-
-@tool
-def tavily_search(query: str) -> dict:
-    """
-    Perform a real-time web search using Tavily.
-    """
-    try:
-        search = TavilySearchResults(max_results=2)
-        results = search.run(query)
-        return {"query": query, "results": results}
-    except Exception as e:
-        return {"error": str(e)}
-    
-#===========================================
-# fetching web context
-#===========================================
-
-def fetch_web_context(state: Ragbot_State):
-    user_query = state["query"]
-
-    enriched_query = f"""
-Fetch the latest, accurate, and up-to-date information about:
-{user_query}
-
-Focus on:
-- recent news
-- official announcements
-- verified sources
-- factual data
-"""
-
-    web_result = tavily_search.run(enriched_query)
-
-    return {
-        "web_context": str(web_result)
-    }
-
-#===========================================
-# db search
-#===========================================
-
-@tool
-def faiss_search(query: str) -> str:
-    """Search the FAISS vectorstore and return relevant documents."""
-    # Check global db variable
-    if db is None:
-        return "No documents have been uploaded yet.", []
-
-    try:
-        results = db.similarity_search(query, k=3)
-        context = "\n\n".join([doc.page_content for doc in results])
-        metadata = [doc.metadata for doc in results]
-        return context, metadata
-    except Exception as e:
-        return f"Error searching vector store: {str(e)}", []
-
-#===========================================
-# router
-#===========================================
-
-
-def router(state: Ragbot_State):
-    if state["RAG"]:
-        return "fetch_context"
-
-    if state["web_search"]:
-        return "fetch_web_context"
-
-    return "chat"
-
-#===========================================
-# fetching context
-#===========================================
-
-def fetch_context(state: Ragbot_State):
-    query = state["query"]
-    context, metadata = faiss_search.invoke({"query": query})
-    return {"context": [context], "metadata": [metadata]}
-
-
-#===========================================
-# system prompt
-#===========================================
-
-
-SYSTEM_PROMPT = SystemMessage(
-    content="""
-You are an intelligent conversational assistant and retrieval-augmented AI system built by Junaid.
-
-Your role is to:
-- Engage naturally in conversation like a friendly, helpful chatbot.
-- Answer general questions using your own knowledge when no external context is provided.
-- When relevant context is provided, use it accurately to answer user questions.
-- Seamlessly switch between casual conversation and knowledge-based answering.
-
-Guidelines:
-- If context is provided and relevant, use it as the primary source of truth.
-- If context is not provided or not relevant, respond using your general knowledge.
-- Do not hallucinate or invent information.
-- If you are unsure or the information is not available, clearly state that.
-- Be clear, concise, and helpful in all responses.
-- Maintain a natural, human-like conversational tone.
-- Never mention internal implementation details such as embeddings, vector databases, or system architecture.
-
-You are designed to provide reliable, accurate, and engaging assistance.
-"""
-)
-
-#===========================================
-# Chat function
-#===========================================
-
-def chat(state:Ragbot_State):
-    query = state['query']
-    context = state['context']
-    metadata = state['metadata']
-    web_context = state['web_context']
-    model_name = state.get('model_name', 'gpt')
-
-    history = state.get("response", [])
-
-    # [CHANGED] Updated Prompt to include History so it remembers your name
-    prompt = f"""
-You are an expert assistant designed to answer user questions using multiple information sources.
-
-Source Priority Rules (STRICT):
-1. **Conversation History**: Check if the answer was provided in previous messages (e.g., user's name, previous topics).
-2. If the provided Context contains the answer, use ONLY the Context.
-3. If the Context does not contain the answer and Web Context is available, use the Web Context.
-4. If neither Context nor Web Context contains the answer, use your general knowledge.
-5. Do NOT invent or hallucinate facts.
-6. If the answer cannot be determined, clearly say so.
-
-User Question:
-{query}
-
-Retrieved Context (Vector Database):
-{context}
-
-Metadata:
-{metadata}
-
-Web Context (Real-time Search):
-{web_context}
-
-Final Answer:
-"""
-
-    selected_llm = get_llm(model_name)
-    messages = [SYSTEM_PROMPT] + history + [HumanMessage(content=prompt)]
-    response = selected_llm.invoke(messages)
-    return {
-        'response': [
-            HumanMessage(content=query), 
-            response
-        ]
-    }
-
-#===========================================
-# Graph Declaration
-#===========================================
-
-# Keeping MemorySaver as requested (Note: RAM only, wipes on restart)
-memory = MemorySaver()
-graph = StateGraph(Ragbot_State)
-
-graph.add_node("fetch_context", fetch_context)
-graph.add_node("fetch_web_context", fetch_web_context)
-graph.add_node("chat", chat)
-
-graph.add_conditional_edges(
-    START,
-    router,
-    {
-        "fetch_context": "fetch_context",
-        "fetch_web_context": "fetch_web_context",
-        "chat": "chat"
-    }
-)
-
-graph.add_edge("fetch_context", "chat")
-graph.add_edge("fetch_web_context", "chat")
-graph.add_edge("chat", END)
-
-app = graph.compile(checkpointer=memory)
-
-
-#===========================================
-# Helper Function
-#===========================================
-
-def ask_bot(query: str, use_rag: bool = False, use_web: bool = False, thread_id: str = "1"):
-    config = {"configurable": {"thread_id": thread_id}}
-    inputs = {
-        "query": query,
-        "RAG": use_rag,
-        "web_search": use_web,
-        "context": [],
-        "metadata": [],
-        "web_context": "",
-    }
-    
-    result = app.invoke(inputs, config=config)
-    last_message = result['response'][-1]
-    
-    return last_message.content
-
-
-"""print("--- Conversation 1 ---")
-# User says hello and gives name
-response = ask_bot("Hi, my name is Junaid", thread_id="session_A")
-print(f"Bot: {response}")
-
-# User asks for name (RAG and Web are OFF)
-response = ask_bot("What is my name?", thread_id="session_A")
+from langgraph.graph import StateGraph, START, END
+from typing import TypedDict, Annotated
+from langchain_groq import ChatGroq
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
+from langgraph.graph.message import add_messages
+from langchain_core.tools import tool
+from dotenv import load_dotenv
+from langgraph.checkpoint.memory import MemorySaver
+import os
+from langchain_community.vectorstores import FAISS
+from langchain_community.tools.tavily_search import TavilySearchResults
+
+load_dotenv()
+
+
+#===========================================
+# Load FAISS DB & Reload Logic [FEATURE ADDED]
+#===========================================
+
+FAISS_DB_PATH = "vectorstore/db_faiss"
+embeddings = OpenAIEmbeddings(model='text-embedding-3-small')
+
+# Global variable for the database
+db = None
+
+def reload_vector_store():
+    """
+    Reloads the FAISS index from disk. 
+    Call this function after a new file is ingested.
+    """
+    global db
+    if os.path.exists(FAISS_DB_PATH):
+        print(f"Loading FAISS from {FAISS_DB_PATH}...")
+        try:
+            db = FAISS.load_local(
+                FAISS_DB_PATH, 
+                embeddings, 
+                allow_dangerous_deserialization=True
+            )
+            print("Vector store loaded successfully.")
+        except Exception as e:
+            print(f"Error loading vector store: {e}")
+            db = None
+    else:
+        print("Warning: No Vector DB found. Please run ingestion first.")
+        db = None 
+
+# Initial Load
+reload_vector_store()
+
+
+#===========================================
+# Class Schema
+#===========================================
+
+class Ragbot_State(TypedDict):
+    query       :   str
+    context     :   list[str]
+    metadata    :   list[dict]
+    RAG         :   bool
+    web_search  :   bool
+    model_name  :   str
+    web_context :   str
+    response    :   Annotated[list[BaseMessage], add_messages]
+
+#===========================================
+# LLM'S
+#===========================================
+
+
+llm_kimi2       =   ChatGroq(model='moonshotai/kimi-k2-instruct-0905', streaming=True, temperature=0.4)
+llm_gpt         =   ChatOpenAI(model='gpt-4.1-nano', streaming=True, temperature=0.2)
+llm_gpt_oss     =   ChatGroq(model='openai/gpt-oss-120b', streaming=True, temperature=0.3)
+llm_lamma4      =   ChatGroq(model='meta-llama/llama-4-scout-17b-16e-instruct', streaming=True, temperature=0.5)
+llm_qwen3       =   ChatGroq(model='qwen/qwen3-32b', streaming=True, temperature=0.5)
+
+def get_llm(model_name: str):
+    if model_name == "kimi2":
+        return llm_kimi2
+    elif model_name == "gpt":
+        return llm_gpt
+    elif model_name == "gpt_oss":
+        return llm_gpt_oss
+    elif model_name == "lamma4":
+        return llm_lamma4
+    elif model_name == "qwen3":
+        return llm_qwen3
+    else:
+        return llm_gpt   # fallback if no match
+
+#===========================================
+# Search tool
+#===========================================
+
+@tool
+def tavily_search(query: str) -> dict:
+    """
+    Perform a real-time web search using Tavily.
+    """
+    try:
+        search = TavilySearchResults(max_results=2)
+        results = search.run(query)
+        return {"query": query, "results": results}
+    except Exception as e:
+        return {"error": str(e)}
+    
+#===========================================
+# fetching web context
+#===========================================
+
+def fetch_web_context(state: Ragbot_State):
+    user_query = state["query"]
+
+    enriched_query = f"""
+Fetch the latest, accurate, and up-to-date information about:
+{user_query}
+
+Focus on:
+- recent news
+- official announcements
+- verified sources
+- factual data
+"""
+
+    web_result = tavily_search.run(enriched_query)
+
+    return {
+        "web_context": str(web_result)
+    }
+
+#===========================================
+# db search
+#===========================================
+
+@tool
+def faiss_search(query: str) -> str:
+    """Search the FAISS vectorstore and return relevant documents."""
+    # Check global db variable
+    if db is None:
+        return "No documents have been uploaded yet.", []
+
+    try:
+        results = db.similarity_search(query, k=3)
+        context = "\n\n".join([doc.page_content for doc in results])
+        metadata = [doc.metadata for doc in results]
+        return context, metadata
+    except Exception as e:
+        return f"Error searching vector store: {str(e)}", []
+
+#===========================================
+# router
+#===========================================
+
+
+def router(state: Ragbot_State):
+    if state["RAG"]:
+        return "fetch_context"
+
+    if state["web_search"]:
+        return "fetch_web_context"
+
+    return "chat"
+
+#===========================================
+# fetching context
+#===========================================
+
+def fetch_context(state: Ragbot_State):
+    query = state["query"]
+    context, metadata = faiss_search.invoke({"query": query})
+    return {"context": [context], "metadata": [metadata]}
+
+
+#===========================================
+# system prompt
+#===========================================
+
+
+SYSTEM_PROMPT = SystemMessage(
+    content="""
+You are an intelligent conversational assistant and retrieval-augmented AI system built by Junaid.
+
+Your role is to:
+- Engage naturally in conversation like a friendly, helpful chatbot.
+- Answer general questions using your own knowledge when no external context is provided.
+- When relevant context is provided, use it accurately to answer user questions.
+- Seamlessly switch between casual conversation and knowledge-based answering.
+
+Guidelines:
+- If context is provided and relevant, use it as the primary source of truth.
+- If context is not provided or not relevant, respond using your general knowledge.
+- Do not hallucinate or invent information.
+- If you are unsure or the information is not available, clearly state that.
+- Be clear, concise, and helpful in all responses.
+- Maintain a natural, human-like conversational tone.
+- Never mention internal implementation details such as embeddings, vector databases, or system architecture.
+
+You are designed to provide reliable, accurate, and engaging assistance.
+"""
+)
+
+#===========================================
+# Chat function
+#===========================================
+
+async def chat(state:Ragbot_State):
+    query = state['query']
+    context = state['context']
+    metadata = state['metadata']
+    web_context = state['web_context']
+    model_name = state.get('model_name', 'gpt')
+
+    history = state.get("response", [])
+
+    # [CHANGED] Updated Prompt to include History so it remembers your name
+    prompt = f"""
+You are an expert assistant designed to answer user questions using multiple information sources.
+
+Source Priority Rules (STRICT):
+1. **Conversation History**: Check if the answer was provided in previous messages (e.g., user's name, previous topics).
+2. If the provided Context contains the answer, use ONLY the Context.
+3. If the Context does not contain the answer and Web Context is available, use the Web Context.
+4. If neither Context nor Web Context contains the answer, use your general knowledge.
+5. Do NOT invent or hallucinate facts.
+6. If the answer cannot be determined, clearly say so.
+
+User Question:
+{query}
+
+Retrieved Context (Vector Database):
+{context}
+
+Metadata:
+{metadata}
+
+Web Context (Real-time Search):
+{web_context}
+
+Final Answer:
+"""
+
+    selected_llm = await get_llm(model_name)
+    messages = [SYSTEM_PROMPT] + history + [HumanMessage(content=prompt)]
+    response = selected_llm.invoke(messages)
+    return {
+        'response': [
+            HumanMessage(content=query), 
+            response
+        ]
+    }
+
+#===========================================
+# Graph Declaration
+#===========================================
+
+# Keeping MemorySaver as requested (Note: RAM only, wipes on restart)
+memory = MemorySaver()
+graph = StateGraph(Ragbot_State)
+
+graph.add_node("fetch_context", fetch_context)
+graph.add_node("fetch_web_context", fetch_web_context)
+graph.add_node("chat", chat)
+
+graph.add_conditional_edges(
+    START,
+    router,
+    {
+        "fetch_context": "fetch_context",
+        "fetch_web_context": "fetch_web_context",
+        "chat": "chat"
+    }
+)
+
+graph.add_edge("fetch_context", "chat")
+graph.add_edge("fetch_web_context", "chat")
+graph.add_edge("chat", END)
+
+app = graph.compile(checkpointer=memory)
+
+
+#===========================================
+# Helper Function
+#===========================================
+
+def ask_bot(query: str, use_rag: bool = False, use_web: bool = False, thread_id: str = "1"):
+    config = {"configurable": {"thread_id": thread_id}}
+    inputs = {
+        "query": query,
+        "RAG": use_rag,
+        "web_search": use_web,
+        "context": [],
+        "metadata": [],
+        "web_context": "",
+    }
+    
+    result = app.invoke(inputs, config=config)
+    last_message = result['response'][-1]
+    
+    return last_message.content
+
+
+"""print("--- Conversation 1 ---")
+# User says hello and gives name
+response = ask_bot("Hi, my name is Junaid", thread_id="session_A")
+print(f"Bot: {response}")
+
+# User asks for name (RAG and Web are OFF)
+response = ask_bot("What is my name?", thread_id="session_A")
 print(f"Bot: {response}")"""