73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
llm.py — LLM interaction layer for The Chaos engine.
|
|
|
|
Provides the low-level call_llm function and environment variable setup
|
|
for provider-specific auth.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
from .state import append_llm_log
|
|
|
|
|
|
def set_llm_env(model: str, api_key: str | None, api_base: str | None) -> None:
|
|
"""Set provider-specific env vars for litellm."""
|
|
prefix = model.split("/")[0].upper()
|
|
key = api_key or "sk-placeholder"
|
|
os.environ[f"{prefix}_API_KEY"] = key
|
|
if api_base:
|
|
os.environ[f"{prefix}_API_BASE"] = api_base
|
|
|
|
|
|
def call_llm(
|
|
messages: list[dict],
|
|
*,
|
|
model: str | None = None,
|
|
temperature: float | None = None,
|
|
timeout: int | None = None,
|
|
max_tokens: int | None = None,
|
|
label: str = "",
|
|
on_debug: callable = None,
|
|
) -> str | None:
|
|
"""Make a single LLM call. Loads config automatically. Returns content text or None on error."""
|
|
from .config import load_config
|
|
cfg = load_config().get("llm", {})
|
|
model = model or cfg.get("model", "ollama/llama3.1")
|
|
temperature = temperature if temperature is not None else cfg.get("temperature", 0.8)
|
|
timeout = timeout if timeout is not None else cfg.get("timeout", 120)
|
|
max_tokens = max_tokens if max_tokens is not None else cfg.get("max_tokens", 4096)
|
|
api_key = cfg.get("api_key")
|
|
api_base = cfg.get("api_base")
|
|
set_llm_env(model, api_key, api_base)
|
|
try:
|
|
import litellm
|
|
except ImportError:
|
|
if on_debug:
|
|
on_debug("llm_error", {"label": label, "error": "litellm not installed"})
|
|
return None
|
|
try:
|
|
response = litellm.completion(
|
|
model=model,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
stream=False,
|
|
timeout=timeout,
|
|
max_tokens=max_tokens,
|
|
)
|
|
content = getattr(response.choices[0].message, 'content', None) or ""
|
|
reasoning = getattr(response.choices[0].message, 'reasoning_content', None) or ""
|
|
if reasoning and reasoning not in content:
|
|
append_llm_log(f"\n--- {label} [reasoning] ---\n{reasoning}")
|
|
text = content or reasoning
|
|
append_llm_log(f"\n--- {label} ---\n{text}")
|
|
return text
|
|
except Exception as e:
|
|
err_msg = f"{type(e).__name__}: {e}"
|
|
append_llm_log(f"\n--- LLM ERROR ({label}) ---\n{err_msg}")
|
|
if on_debug:
|
|
on_debug("llm_error", {"label": label, "error": err_msg})
|
|
return None
|