69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
|
|
import os, json, httpx
|
|
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():
|
|
# Prefer Docker secret if available
|
|
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)")
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"ok": True}
|
|
|
|
@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},
|
|
],
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
r = await client.post(url, headers=headers, json=body)
|
|
if r.status_code >= 400:
|
|
raise HTTPException(502, f"OpenAI error: {r.text}")
|
|
data = r.json()
|
|
|
|
try:
|
|
content = data["choices"][0]["message"]["content"]
|
|
cmd = json.loads(content)
|
|
except Exception as e:
|
|
raise HTTPException(500, f"Failed to parse model JSON: {e}; raw={str(data)[:300]}")
|
|
|
|
# Forward to the agent
|
|
async with httpx.AsyncClient(timeout=15) as client:
|
|
r = await client.post(f"{AGENT_URL}/command", json=cmd)
|
|
if r.status_code >= 400:
|
|
raise HTTPException(r.status_code, f"Agent error: {r.text}")
|
|
return r.json()
|