Update app.py

app.py

@@ -2,7 +2,9 @@ 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 typing import TypedDict, Annotated, List
+from typing import TypedDict, Annotated, List, Optional
+import json
+import time
 
 # Configuration
 class MedicalConfig:
@@ -12,12 +14,19 @@ class MedicalConfig:
 3. Suggest tests only when necessary
 4. Use structured actions:
    - order_lab_test: {test_name, reason}
-   - prescribe_medication: {name, dosage, frequency}"""
+   - prescribe_medication: {name, dosage, frequency}
+5. 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]
 
 class MedicalAgent:
     def __init__(self):
@@ -30,23 +39,37 @@ class MedicalAgent:
             "research": TavilySearchResults(max_results=3)
         }
 
-    def analyze_patient(self, patient_data):
-        response = self.model.invoke([
-            SystemMessage(content=MedicalConfig.SYSTEM_PROMPT),
-            HumanMessage(content=f"Patient Data: {patient_data}")
-        ])
-        return response
+    def analyze_patient(self, patient_data: dict) -> Optional[dict]:
+        try:
+            response = self.model.invoke([
+                SystemMessage(content=MedicalConfig.SYSTEM_PROMPT),
+                HumanMessage(content=f"Patient Data: {json.dumps(patient_data)}")
+            ])
+            return response
+        except Exception as e:
+            st.error(f"Error analyzing patient data: {str(e)}")
+            return None
+
+    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 "Unknown action"
+        except Exception as e:
+            return f"Error processing action: {str(e)}"
 
-    def process_action(self, action):
-        if action['name'] in self.tools['medical_actions']:
-            return self.tools['medical_actions'][action['name']](action['args'])
-        return "Unknown action"
+    def order_lab_test(self, params: dict) -> str:
+        return f"✅ Lab ordered: {params['test_name']} ({params['reason']})"
 
-    def order_lab_test(self, params):
-        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 prescribe_medication(self, params):
-        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():
@@ -60,7 +83,7 @@ def main():
     with st.sidebar:
         st.header("Patient Intake")
         patient_data = {
-            "symptoms": st.multiselect("Symptoms", ["Fever", "Cough", "Chest Pain"]),
+            "symptoms": st.multiselect("Symptoms", ["Fever", "Cough", "Chest Pain", "Shortness of Breath"]),
             "history": {
                 "conditions": st.text_input("Medical History"),
                 "medications": st.text_input("Current Medications")
@@ -74,16 +97,25 @@ def main():
     # 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 hasattr(response, 'tool_calls'):
-                for action in response.tool_calls:
-                    result = st.session_state.agent.process_action(action)
-                    st.success(result)
-            else:
-                st.info(response.content)
+            if response:
+                if hasattr(response, 'tool_calls'):
+                    for action in response.tool_calls:
+                        result = st.session_state.agent.process_action(action)
+                        st.success(result)
+                else:
+                    st.info(response.content)
 
+# Run the app
 if __name__ == "__main__":
-    main() 
+    main()