Update inference.py

inference.py

@@ -1,24 +1,36 @@
+import os
 import requests
+from openai import OpenAI
 
-BASE_URL = "https://kevanthonyp-it-support-triage.hf.space"
+BASE_URL = os.environ.get("BASE_URL", "https://kevanthonyp-it-support-triage.hf.space")
+
+client = OpenAI(
+    base_url=os.environ["API_BASE_URL"],
+    api_key=os.environ["API_KEY"],
+)
 
 TASKS = ["task_easy", "task_medium", "task_hard"]
 
 def agent_policy(observation):
-    text = str(observation).lower()
-    if "password" in text or "hacked" in text:
-        return {
-            "category": "security",
-            "priority": "high",
-            "response": "Disconnect system and contact IT."
-        }
-    elif "not working" in text:
-        return {
-            "category": "hardware",
-            "priority": "medium",
-            "response": "Restart device."
-        }
-    else:
+    prompt = f"""You are an IT support triage agent. Given this support ticket observation, respond with a JSON object containing:
+- category: one of "security", "hardware", "software", "network"
+- priority: one of "high", "medium", "low"
+- response: a brief action to resolve the issue
+
+Observation: {observation}
+
+Respond with only valid JSON, no markdown."""
+
+    completion = client.chat.completions.create(
+        model="openai/gpt-4o-mini",
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=200,
+    )
+
+    import json
+    try:
+        return json.loads(completion.choices[0].message.content)
+    except Exception:
         return {
             "category": "software",
             "priority": "low",