Update app.py

app.py

@@ -1,13 +1,7 @@
-# app.py
 import streamlit as st
-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 langchain_core.messages import HumanMessage, SystemMessage
 from typing import TypedDict, Annotated, List
 
 # Configuration
@@ -18,19 +12,12 @@ class MedicalConfig:
 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')}
-    }
+   - prescribe_medication: {name, dosage, frequency}"""
 
 # 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):
@@ -42,55 +29,18 @@ class MedicalAgent:
             },
             "research": TavilySearchResults(max_results=3)
         }
-        self.checkpointer = SqliteSaver(sqlite3.connect("medical.db"))
-        self._build_graph()
-
-    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)
-        
-        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")
-        
-        graph.set_entry_point("analyze")
-        self.graph = graph.compile(checkpointer=self.checkpointer)
 
-    def analyze_patient(self, state: AgentState):
+    def analyze_patient(self, patient_data):
         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']}")
+            HumanMessage(content=f"Patient Data: {patient_data}")
         ])
-        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}]}
+        return response
 
-    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 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 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):
         return f"Lab ordered: {params['test_name']} ({params['reason']})"
@@ -124,58 +74,16 @@ def main():
     # Main interface
     st.title("AI-Powered Clinical Support System")
     
-    col1, col2 = st.columns([2, 1])
-    
-    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 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 __name__ == "__main__":
     main()