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. Provides a simple async interface for handlers to call LLM APIs.
Uses harnessed_agent's llm_chat under the hood when available, Looks up model config from the llm database table first,
falls back to direct HTTP calls. falls back to environment variables.
""" """
import json import json
@ -11,15 +12,54 @@ import os
logger = logging.getLogger("pipeline.llm_bridge") 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: async def llm_call(prompt: str, model: str = None, temperature: float = 0.7) -> str:
"""Call LLM and return text response. """Call LLM and return text response.
Tries multiple backends: Backend priority:
1. harnessed_agent.llm_chat (if loaded in ServerEnv) 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: try:
from ahserver.serverenv import ServerEnv from ahserver.serverenv import ServerEnv
env = ServerEnv() env = ServerEnv()
@ -31,32 +71,42 @@ async def llm_call(prompt: str, model: str = None, temperature: float = 0.7) ->
except Exception: except Exception:
pass pass
# Fallback: direct HTTP call to OpenAI-compatible endpoint # Priority 2: DB llm table
import aiohttp cfg = await _get_model_config(model)
if cfg.get("api_key") and cfg.get("api_base"):
api_base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1") api_base = cfg["api_base"]
api_key = os.environ.get("LLM_API_KEY", "") api_key = cfg["api_key"]
model = model or os.environ.get("LLM_MODEL", "gpt-4o-mini") 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: 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 = { headers = {
"Authorization": f"Bearer {api_key}", "Authorization": f"Bearer {api_key}",
"Content-Type": "application/json", "Content-Type": "application/json",
} }
payload = { payload = {
"model": model, "model": model_id,
"messages": [{"role": "user", "content": prompt}], "messages": [{"role": "user", "content": prompt}],
"temperature": temperature, "temperature": temperature,
} }
url = api_base.rstrip("/") + "/chat/completions"
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.post( 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: ) as resp:
if resp.status != 200: if resp.status != 200:
text = await resp.text() 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() data = await resp.json()
return data["choices"][0]["message"]["content"] return data["choices"][0]["message"]["content"]