Files
ai-ops-agent/relay/relay.py
2025-09-18 15:28:36 -04:00

96 lines
3.6 KiB
Python

import os, json, httpx, traceback
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
AGENT_URL = os.getenv("AGENT_URL", "http://ai-agent:8080")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
def _read_api_key():
path = os.getenv("OPENAI_API_KEY_FILE", "/run/secrets/openai_api_key")
if os.path.exists(path):
return open(path, "r").read().strip()
return os.getenv("OPENAI_API_KEY", "")
SYSTEM_PROMPT = (
"You are an ops command planner. Convert the user's intent into a STRICT JSON object "
"with fields: action (scale|restart_service), params (dict). No prose. Examples: "
'{"action":"scale","params":{"service":"weblabs_php","replicas":3}} '
'or {"action":"restart_service","params":{"service":"weblabs_php"}}. '
"Only produce valid JSON."
)
class ChatIn(BaseModel):
prompt: str
app = FastAPI(title="AI Relay (LLM -> Agent)")
_last_debug = {"openai_request": None, "openai_response": None, "agent_request": None, "agent_response": None, "error": None}
@app.get("/health")
def health():
return {"ok": True, "model": OPENAI_MODEL, "agent_url": AGENT_URL}
@app.get("/last-raw")
def last_raw():
# Expose last request/response for debugging
return _last_debug
@app.post("/chat")
async def chat(inp: ChatIn):
api_key = _read_api_key()
if not api_key:
raise HTTPException(500, "Missing OPENAI_API_KEY (env or secret).")
url = "https://api.openai.com/v1/chat/completions"
headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
body = {
"model": OPENAI_MODEL,
"response_format": {"type": "json_object"},
"temperature": 0.1,
"messages": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": inp.prompt},
],
}
_last_debug["openai_request"] = {"url": url, "body": body}
_last_debug["openai_response"] = None
_last_debug["agent_request"] = None
_last_debug["agent_response"] = None
_last_debug["error"] = None
try:
async with httpx.AsyncClient(timeout=30) as client:
r = await client.post(url, headers=headers, json=body)
_last_debug["openai_response"] = {"status": r.status_code, "text": r.text[:500]}
r.raise_for_status()
data = r.json()
except httpx.RequestError as e:
_last_debug["error"] = f"OpenAI network error: {str(e)}"
raise HTTPException(502, f"OpenAI network error: {e}")
except httpx.HTTPStatusError as e:
_last_debug["error"] = f"OpenAI HTTP error: {e.response.text}"
raise HTTPException(502, f"OpenAI error: {e.response.text}")
try:
content = data["choices"][0]["message"]["content"]
cmd = json.loads(content)
except Exception as e:
_last_debug["error"] = f"Parse error: {str(e)}; raw={str(data)[:300]}"
raise HTTPException(500, f"Failed to parse model JSON: {e}; raw={str(data)[:300]}")
_last_debug["agent_request"] = {"url": f"{AGENT_URL}/command", "json": cmd}
try:
async with httpx.AsyncClient(timeout=15) as client:
r = await client.post(f"{AGENT_URL}/command", json=cmd)
_last_debug["agent_response"] = {"status": r.status_code, "text": r.text[:500]}
r.raise_for_status()
return r.json()
except httpx.RequestError as e:
_last_debug["error"] = f"Agent network error: {str(e)}"
raise HTTPException(502, f"Agent network error: {e}")
except httpx.HTTPStatusError as e:
_last_debug["error"] = f"Agent HTTP error: {e.response.text}"
raise HTTPException(e.response.status_code, f"Agent error: {e.response.text}")