Langgraph-Agent V0.5

.gitignore

@@ -2,4 +2,8 @@ agent_advance.py
 *.ipynb
 __pycache__/
 hackathon-healthcare-solutions-9e6f46d0a21e.json
-venv/
+venv/
+.vscode/
+.env
+test.py
+app.py

agent.py

@@ -3,7 +3,6 @@ from dotenv import load_dotenv
 from langchain_community.document_loaders import TextLoader, DirectoryLoader, UnstructuredPDFLoader, UnstructuredWordDocumentLoader
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
 import json
 from google.oauth2 import service_account
@@ -14,36 +13,31 @@ GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 if GEMINI_API_KEY is None:
     GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 
-conf = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
+prod = os.environ.get("PROD")
+if prod == "true":
+    conf = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
 
-service_account_info = json.loads(conf)
-service_account_info = eval(service_account_info)
-
-credentials = service_account.Credentials.from_service_account_info(service_account_info)
+    service_account_info = json.loads(conf)
+    service_account_info = eval(service_account_info)
 
+    credentials = service_account.Credentials.from_service_account_info(service_account_info)
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", GEMINI_API_KEY=GEMINI_API_KEY, temperature=0.7, credentials=credentials)
+else:
+    # Initialize the language model with your API key
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", GEMINI_API_KEY=GEMINI_API_KEY, temperature=0.7)
 DOCUMENT_DIR = 'document/'
 COLLECTION_NAME = "health_documents"
 
-llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", GEMINI_API_KEY=GEMINI_API_KEY, temperature=0.7, credentials=credentials)
 
 
 print("Models initialized successfully.")
 
-import os
-from dotenv import load_dotenv
-from langchain_community.llms import HuggingFacePipeline
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.document_loaders import TextLoader, DirectoryLoader, UnstructuredPDFLoader, UnstructuredWordDocumentLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
-from langchain_community.vectorstores.utils import filter_complex_metadata
 from langchain_community.vectorstores import Chroma
 import torch
 from constants import CHROMA_PATH
 
 # Load environment variables (if needed)
-load_dotenv()
 
 # Define the directory containing the documents
 DOCUMENT_DIR = 'document/'
@@ -191,7 +185,6 @@ def create_health_agent(vector_store):
             docs = retriever.get_relevant_documents(query)
             context = "\n".join([doc.page_content for doc in docs])
 
