Update app2.py

app2.py

@@ -5,7 +5,7 @@ import json
 import re
 
 MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
-
+client = Mistral(api_key=MISTRAL_API_KEY) if MISTRAL_API_KEY else None
 PROMPT_TEMPLATE = """
 You are a clinical triage assistant. A patient describes their symptoms as follows:
 
@@ -64,23 +64,46 @@ def triage_response(symptoms: str) -> dict:
     
     prompt = build_prompt(symptoms)
 
-    body = {
-        "model": "mistral-large-latest",
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.3
-    }
-
     try:
-        response = requests.post(MISTRAL_API_URL, headers=headers, json=body)
-        response.raise_for_status()
-        raw_output = response.json()["choices"][0]["message"]["content"]
+        # Create chat completion using Mistral client
+        response = client.chat.complete(
+            model="mistral-large-latest",
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ],
+            temperature=0.3,
+            max_tokens=500
+        )
+        
+        # Extract content from response
+        raw_output = response.choices[0].message.content
         cleaned_output = extract_json(raw_output)
+        
         try:
-            return json.loads(cleaned_output)
-        except Exception as e:
-            return {"error": "Invalid JSON format", "raw_output": raw_output, "exception": str(e)}
+            result = json.loads(cleaned_output)
+            # Validate required fields
+            required_fields = ["urgency_level", "possible_condition", "recommended_action"]
+            if all(field in result for field in required_fields):
+                return result
+            else:
+                return {
+                    "error": "Missing required fields in response", 
+                    "raw_output": raw_output,
+                    "parsed_result": result
+                }
+        except json.JSONDecodeError as e:
+            return {
+                "error": "Invalid JSON format", 
+                "raw_output": raw_output, 
+                "exception": str(e)
+            }
+            
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": f"Mistral API error: {str(e)}"}
+
 
 # Launch Gradio app with MCP support
 demo = gr.Interface(