Update app.py

app.py

@@ -1,103 +1,112 @@
 # app.py
 import streamlit as st
-from uuid import uuid4
 import sqlite3
 import plotly.express as px
+from uuid import uuid4
+from langgraph.graph import START, StateGraph, END
 from langchain_groq import ChatGroq
 from langgraph.checkpoint.sqlite import SqliteSaver
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
+from typing import TypedDict, Annotated, List
+
+# Configuration
+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. Research using tavily_search when uncertain"""
+    
+    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]
 
-# Initialize components
-def init_components():
-    model = ChatGroq(temperature=0.1, model="Llama-3.3-70b-Specdec")
-    tool = TavilySearchResults(max_results=3)
-    conn = sqlite3.connect("medical_checkpoints.db", check_same_thread=False)
-    memory = SqliteSaver(conn)
-    return model, tool, memory
+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.checkpointer = SqliteSaver(sqlite3.connect("medical.db"))
+        self._build_graph()
 
-# Medical Agent Class
-class AdvancedMedicalAgent:
-    def __init__(self, model, tools, memory):
-        self.model = model.bind_tools(tools)
-        self.tools = {t.name: t for t in tools}
-        self.memory = memory
+    def _build_graph(self):
+        graph = StateGraph(AgentState)
+        graph.add_node("analyze", self.analyze_patient)
+        graph.add_node("execute", self.execute_actions)
+        graph.add_node("safety_check", self.safety_checks)
         
-    def process_query(self, patient_data, thread_id):
-        # Initialize or retrieve conversation thread
-        if 'history' not in st.session_state:
-            st.session_state.history = []
-            
-        # Add patient data visualization
-        self.display_patient_dashboard(patient_data)
+        graph.add_edge(START, "analyze")
+        graph.add_conditional_edges(
+            "analyze",
+            self.route_actions,
+            {"medical": "safety_check", "research": "execute"}
+        )
+        graph.add_edge("safety_check", "execute")
+        graph.add_edge("execute", "analyze")
         
-        # Process medical inquiry
-        response = self.model.invoke([
-            SystemMessage(content=self.get_system_prompt()),
-            HumanMessage(content=f"Patient Case:\n{patient_data}")
+        graph.set_entry_point("analyze")
+        self.graph = graph.compile(checkpointer=self.checkpointer)
+
+    def analyze_patient(self, state: AgentState):
+        response = self.model.bind_tools(list(self.tools['medical_actions'].keys())).invoke([
+            SystemMessage(content=MedicalConfig.SYSTEM_PROMPT),
+            HumanMessage(content=f"Patient Data: {state['patient_data']}")
         ])
-        
-        # Handle tool calls
-        if hasattr(response, 'tool_calls'):
-            self.handle_medical_actions(response.tool_calls, thread_id)
-            
-        st.session_state.history.append(("AI", response.content))
-        
-    def display_patient_dashboard(self, data):
-        tab1, tab2, tab3 = st.tabs(["Vitals", "History", "Timeline"])
-        
-        with tab1:
-            fig = px.line(
-                x=["Temperature", "BP", "Heart Rate"],
-                y=[data['vitals']['temp'], 130, 85],  # Example data
-                title="Vital Signs"
-            )
-            st.plotly_chart(fig)
-            
-        with tab2:
-            st.json(data['history'])
-            
-        with tab3:
-            st.vega_lite_chart({
-                "mark": {"type": "circle", "tooltip": True},
-                "encoding": {
-                    "x": {"field": "timestamp", "type": "temporal"},
-                    "y": {"field": "severity", "type": "quantitative"}
-                }
-            })
+        return {"messages": [response]}
+    
+    def safety_checks(self, state: AgentState):
+        current_action = state['messages'][-1].tool_calls[0]
+        risk = self.check_contraindications(state['patient_data'], current_action)
+        return {"approvals": [{"action": current_action, "risk": risk}]}
 
-    def handle_medical_actions(self, tool_calls, thread_id):
-        for call in tool_calls:
-            st.warning(f"🩺 Pending Medical Action: {call['name']}")
-            
-            if st.button("Review Action Details"):
-                st.write(f"**Action Type:** {call['name']}")
-                st.write(f"**Parameters:** {call['args']}")
-                
-                if st.checkbox("I approve this action"):
-                    result = self.tools[call['name']].invoke(call['args'])
-                    st.session_state.history.append(
-                        ("System", f"Action executed: {result}")
-                    )
-                else:
-                    st.session_state.history.append(
-                        ("System", "Action cancelled by clinician")
-                    )
+    def execute_actions(self, state: AgentState):
+        results = []
+        for action in state.get('approvals', []):
+            if action['risk']:
+                results.append(f"Blocked dangerous action: {action['action']['name']}")
+            else:
+                tool = self.tools['medical_actions'][action['action']['name']]
+                results.append(tool(action['action']['args']))
+        return {"messages": [{"content": "\n".join(results)}]}
 
-    def get_system_prompt(self):
-        return """You are an advanced medical AI assistant. Follow these steps:
-1. Analyze patient data using latest clinical guidelines
-2. Consider drug interactions and contraindications
-3. Propose differential diagnosis (3 possibilities)
-4. Suggest evidence-based treatment options
-5. Identify necessary lab tests with reasoning
-6. Flag high-risk factors in RED
-7. Maintain audit trail of all decisions"""
+    def check_contraindications(self, patient_data, action):
+        # Implement actual medical safety checks
+        if 'prescribe' in action['name']:
+            return any(drug in patient_data['medications'] 
+                      for drug in ['warfarin', 'insulin'])
+        return False
+
+    def order_lab_test(self, params):
+        return f"Lab ordered: {params['test_name']} ({params['reason']})"
+
+    def prescribe_medication(self, params):
+        return f"Prescribed: {params['name']} {params['dosage']} {params['frequency']}"
 
 # Streamlit UI
 def main():
-    st.set_page_config(page_title="AI Clinical Assistant", layout="wide")
+    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 = {
@@ -107,25 +116,66 @@ def main():
                 "medications": st.text_input("Current Medications")
             },
             "vitals": {
-                "temp": st.number_input("Temperature (°C)", 36.0, 42.0, 37.0),
-                "bp": st.text_input("Blood Pressure (mmHg)", "120/80")
+                "temp": st.number_input("Temp (°C)", 35.0, 42.0, 37.0),
+                "bp": st.text_input("BP (mmHg)", "120/80")
             }
         }
     
-    st.title("🧠 AI-Powered Clinical Decision Support")
+    # Main interface
+    st.title("AI-Powered Clinical Support System")
     
-    model, tool, memory = init_components()
-    agent = AdvancedMedicalAgent(model, [tool], memory)
+    col1, col2 = st.columns([2, 1])
     
-    if st.button("Start Analysis"):
-        with st.spinner("Analyzing patient case..."):
-            agent.process_query(patient_data, str(uuid4()))
-            
-    if 'history' in st.session_state:
-        st.subheader("Clinical Decision Log")
-        for role, content in st.session_state.history:
-            with st.chat_message(role):
-                st.markdown(content)
+    with col1:
+        st.subheader("Patient Overview")
+        self.render_patient_dashboard(patient_data)
+        
+        if st.button("Start Analysis"):
+            thread_id = f"case_{uuid4().hex[:8]}"
+            with st.spinner("Analyzing..."):
+                st.session_state.agent.graph.invoke(
+                    {"patient_data": patient_data},
+                    {"configurable": {"thread_id": thread_id}}
+                )
+    
+    with col2:
+        st.subheader("Action Center")
+        self.render_approval_interface()
+
+def render_patient_dashboard(self, data):
+    tab1, tab2 = st.tabs(["Vitals", "Timeline"])
+    
+    with tab1:
+        fig = px.line(
+            x=["Temperature", "Blood Pressure"],
+            y=[data['vitals']['temp'], 120],
+            title="Vital Signs"
+        )
+        st.plotly_chart(fig)
+    
+    with tab2:
+        st.vega_lite_chart({
+            "mark": {"type": "line", "point": True},
+            "encoding": {
+                "x": {"field": "hour", "type": "ordinal"},
+                "y": {"field": "temp", "type": "quantitative"}
+            },
+            "data": {"values": [{"hour": i, "temp": 36.5 + i/10} for i in range(24)]}
+        })
+
+def render_approval_interface(self):
+    if 'current_action' in st.session_state:
+        action = st.session_state.current_action
+        st.warning("Action Requires Approval")
+        st.write(f"**Type:** {action['name'].replace('_', ' ').title()}")
+        st.json(action['args'])
+        
+        if st.button("Approve"):
+            self.process_approval(True)
+        if st.button("Reject"):
+            self.process_approval(False)
+    else:
+        st.info("No pending actions")
 
 if __name__ == "__main__":
     main()