64 lines
1.9 KiB
Python
64 lines
1.9 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,
|
|
temperature: float,
|
|
timeout: int,
|
|
max_tokens: int,
|
|
label: str = "",
|
|
on_debug: callable = None,
|
|
) -> str | None:
|
|
"""Make a single LLM call. Returns content text or None on error."""
|
|
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
|