-            # Prepare inputs for the LLM chain
             llm_inputs = {
                 'context': context,
                 'question': query,
@@ -203,12 +196,11 @@ def create_health_agent(vector_store):
             result = self.llm_chain(llm_inputs)
             return {'result': result['text']}
 
-    # Create the LLM chain
     llm_chain = LLMChain(llm=llm, prompt=PROMPT)
 
-    # Create and return the custom chain
     return CustomRetrievalQA(retriever=retriever, llm_chain=llm_chain, user_data=None)
 
+# from langgraph_agent import initialize_health_agent
 
 def agent_with_db():
     # 1. Load documents

fetch.py

@@ -5,7 +5,7 @@ import json
 Contains the example code to retrieve response from the server in python-requests"""
 
 ## without previous_state
-url = "https://arpit-bansal-healthbridge.hf.space/"
+url = "https://arpit-bansal-healthbridge.hf.space/retrieve"
 headers = {
     "accept": "application/json",
     "Content-Type": "application/json"

langgraph_agent.py

@@ -0,0 +1,200 @@
+# Import LangGraph components
+from langgraph.graph import StateGraph, END
+from typing import TypedDict, List, Dict, Any, Annotated, Union, Literal
+import operator
+from pydantic import BaseModel, Field
+from agent import LLMChain, PromptTemplate, llm, DOCUMENT_DIR, load_documents, split_documents, CHROMA_PATH, load_vectordb, create_and_store_embeddings
+import os
+# Define state schema
+class AgentState(TypedDict):
+    query: str
+    previous_conversation: str
+    user_data: Dict[str, Any]
+    requires_rag: bool
+    context: List[str]
+    response: str
+    
+# Define tools and nodes for the LangGraph
+
+def query_classifier(state: AgentState) -> AgentState:
+    """Determine if the query requires RAG retrieval."""
+    query_lower = state["query"].lower()
+    rag_keywords = [
+        "scheme", "schemes", "program", "programs", "policy", "policies", 
+        "public health engineering", "phe", "public health", "government", 
+        "benefit", "financial", "assistance", "aid", "initiative"
+    ]
+    
+    state["requires_rag"] = any(keyword in query_lower for keyword in rag_keywords)
+    return state
+
+def retrieve_documents(state: AgentState) -> AgentState:
+    """Retrieve documents from vector store if needed."""
+    if state["requires_rag"]:
+        # Get the global vector_store variable
+        # This assumes vector_store is accessible in this scope
+        docs = vector_store.as_retriever(search_kwargs={"k": 5}).get_relevant_documents(state["query"])
+        state["context"] = [doc.page_content for doc in docs]
+    else:
+        state["context"] = []
+    return state
+
+def generate_response(state: AgentState) -> AgentState:
+    """Generate response with or without context."""
+    # style = state["user_data"].get("style", "normal") if isinstance(state["user_data"], dict) else "normal"
+    
+    base_prompt = """You are a helpful health assistant. Who will talk to the user as human and resolve their queries.
+
+Use Previous_Conversation to maintain consistency in the conversation.
+These are Previous_Conversation between you and user.
+Previous_Conversation: {previous_conversation}
+
+These are info about the person.
+User_Data: {user_data}
+
+Points to Adhere:
+1. Only tell the schemes if user specifically asked, otherwise don't share schemes information.
+2. If the user asks about schemes, Ask about what state they belong to first.
+3. You can act as a mental-health counselor if needed. 
+4. Give precautions and natural-remedies for the diseases, if user asked or it's needed, only for Common diseases include the common cold, flu etc.
+5. Ask the preferred language of the user, In the starting of the conversation.
+6. Give the answer in a friendly and conversational tone.
+7. Style to answer in {style} way.
+Question: {question}
+"""
+    
+    if state["requires_rag"] and state["context"]:
+        # Add context to prompt if we're using RAG
+        context = "\n".join(state["context"])
+        prompt_template = base_prompt + "\nContext from knowledge base:\n{context}\n\nAnswer:"
+        prompt = PromptTemplate(
+            template=prompt_template,
+            input_variables=["context", "question", "previous_conversation", "user_data", "style"]
+        )
+        
+        llm_chain = LLMChain(llm=llm, prompt=prompt)
+        result = llm_chain({
+            'context': context,
+            'question': state["query"],
+            'previous_conversation': state["previous_conversation"],
+            'user_data': state["user_data"],
+            'style': state["user_data"].get("style", "normal")
+        })
+    else:
+        # Answer directly without context
+        prompt_template = base_prompt + "\nAnswer:"
+        prompt = PromptTemplate(
+            template=prompt_template,
+            input_variables=["question", "previous_conversation", "user_data", "style"]
+        )
+        
+        llm_chain = LLMChain(llm=llm, prompt=prompt)
+        result = llm_chain({
+            'question': state["query"],
+            'previous_conversation': state["previous_conversation"],
+            'user_data': state["user_data"],
+            'style': state["user_data"].get("style", "normal")
+        })
+    
+    state["response"] = result["text"]
+    return state
+
+def create_agent_workflow():
+    """Create the LangGraph workflow for the health agent."""
+    # Initialize the state graph
+    workflow = StateGraph(AgentState)
+    
+    # Add nodes
+    workflow.add_node("classifier", query_classifier)
+    workflow.add_node("retriever", retrieve_documents)
+    workflow.add_node("responder", generate_response)
+    
+    # Create edges
+    workflow.add_edge("classifier", "retriever")
+    workflow.add_edge("retriever", "responder")
+    workflow.add_edge("responder", END)
+    
+    # Set the entry point
+    workflow.set_entry_point("classifier")
+    
+    # Compile the graph
+    return workflow.compile()
+
+def agent_with_db():
+    # Load or create vector store
+    global vector_store
+    vector_store = None
+    try:
+        vector_store = load_vectordb(CHROMA_PATH)
+    except ValueError:
+        pass
+    
+    UPDATE_DB = os.getenv("UPDATE_DB", "false")
+    if UPDATE_DB.lower() == "true" or vector_store is None:
+        print("Loading and processing documents...")
+        documents = load_documents(DOCUMENT_DIR)
+        chunks = split_documents(documents)
+        
+        try:
+            vector_store = create_and_store_embeddings(chunks)
+        except Exception as e:
+            print(f"Error creating embeddings: {e}")
+            return None
+
+    print("Creating the LangGraph health agent workflow...")
+    agent_workflow = create_agent_workflow()
+    
+    class HealthAgent:
+        def __init__(self, workflow):
+            self.workflow = workflow
+            self.conversation_history = ""
+        
+        def __call__(self, input_data):
+            # Handle both dictionary input and direct arguments
+            if isinstance(input_data, dict):
+                query = input_data.get("query", "")
+                previous_conversation = input_data.get("previous_conversation", "")
+                user_data = input_data.get("user_data", {})
+                style = input_data.get("style", "normal")
+            else:
+                # Assume it's a direct query string
+                query = input_data
+                previous_conversation = ""
+                user_data = {}
+                style = "normal"
+            
+            # Store previous conversation if provided
+            if previous_conversation:
+                self.conversation_history = previous_conversation
+            
+            # Update conversation history
+            if self.conversation_history:
+                self.conversation_history += f"\nUser: {query}\n"
+            else:
+                self.conversation_history = f"User: {query}\n"
+            
+            if "style" not in user_data:
+                user_data["style"] = style
+            # Prepare initial state
+            initial_state = {
+                "query": query,
+                "previous_conversation": self.conversation_history,
+                "user_data": user_data,
+                "requires_rag": False,
+                "context": [],
+                "response": "",
+                # "style": style
+            }
+            print("Initial state:", initial_state)
+            
+            # Run the workflow
+            final_state = self.workflow.invoke(initial_state)
+            print("Final state:", final_state)
+            
+            # Update conversation history with response
+            self.conversation_history += f"Assistant: {final_state['response']}\n"
+            
+            # Return in the expected format
+            return {"result": final_state["response"]}
+        
+    return HealthAgent(agent_workflow)

main.py

@@ -1,8 +1,9 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 import os
 from fastapi import HTTPException
-from agent import agent_with_db
+# from agent import agent_with_db
+from langgraph_agent import agent_with_db
 from schemas import request
 from dotenv import load_dotenv
 load_dotenv()
@@ -44,7 +45,7 @@ async def parse_user_data(user_data):
     return user_info
 
 @app.post("/retrieve", status_code=200)
-async def retrieve(request:request):
+async def retrieve(request:request, url:Request):
     try:
         prev_conv = request.previous_state
         user_info = await parse_user_data(request.user_data)
@@ -53,9 +54,14 @@ async def retrieve(request:request):
             prev_conv = "No previous conversation available, first time"
         query = request.query
         prev_conv = str(prev_conv)
-        user_info = str(user_info)
-        response = agent({"query": query, "previous_conversation": prev_conv, "user_data": user_info})
-
+        # user_info = str(user_info) # Was needed in Old-Rag not needed in LangGraph-Rag.
+        # Did a mistake by choosing to string format for Old-Rag
+        response = agent({"query": query, "previous_conversation": prev_conv, "user_data": user_info, "style": request.user_data["style"]})
+        origin = url.headers.get('origin')
+        if origin is None:
+            origin = url.headers.get('referer')
+        print("origin: ", origin)
+        print("response: ", response)
         return {"response": response["result"]}
     
     except Exception as e: