Refactor app.py to streamline conversation state management and update requirements.txt for package versions

app.py

@@ -2,9 +2,6 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from langgraph.graph import StateGraph, END
-from typing import TypedDict, List, Tuple
-import json
 
 # Model configuration
 LLAMA_MODEL = "meta-llama/Llama-2-7b-chat-hf"
@@ -36,7 +33,6 @@ Patient information: {patient_info}
 <|im_start|>assistant
 """
 
-# Load models
 print("Loading Llama-2 model...")
 tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
 model = AutoModelForCausalLM.from_pretrained(
@@ -55,16 +51,9 @@ meditron_model = AutoModelForCausalLM.from_pretrained(
 )
 print("Meditron model loaded successfully!")
 
-# Define the state for LangGraph
-class ConversationState(TypedDict):
-    messages: List[str]
-    history: List[Tuple[str, str]]
-    current_message: str
-    conversation_turns: int
-    patient_data: List[str]
-    llama_response: str
-    final_response: str
-    should_get_suggestions: bool
+# Conversation state tracking
+conversation_turns = {}
+patient_data = {}
 
 def build_llama2_prompt(system_prompt, history, user_input):
     """Format the conversation history and user input for Llama-2 chat models."""
@@ -97,29 +86,25 @@ def get_meditron_suggestions(patient_info):
     suggestion = meditron_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return suggestion
 
-# LangGraph Node Functions
-def initialize_conversation(state: ConversationState) -> ConversationState:
-    """Initialize or update conversation state."""
-    # Update conversation turns
-    state["conversation_turns"] = state.get("conversation_turns", 0) + 1
-    
-    # Add current message to patient data
-    if "patient_data" not in state:
-        state["patient_data"] = []
-    state["patient_data"].append(state["current_message"])
+@spaces.GPU
+def generate_response(message, history):
+    """Generate a response using both models."""
+    # Track conversation turns
+    session_id = "default-session"
+    if session_id not in conversation_turns:
+        conversation_turns[session_id] = 0
+    conversation_turns[session_id] += 1
     
-    # Determine if we should get suggestions (after 4 turns)
-    state["should_get_suggestions"] = state["conversation_turns"] >= 4
+    # Store the entire conversation for reference
+    if session_id not in patient_data:
+        patient_data[session_id] = []
+    patient_data[session_id].append(message)
     
-    return state
-
-def generate_llama_response(state: ConversationState) -> ConversationState:
-    """Generate response using Llama-2 model."""
     # Build the prompt with proper Llama-2 formatting
-    prompt = build_llama2_prompt(SYSTEM_PROMPT, state["history"], state["current_message"])
+    prompt = build_llama2_prompt(SYSTEM_PROMPT, history, message)
     
     # Add summarization instruction after 4 turns
-    if state["conversation_turns"] >= 4:
+    if conversation_turns[session_id] >= 4:
         prompt = prompt.replace("[/INST] ", "[/INST] Now summarize what you've learned and suggest when professional care may be needed. ")
     
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -140,109 +125,29 @@ def generate_llama_response(state: ConversationState) -> ConversationState:
     full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
     llama_response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
     
-    state["llama_response"] = llama_response
-    return state
-
-def generate_medicine_suggestions(state: ConversationState) -> ConversationState:
-    """Generate medicine suggestions using Meditron model."""
-    # Collect full patient conversation
-    full_patient_info = "\n".join(state["patient_data"]) + "\n\nSummary: " + state["llama_response"]
-    
-    # Get medicine suggestions
-    medicine_suggestions = get_meditron_suggestions(full_patient_info)
-    
-    # Format final response
-    final_response = (
-        f"{state['llama_response']}\n\n"
-        f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
-        f"{medicine_suggestions}"
-    )
-    
-    state["final_response"] = final_response
-    return state
-
-def finalize_response(state: ConversationState) -> ConversationState:
-    """Finalize the response without medicine suggestions."""
-    state["final_response"] = state["llama_response"]
-    return state
-
-def should_get_suggestions(state: ConversationState) -> str:
-    """Conditional edge to determine next step."""
-    if state["should_get_suggestions"]:
-        return "get_suggestions"
-    else:
-        return "finalize"
-
-# Create the LangGraph workflow
-def create_medical_workflow():
-    """Create the LangGraph workflow for medical assistant."""
-    workflow = StateGraph(ConversationState)
-    
-    # Add nodes
-    workflow.add_node("initialize", initialize_conversation)
-    workflow.add_node("generate_llama", generate_llama_response)
-    workflow.add_node("get_suggestions", generate_medicine_suggestions)
-    workflow.add_node("finalize", finalize_response)
-    
-    # Define the flow
-    workflow.set_entry_point("initialize")
-    workflow.add_edge("initialize", "generate_llama")
-    workflow.add_conditional_edges(
-        "generate_llama",
-        should_get_suggestions,
-        {
-            "get_suggestions": "get_suggestions",
-            "finalize": "finalize"
-        }
-    )
-    workflow.add_edge("get_suggestions", END)
-    workflow.add_edge("finalize", END)
-    
-    return workflow.compile()
-
-# Initialize the workflow
-medical_workflow = create_medical_workflow()
-
-# Conversation state tracking (for Gradio session management)
-conversation_states = {}
-
-@spaces.GPU
-def generate_response(message, history):
-    """Generate a response using the LangGraph workflow."""
-    session_id = "default-session"
-    
-    # Initialize or get existing conversation state
-    if session_id not in conversation_states:
-        conversation_states[session_id] = {
-            "messages": [],
-            "history": [],
-            "conversation_turns": 0,
-            "patient_data": []
-        }
-    
-    # Update state with current message and history
-    state = conversation_states[session_id].copy()
-    state["current_message"] = message
-    state["history"] = history
-    
-    # Run the workflow
-    result = medical_workflow.invoke(state)
-    
-    # Update the stored conversation state
-    conversation_states[session_id] = {
-        "messages": result["messages"] if "messages" in result else [],
-        "history": history,
-        "conversation_turns": result["conversation_turns"],
-        "patient_data": result["patient_data"]
-    }
+    # After 4 turns, add medicine suggestions from Meditron
+    if conversation_turns[session_id] >= 4:
+        # Collect full patient conversation
+        full_patient_info = "\n".join(patient_data[session_id]) + "\n\nSummary: " + llama_response
+        
+        # Get medicine suggestions
+        medicine_suggestions = get_meditron_suggestions(full_patient_info)
+        
+        # Format final response
+        final_response = (
+            f"{llama_response}\n\n"
+            f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
+            f"{medicine_suggestions}"
+        )
+        return final_response
     
-    return result["final_response"]
+    return llama_response
 
 # Create the Gradio interface
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="Medical Assistant with LangGraph & Medicine Suggestions",
-    description="Tell me about your symptoms, and after gathering enough information, I'll suggest potential remedies using an AI workflow.",
+    title="Medical Assistant with Medicine Suggestions",
+    description="Tell me about your symptoms, and after gathering enough information, I'll suggest potential remedies.",
     examples=[
         "I have a cough and my throat hurts",
         "I've been having headaches for a week",

requirements.txt

@@ -1,18 +1,6 @@
-# Core packages
-gradio==4.24.0
-spaces==0.21.1
-
-# Transformers & tokenization
-transformers==4.40.1
-torch>=2.1.0
-
-# LangGraph
-langgraph==0.0.41
-
-# Optional but often required for transformers
-accelerate==0.30.1
-sentencepiece==0.1.99
-protobuf==4.25.3
-
-# Utility
-typing-extensions>=4.5.0
+gradio>=4.0
+torch>=2.1
+transformers>=4.37
+spaces
+sentencepiece
+accelerate>=0.21.0