Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from langchain_groq import ChatGroq
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.tools import tool
 from typing import TypedDict, Annotated, List, Optional
 import json
 
@@ -10,39 +11,55 @@ class MedicalConfig:
     SYSTEM_PROMPT = """You are an AI clinical assistant. Follow these rules:
 1. Analyze patient data using latest medical guidelines
 2. Always check for drug interactions
-3. Suggest tests only when necessary
-4. Use structured actions:
-   - order_lab_test: {test_name, reason}
-   - prescribe_medication: {name, dosage, frequency}
-5. Flag high-risk conditions immediately"""
+3. Use structured actions when needed:
+   - lab_order: Order lab tests
+   - prescribe: Prescribe medication
+4. Flag high-risk conditions immediately"""
 
     RED_FLAGS = {
         "symptoms": ["chest pain", "shortness of breath", "severe headache"],
         "vitals": {"temp": (38.5, "°C"), "hr": (120, "bpm"), "bp": ("180/120", "mmHg")}
     }
 
-# State Management
-class AgentState(TypedDict):
-    messages: Annotated[List[dict], lambda l, r: l + r]
-    patient_data: dict
-    approvals: Annotated[List[dict], lambda l, r: l + r]
+# Define tools with proper schemas
+@tool
+def order_lab_test(test_name: str, reason: str) -> str:
+    """Orders a lab test with specified parameters.
+    
+    Args:
+        test_name: Name of the lab test to order
+        reason: Clinical justification for the test
+        
+    Returns:
+        Confirmation message with test details
+    """
+    return f"Lab ordered: {test_name} ({reason})"
+
+@tool
+def prescribe_medication(name: str, dosage: str, frequency: str) -> str:
+    """Prescribes medication with specific dosage instructions.
+    
+    Args:
+        name: Name of the medication
+        dosage: Dosage amount (e.g., '500mg')
+        frequency: Administration frequency (e.g., 'every 6 hours')
+        
+    Returns:
+        Confirmation message with prescription details
+    """
+    return f"Prescribed: {name} {dosage} {frequency}"
+
+# Initialize tools
+tools = [order_lab_test, prescribe_medication, TavilySearchResults(max_results=3)]
 
 class MedicalAgent:
     def __init__(self):
         self.model = ChatGroq(temperature=0.1, model="Llama-3.3-70b-Specdec")
-        self.tools = {
-            "medical_actions": {
-                "order_lab_test": self.order_lab_test,
-                "prescribe_medication": self.prescribe_medication
-            },
-            "research": TavilySearchResults(max_results=3)
-        }
+        self.model_with_tools = self.model.bind_tools(tools)
 
     def analyze_patient(self, patient_data: dict) -> Optional[dict]:
         try:
-            # Bind tools to the model
-            model_with_tools = self.model.bind_tools(list(self.tools["medical_actions"].keys()))
-            response = model_with_tools.invoke([
+            response = self.model_with_tools.invoke([
                 SystemMessage(content=MedicalConfig.SYSTEM_PROMPT),
                 HumanMessage(content=f"Patient Data: {json.dumps(patient_data)}")
             ])
@@ -53,72 +70,16 @@ class MedicalAgent:
 
     def process_action(self, action: dict) -> str:
         try:
-            if action['name'] in self.tools['medical_actions']:
-                return self.tools['medical_actions'][action['name']](action['args'])
-            return f"Unknown action: {action['name']}"
+            tool_name = action['name']
+            args = action['args']
+            
+            if tool_name == "order_lab_test":
+                return order_lab_test.invoke(args)
+            elif tool_name == "prescribe_medication":
+                return prescribe_medication.invoke(args)
+            else:
+                return f"Unknown action: {tool_name}"
         except Exception as e:
             return f"Error processing action: {str(e)}"
 
-    def order_lab_test(self, params: dict) -> str:
-        return f"✅ Lab ordered: {params['test_name']} ({params['reason']})"
-
-    def prescribe_medication(self, params: dict) -> str:
-        return f"✅ Prescribed: {params['name']} {params['dosage']} {params['frequency']}"
-
-    def check_red_flags(self, patient_data: dict) -> List[str]:
-        flags = []
-        for symptom in patient_data.get("symptoms", []):
-            if symptom.lower() in MedicalConfig.RED_FLAGS["symptoms"]:
-                flags.append(f"🚨 High-risk symptom: {symptom}")
-        return flags
-
-# Streamlit UI
-def main():
-    st.set_page_config(page_title="AI Clinic", layout="wide")
-    
-    # Initialize agent
-    if 'agent' not in st.session_state:
-        st.session_state.agent = MedicalAgent()
-    
-    # Patient intake
-    with st.sidebar:
-        st.header("Patient Intake")
-        patient_data = {
-            "symptoms": st.multiselect("Symptoms", ["Fever", "Cough", "Chest Pain", "Shortness of Breath"]),
-            "history": {
-                "conditions": st.text_input("Medical History"),
-                "medications": st.text_input("Current Medications")
-            },
-            "vitals": {
-                "temp": st.number_input("Temp (°C)", 35.0, 42.0, 37.0),
-                "bp": st.text_input("BP (mmHg)", "120/80")
-            }
-        }
-    
-    # Main interface
-    st.title("AI-Powered Clinical Support System")
-    
-    # Check for red flags
-    red_flags = st.session_state.agent.check_red_flags(patient_data)
-    if red_flags:
-        st.warning("High-Risk Conditions Detected")
-        for flag in red_flags:
-            st.error(flag)
-    
-    if st.button("Start Analysis"):
-        with st.spinner("Analyzing patient data..."):
-            response = st.session_state.agent.analyze_patient(patient_data)
-            
-            if response:
-                st.subheader("Analysis Results")
-                if hasattr(response, 'tool_calls') and response.tool_calls:
-                    for action in response.tool_calls:
-                        result = st.session_state.agent.process_action(action)
-                        st.success(result)
-                else:
-                    st.info("No specific actions recommended. Here's the analysis:")
-                    st.write(response.content)
-
-# Run the app
-if __name__ == "__main__":
-    main()
+# Rest of the Streamlit UI code remains the same...