fix: llm_bridge reads model config from llm DB table, falls back to env vars

This commit is contained in:
yumoqing 2026-08-01 10:25:33 +08:00
parent 23047eed95
commit d535e61f29

View File

@ -1,8 +1,9 @@
"""LLM bridge for pipeline handlers.
"""
LLM bridge for pipeline handlers.
Provides a simple async interface for handlers to call LLM APIs.
Uses harnessed_agent's llm_chat under the hood when available,
falls back to direct HTTP calls.
Looks up model config from the llm database table first,
falls back to environment variables.
"""
import json
@ -11,15 +12,54 @@ import os
logger = logging.getLogger("pipeline.llm_bridge")
# Cache: model_name -> {api_base, api_key, model_id}
_model_cache: dict = {}
async def _get_model_config(model_name: str = None) -> dict:
"""Look up model config from llm table. Returns dict with api_base, api_key, model_id."""
global _model_cache
if model_name and model_name in _model_cache:
return _model_cache[model_name]
try:
from sqlor.dbpools import DBPools
db = DBPools()
dbname = "pipeline"
async with db.sqlorContext(dbname) as sor:
if model_name:
sql = "SELECT api_base, api_key, model_id FROM llm WHERE name=${name}$ AND status='active' LIMIT 1"
recs = await sor.sqlExe(sql, {"name": model_name})
else:
sql = "SELECT api_base, api_key, model_id, name FROM llm WHERE status='active' ORDER BY id LIMIT 1"
recs = await sor.sqlExe(sql, {})
if recs:
r = recs[0]
cfg = {
"api_base": getattr(r, "api_base", "") or "",
"api_key": getattr(r, "api_key", "") or "",
"model_id": getattr(r, "model_id", "") or "",
}
cache_key = model_name or getattr(r, "name", "")
if cache_key:
_model_cache[cache_key] = cfg
return cfg
except Exception as e:
logger.warning("llm_bridge: DB lookup failed: %s", e)
return {}
async def llm_call(prompt: str, model: str = None, temperature: float = 0.7) -> str:
"""Call LLM and return text response.
Tries multiple backends:
Backend priority:
1. harnessed_agent.llm_chat (if loaded in ServerEnv)
2. Direct OpenAI-compatible API call
2. DB llm table (api_base + api_key)
3. Environment variables (LLM_API_BASE, LLM_API_KEY, LLM_MODEL)
"""
# Try harnessed_agent first
# Priority 1: harnessed_agent
try:
from ahserver.serverenv import ServerEnv
env = ServerEnv()
@ -31,32 +71,42 @@ async def llm_call(prompt: str, model: str = None, temperature: float = 0.7) ->
except Exception:
pass
# Fallback: direct HTTP call to OpenAI-compatible endpoint
import aiohttp
api_base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1")
api_key = os.environ.get("LLM_API_KEY", "")
model = model or os.environ.get("LLM_MODEL", "gpt-4o-mini")
# Priority 2: DB llm table
cfg = await _get_model_config(model)
if cfg.get("api_key") and cfg.get("api_base"):
api_base = cfg["api_base"]
api_key = cfg["api_key"]
model_id = cfg.get("model_id") or model or "default"
logger.info("llm_bridge: using DB model config for %s -> %s", model, api_base)
else:
# Priority 3: Environment variables
api_base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1")
api_key = os.environ.get("LLM_API_KEY", "")
model_id = model or os.environ.get("LLM_MODEL", "gpt-4o-mini")
if not api_key:
raise ValueError("No LLM API configured (set LLM_API_KEY env var)")
raise ValueError("No LLM API configured. Please add a model in the llm table or set LLM_API_KEY env var.")
import aiohttp
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
payload = {
"model": model,
"model": model_id,
"messages": [{"role": "user", "content": prompt}],
"temperature": temperature,
}
url = api_base.rstrip("/") + "/chat/completions"
async with aiohttp.ClientSession() as session:
async with session.post(
f"{api_base}/chat/completions", headers=headers, json=payload, timeout=aiohttp.ClientTimeout(total=120)
url, headers=headers, json=payload,
timeout=aiohttp.ClientTimeout(total=120)
) as resp:
if resp.status != 200:
text = await resp.text()
raise ValueError(f"LLM API error {resp.status}: {text[:200]}")
raise ValueError(f"LLM API error {resp.status}: {text[:300]}")
data = await resp.json()
return data["choices"][0]["message"]["content"]