Update app.py

app.py

@@ -330,24 +330,20 @@ try:
 except Exception as e:
     print(f"[ERROR] Workflow execution failed: {e}")
 '''
-
-from typing import Annotated, Dict, Any
-import functools
+from typing import Dict
+from pydantic import BaseModel
 from langgraph.graph import StateGraph, END
+from langchain_core.prompts import PromptTemplate
 from langchain_core.messages import HumanMessage
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_experimental.tools import PythonREPLTool
 from langchain_huggingface import HuggingFacePipeline
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-# Define the agents and tools
+# Define HuggingFace LLM
 def create_llm():
-    """Create the HuggingFace LLM pipeline."""
     model_name = "Qwen/Qwen2.5-7B-Instruct-1M"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
-    
     llm_pipeline = pipeline(
         task="text-generation",
         model=model,
@@ -358,19 +354,13 @@ def create_llm():
     return HuggingFacePipeline(pipeline=llm_pipeline)
 
 # Mock Tools for Registration and Scheduling
-registration_tool = TavilySearchResults(max_results=1)  # A mock registration tool
-scheduling_tool = PythonREPLTool()  # A mock scheduling tool
-
-# Define the agents
-def registration_agent(visitor_details: Dict[str, str]) -> Dict[str, Any]:
+def registration_agent(visitor_details: Dict[str, str]) -> Dict[str, bool]:
     """Check if the visitor is registered."""
     visitor_name = visitor_details.get("visitor_name")
     visitor_mobile = visitor_details.get("visitor_mobile")
 
-    # Mock registration lookup (replace this with actual database/API call)
+    # Mock registration lookup
     registered_visitors = [{"visitor_name": "John Doe", "visitor_mobile": "1234567890"}]
-
-    # Check if visitor exists
     is_registered = any(
         v["visitor_name"] == visitor_name and v["visitor_mobile"] == visitor_mobile
         for v in registered_visitors
@@ -382,75 +372,75 @@ def scheduling_agent(scheduling_details: Dict[str, str]) -> Dict[str, str]:
     doctor_name = scheduling_details.get("doctor_name")
     department_name = scheduling_details.get("department_name")
 
-    # Mock scheduling logic (replace with actual logic/API integration)
+    # Mock scheduling logic
     appointment_status = "Scheduled successfully"
     return {"status": appointment_status, "doctor_name": doctor_name, "department_name": department_name}
 
+# Define the VisitorState schema
+class VisitorState(BaseModel):
+    visitor_name: str = ""
+    visitor_mobile: str = ""
+    doctor_name: str = ""
+    department_name: str = ""
+    messages: list = []
 
-# Define Langgraph States
-class VisitorState:
-    messages: Annotated[list, functools.partial(lambda x: [])]
-
-
-def input_state(state):
-    """State to input visitor details."""
-    return {"messages": [HumanMessage(content="Please provide your name and mobile number.")]}
+def input_state(state: VisitorState) -> dict:
+    """Initial state to collect visitor details."""
+    return {"messages": ["Please provide your name and mobile number."], "next": "RegistrationState"}
 
-def registration_state(state):
+def registration_state(state: VisitorState) -> dict:
     """State to check visitor registration."""
-    visitor_name = state["visitor_name"]
-    visitor_mobile = state["visitor_mobile"]
-    result = registration_agent({"visitor_name": visitor_name, "visitor_mobile": visitor_mobile})
-
+    result = registration_agent(
+        {"visitor_name": state.visitor_name, "visitor_mobile": state.visitor_mobile}
+    )
     if result["registered"]:
         return {"messages": ["Visitor is registered."], "next": "SchedulingState"}
     else:
         return {"messages": ["Visitor not found in records. Please register first."], "next": END}
 
-def scheduling_state(state):
-    """State for scheduling appointment."""
-    doctor_name = state["doctor_name"]
-    department_name = state["department_name"]
-    result = scheduling_agent({"doctor_name": doctor_name, "department_name": department_name})
-
+def scheduling_state(state: VisitorState) -> dict:
+    """State to schedule an appointment."""
+    result = scheduling_agent(
+        {"doctor_name": state.doctor_name, "department_name": state.department_name}
+    )
     return {
-        "messages": [f"Appointment {result['status']} with Dr.{result['doctor_name']} in {result['department_name']} department."],
-        "next": END
+        "messages": [
+            f"Appointment {result['status']} with Dr.{result['doctor_name']} in {result['department_name']} department."
+        ],
+        "next": END,
     }
 
-
-# Build Langgraph Workflow
+# Define Workflow
 workflow = StateGraph(VisitorState)
-
-# Add nodes
 workflow.add_node("InputState", input_state)
 workflow.add_node("RegistrationState", registration_state)
 workflow.add_node("SchedulingState", scheduling_state)
 
 # Define edges
 workflow.add_edge("InputState", "RegistrationState")
-workflow.add_edge("RegistrationState", "SchedulingState")
 workflow.add_conditional_edges(
     "RegistrationState",
-    lambda x: x.get("next"),
-    {"SchedulingState": "SchedulingState", END: END}
+    lambda state: state.get("next"),
+    {
+        "SchedulingState": "SchedulingState",
+        END: END,
+    }
 )
-
 workflow.set_entry_point("InputState")
 compiled_graph = workflow.compile()
 
 # Gradio Frontend
 def gradio_interface(visitor_name, visitor_mobile, doctor_name, department_name):
-    """Gradio interface to interact with the workflow."""
-    state = {
-        "visitor_name": visitor_name,
-        "visitor_mobile": visitor_mobile,
-        "doctor_name": doctor_name,
-        "department_name": department_name,
-    }
+    """Gradio interface function."""
+    state = VisitorState(
+        visitor_name=visitor_name,
+        visitor_mobile=visitor_mobile,
+        doctor_name=doctor_name,
+        department_name=department_name,
+    )
 
-    # Invoke workflow
-    result = compiled_graph.invoke(state)
+    # Execute workflow
+    result = compiled_graph.invoke(state.dict())
     return result["messages"][0]
 
 iface = gr.Interface(