diff --git a/.env.example b/.env.example index 931764b..b06bcec 100644 --- a/.env.example +++ b/.env.example @@ -26,6 +26,21 @@ ZEP_API_KEY=your_zep_api_key_here # LLM_BOOST_BASE_URL= # LLM_BOOST_MODEL_NAME= +# Per-task model overrides — pick a cheaper/stronger model per agent task. +# Anything left unset falls back to LLM_MODEL_NAME. +# LLM_MODEL_PROFILE= # agent persona generation +# LLM_MODEL_CONFIG= # simulation config generation +# LLM_MODEL_REPORT= # report-writing agent +# LLM_MODEL_ONTOLOGY= # entity/relationship ontology +# LLM_MODEL_SENTIMENT= # sentiment/topic classification +# LLM_MODEL_TOOLS= # Zep graph tool calls + +# Backup model used automatically when a primary LLM call fails. May live on a +# different provider — base URL / key default to the primary connection. +# LLM_FALLBACK_MODEL= +# LLM_FALLBACK_BASE_URL= +# LLM_FALLBACK_API_KEY= + # API key that callers must supply in X-Api-Key or Authorization: Bearer headers. # Leave empty (default) for open access — suitable for localhost development. # For public / self-hosted deployments set a strong random value: diff --git a/backend/app/config.py b/backend/app/config.py index 7cd80ca..de30f13 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -1,7 +1,7 @@ """Application configuration loaded from the project root `.env` file.""" import os -from typing import List +from typing import Dict, List, Optional from dotenv import load_dotenv @@ -38,6 +38,24 @@ class Config: LLM_BASE_URL = os.environ.get('LLM_BASE_URL', 'https://api.openai.com/v1') LLM_MODEL_NAME = os.environ.get('LLM_MODEL_NAME', 'gpt-4o-mini') + # Per-task model overrides (issue #15). Each agent task may pick a model with + # different cost/quality trade-offs via LLM_MODEL_; anything unset falls + # back to LLM_MODEL_NAME. Recognised task keys are the dict keys below. + LLM_TASK_MODELS: Dict[str, Optional[str]] = { + 'profile': os.environ.get('LLM_MODEL_PROFILE'), + 'config': os.environ.get('LLM_MODEL_CONFIG'), + 'report': os.environ.get('LLM_MODEL_REPORT'), + 'ontology': os.environ.get('LLM_MODEL_ONTOLOGY'), + 'sentiment': os.environ.get('LLM_MODEL_SENTIMENT'), + 'tools': os.environ.get('LLM_MODEL_TOOLS'), + } + + # Optional backup model used when a primary LLM call fails. May live on a + # different provider (separate base URL / key); both default to the primary. + LLM_FALLBACK_MODEL = os.environ.get('LLM_FALLBACK_MODEL', '').strip() + LLM_FALLBACK_BASE_URL = os.environ.get('LLM_FALLBACK_BASE_URL', '').strip() + LLM_FALLBACK_API_KEY = os.environ.get('LLM_FALLBACK_API_KEY', '').strip() + ZEP_API_KEY = os.environ.get('ZEP_API_KEY') APP_API_KEY = os.environ.get('APP_API_KEY', '').strip() @@ -82,6 +100,33 @@ class Config: FRONTEND_DIST_DIR = os.path.join(os.path.dirname(__file__), '../../frontend/dist') TASK_DATA_DIR = os.path.join(os.path.dirname(__file__), '../uploads/tasks') + @classmethod + def model_for_task(cls, task: Optional[str] = None) -> str: + """Return the model configured for an agent task, or the global default.""" + if task: + override = cls.LLM_TASK_MODELS.get(task) + if override: + return override + return cls.LLM_MODEL_NAME + + @classmethod + def llm_settings(cls, task: Optional[str] = None) -> Dict[str, Optional[str]]: + """Resolve primary + fallback LLM connection settings for an agent task. + + The returned dict is directly consumable as ``LLMClient(**settings)``. + Fallback fields are ``None`` unless ``LLM_FALLBACK_MODEL`` is set, and a + fallback may target a different provider (its own base URL / API key). + """ + fallback_model = cls.LLM_FALLBACK_MODEL or None + return { + 'api_key': cls.LLM_API_KEY, + 'base_url': cls.LLM_BASE_URL, + 'model': cls.model_for_task(task), + 'fallback_model': fallback_model, + 'fallback_base_url': (cls.LLM_FALLBACK_BASE_URL or cls.LLM_BASE_URL) if fallback_model else None, + 'fallback_api_key': (cls.LLM_FALLBACK_API_KEY or cls.LLM_API_KEY) if fallback_model else None, + } + @classmethod def validate(cls): """Validate required runtime configuration.""" diff --git a/backend/app/services/oasis_profile_generator.py b/backend/app/services/oasis_profile_generator.py index 25cfc22..713254e 100644 --- a/backend/app/services/oasis_profile_generator.py +++ b/backend/app/services/oasis_profile_generator.py @@ -191,7 +191,7 @@ def __init__( ): self.api_key = api_key or Config.LLM_API_KEY self.base_url = base_url or Config.LLM_BASE_URL - self.model_name = model_name or Config.LLM_MODEL_NAME + self.model_name = model_name or Config.model_for_task('profile') if not self.api_key: raise ValueError("LLM_API_KEY is not configured") diff --git a/backend/app/services/ontology_generator.py b/backend/app/services/ontology_generator.py index 9731708..4b166e9 100644 --- a/backend/app/services/ontology_generator.py +++ b/backend/app/services/ontology_generator.py @@ -162,7 +162,7 @@ class OntologyGenerator: """ def __init__(self, llm_client: Optional[LLMClient] = None): - self.llm_client = llm_client or LLMClient() + self.llm_client = llm_client or LLMClient.for_task('ontology') def generate( self, diff --git a/backend/app/services/report_agent.py b/backend/app/services/report_agent.py index 6eb8216..cba2a30 100644 --- a/backend/app/services/report_agent.py +++ b/backend/app/services/report_agent.py @@ -908,7 +908,7 @@ def __init__( self.simulation_id = simulation_id self.simulation_requirement = simulation_requirement - self.llm = llm_client or LLMClient() + self.llm = llm_client or LLMClient.for_task('report') self.zep_tools = zep_tools or ZepToolsService() # Tool definitions diff --git a/backend/app/services/sentiment_analyzer.py b/backend/app/services/sentiment_analyzer.py index 6a08f45..0845867 100644 --- a/backend/app/services/sentiment_analyzer.py +++ b/backend/app/services/sentiment_analyzer.py @@ -123,7 +123,7 @@ def __init__( api_key=Config.LLM_API_KEY, base_url=Config.LLM_BASE_URL, ) - self._model = Config.LLM_MODEL_NAME + self._model = Config.model_for_task('sentiment') self._archetype_map: Optional[Dict[int, str]] = None self._advocate_threshold = advocate_threshold self._detractor_threshold = detractor_threshold diff --git a/backend/app/services/simulation_config_generator.py b/backend/app/services/simulation_config_generator.py index 9cc0459..2ff40a0 100644 --- a/backend/app/services/simulation_config_generator.py +++ b/backend/app/services/simulation_config_generator.py @@ -230,7 +230,7 @@ def __init__( ): self.api_key = api_key or Config.LLM_API_KEY self.base_url = base_url or Config.LLM_BASE_URL - self.model_name = model_name or Config.LLM_MODEL_NAME + self.model_name = model_name or Config.model_for_task('config') if not self.api_key: raise ValueError("LLM_API_KEY is not configured") diff --git a/backend/app/services/zep_tools.py b/backend/app/services/zep_tools.py index e8a7cf3..85c8690 100644 --- a/backend/app/services/zep_tools.py +++ b/backend/app/services/zep_tools.py @@ -435,7 +435,7 @@ def __init__(self, api_key: Optional[str] = None, llm_client: Optional[LLMClient def llm(self) -> LLMClient: """Lazily initialize the LLM client""" if self._llm_client is None: - self._llm_client = LLMClient() + self._llm_client = LLMClient.for_task('tools') return self._llm_client def _call_with_retry(self, func, operation_name: str, max_retries: int = None): diff --git a/backend/app/utils/llm_client.py b/backend/app/utils/llm_client.py index 50c8d5b..818571f 100644 --- a/backend/app/utils/llm_client.py +++ b/backend/app/utils/llm_client.py @@ -7,17 +7,23 @@ from openai import OpenAI from ..config import Config +from .logger import get_logger from .llm_sanitizer import sanitize_content, parse_json +logger = get_logger('mirofish.llm_client') + class LLMClient: - """LLM client""" + """LLM client with optional per-task model and backup-model fallback.""" def __init__( self, api_key: Optional[str] = None, base_url: Optional[str] = None, - model: Optional[str] = None + model: Optional[str] = None, + fallback_model: Optional[str] = None, + fallback_base_url: Optional[str] = None, + fallback_api_key: Optional[str] = None ): self.api_key = api_key or Config.LLM_API_KEY self.base_url = base_url or Config.LLM_BASE_URL @@ -31,6 +37,52 @@ def __init__( base_url=self.base_url ) + # Optional backup model engaged only when a primary call raises. It may + # target a different provider, so it keeps its own URL/key (defaulting + # to the primary connection). The client is built lazily on first use. + self.fallback_model = fallback_model + self.fallback_base_url = fallback_base_url or self.base_url + self.fallback_api_key = fallback_api_key or self.api_key + self._fallback_client: Optional[OpenAI] = None + + @classmethod + def for_task(cls, task: Optional[str] = None, **overrides) -> "LLMClient": + """Build a client using the per-task model + fallback from Config. + + Keyword overrides (e.g. ``model=...``) take precedence over the resolved + settings when not ``None``, so callers can still force a specific model. + """ + settings = Config.llm_settings(task) + settings.update({k: v for k, v in overrides.items() if v is not None}) + return cls(**settings) + + def _get_fallback_client(self) -> Optional[OpenAI]: + """Lazily build the backup client, or None when no fallback is configured.""" + if not self.fallback_model or self.fallback_model == self.model: + return None + if self._fallback_client is None: + self._fallback_client = OpenAI( + api_key=self.fallback_api_key, + base_url=self.fallback_base_url + ) + return self._fallback_client + + def _create(self, kwargs: Dict[str, Any]) -> str: + """Call the primary model, falling back to the backup model on failure.""" + try: + response = self.client.chat.completions.create(model=self.model, **kwargs) + return response.choices[0].message.content + except Exception as primary_error: + fallback = self._get_fallback_client() + if fallback is None: + raise + logger.warning( + "Primary model '%s' failed (%s); falling back to '%s'", + self.model, primary_error, self.fallback_model + ) + response = fallback.chat.completions.create(model=self.fallback_model, **kwargs) + return response.choices[0].message.content + def chat( self, messages: List[Dict[str, str]], @@ -51,7 +103,6 @@ def chat( Model response text """ kwargs = { - "model": self.model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens, @@ -60,10 +111,10 @@ def chat( if response_format: kwargs["response_format"] = response_format - response = self.client.chat.completions.create(**kwargs) - content = response.choices[0].message.content - # Route every response through the shared guardrail so reasoning-model - # artefacts (e.g. blocks from MiniMax/GLM) never leak downstream. + # _create selects the model (per-task) and falls back to the backup + # model on failure. Route every response through the shared guardrail so + # reasoning-model artefacts (e.g. blocks) never leak downstream. + content = self._create(kwargs) return sanitize_content(content) def chat_json( diff --git a/backend/tests/test_llm_task_config.py b/backend/tests/test_llm_task_config.py new file mode 100644 index 0000000..3c563f6 --- /dev/null +++ b/backend/tests/test_llm_task_config.py @@ -0,0 +1,191 @@ +"""Tests for per-task model selection and backup-model fallback (issue #15).""" + +import pytest + +from app.config import Config +from app.utils.llm_client import LLMClient + + +# --- Config.model_for_task --------------------------------------------------- + +def test_model_for_task_uses_override(monkeypatch): + monkeypatch.setattr(Config, "LLM_MODEL_NAME", "global-model") + monkeypatch.setattr( + Config, "LLM_TASK_MODELS", {"report": "report-model", "profile": None} + ) + assert Config.model_for_task("report") == "report-model" + + +def test_model_for_task_falls_back_to_global(monkeypatch): + monkeypatch.setattr(Config, "LLM_MODEL_NAME", "global-model") + monkeypatch.setattr(Config, "LLM_TASK_MODELS", {"profile": None}) + # Unset task override and unknown task both resolve to the global default. + assert Config.model_for_task("profile") == "global-model" + assert Config.model_for_task("does-not-exist") == "global-model" + assert Config.model_for_task(None) == "global-model" + + +def test_llm_settings_includes_fallback_when_configured(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + monkeypatch.setattr(Config, "LLM_BASE_URL", "https://primary/v1") + monkeypatch.setattr(Config, "LLM_MODEL_NAME", "global-model") + monkeypatch.setattr(Config, "LLM_TASK_MODELS", {"report": "report-model"}) + monkeypatch.setattr(Config, "LLM_FALLBACK_MODEL", "backup-model") + monkeypatch.setattr(Config, "LLM_FALLBACK_BASE_URL", "https://backup/v1") + monkeypatch.setattr(Config, "LLM_FALLBACK_API_KEY", "backup-key") + + settings = Config.llm_settings("report") + assert settings == { + "api_key": "primary-key", + "base_url": "https://primary/v1", + "model": "report-model", + "fallback_model": "backup-model", + "fallback_base_url": "https://backup/v1", + "fallback_api_key": "backup-key", + } + + +def test_llm_settings_fallback_defaults_to_primary_connection(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + monkeypatch.setattr(Config, "LLM_BASE_URL", "https://primary/v1") + monkeypatch.setattr(Config, "LLM_MODEL_NAME", "global-model") + monkeypatch.setattr(Config, "LLM_TASK_MODELS", {}) + monkeypatch.setattr(Config, "LLM_FALLBACK_MODEL", "backup-model") + monkeypatch.setattr(Config, "LLM_FALLBACK_BASE_URL", "") + monkeypatch.setattr(Config, "LLM_FALLBACK_API_KEY", "") + + settings = Config.llm_settings("report") + assert settings["fallback_base_url"] == "https://primary/v1" + assert settings["fallback_api_key"] == "primary-key" + + +def test_llm_settings_no_fallback_when_unset(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + monkeypatch.setattr(Config, "LLM_MODEL_NAME", "global-model") + monkeypatch.setattr(Config, "LLM_TASK_MODELS", {}) + monkeypatch.setattr(Config, "LLM_FALLBACK_MODEL", "") + + settings = Config.llm_settings("report") + assert settings["fallback_model"] is None + assert settings["fallback_base_url"] is None + assert settings["fallback_api_key"] is None + + +# --- Fake OpenAI client ------------------------------------------------------ + +class _FakeMessage: + def __init__(self, content): + self.content = content + + +class _FakeChoice: + def __init__(self, content): + self.message = _FakeMessage(content) + self.finish_reason = "stop" + + +class _FakeCompletion: + def __init__(self, content): + self.choices = [_FakeChoice(content)] + + +class _RecordingClient: + """Returns canned content (or raises) and records the model it was called with.""" + + def __init__(self, content=None, error=None): + self._content = content + self._error = error + self.calls = [] + + class _Completions: + def create(_self, **kwargs): + self.calls.append(kwargs.get("model")) + if self._error is not None: + raise self._error + return _FakeCompletion(self._content) + + class _Chat: + completions = _Completions() + + self.chat = _Chat() + + +# --- LLMClient.for_task ------------------------------------------------------ + +def test_for_task_builds_client_with_task_model(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + monkeypatch.setattr(Config, "LLM_MODEL_NAME", "global-model") + monkeypatch.setattr(Config, "LLM_TASK_MODELS", {"report": "report-model"}) + monkeypatch.setattr(Config, "LLM_FALLBACK_MODEL", "") + + client = LLMClient.for_task("report") + assert client.model == "report-model" + assert client.fallback_model is None + + +def test_for_task_override_takes_precedence(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + monkeypatch.setattr(Config, "LLM_MODEL_NAME", "global-model") + monkeypatch.setattr(Config, "LLM_TASK_MODELS", {"report": "report-model"}) + monkeypatch.setattr(Config, "LLM_FALLBACK_MODEL", "") + + client = LLMClient.for_task("report", model="forced-model") + assert client.model == "forced-model" + + +# --- Fallback behaviour ------------------------------------------------------ + +def test_chat_falls_back_to_backup_on_primary_failure(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + client = LLMClient( + api_key="primary-key", + model="primary-model", + fallback_model="backup-model", + ) + client.client = _RecordingClient(error=RuntimeError("primary down")) + backup = _RecordingClient(content="from backup") + client._fallback_client = backup + + result = client.chat([{"role": "user", "content": "hi"}]) + assert result == "from backup" + assert backup.calls == ["backup-model"] + + +def test_chat_no_fallback_propagates_error(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + client = LLMClient(api_key="primary-key", model="primary-model") + client.client = _RecordingClient(error=RuntimeError("primary down")) + + with pytest.raises(RuntimeError, match="primary down"): + client.chat([{"role": "user", "content": "hi"}]) + + +def test_chat_uses_primary_when_it_succeeds(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + client = LLMClient( + api_key="primary-key", + model="primary-model", + fallback_model="backup-model", + ) + primary = _RecordingClient(content="from primary") + client.client = primary + backup = _RecordingClient(content="from backup") + client._fallback_client = backup + + result = client.chat([{"role": "user", "content": "hi"}]) + assert result == "from primary" + assert primary.calls == ["primary-model"] + assert backup.calls == [] # backup never touched + + +def test_fallback_skipped_when_same_as_primary(monkeypatch): + monkeypatch.setattr(Config, "LLM_API_KEY", "primary-key") + client = LLMClient( + api_key="primary-key", + model="same-model", + fallback_model="same-model", + ) + client.client = _RecordingClient(error=RuntimeError("down")) + # No distinct backup model -> error should propagate, not loop on itself. + with pytest.raises(RuntimeError, match="down"): + client.chat([{"role": "user", "content": "hi"}])