app.py

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

# 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. 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")}
    }

# 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.model_with_tools = self.model.bind_tools(tools)

    def analyze_patient(self, patient_data: dict) -> Optional[dict]:
        try:
            response = self.model_with_tools.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:
            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)}"

# Rest of the Streamlit UI code remains the same...