221 lines
6.6 KiB
Python
221 lines
6.6 KiB
Python
"""SDK model configuration helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import TYPE_CHECKING
|
|
|
|
from agents import set_default_openai_api, set_default_openai_key, set_tracing_disabled
|
|
from agents.models.multi_provider import MultiProvider
|
|
from agents.retry import (
|
|
ModelRetryBackoffSettings,
|
|
ModelRetrySettings,
|
|
retry_policies,
|
|
)
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from agents.models.interface import ModelProvider
|
|
|
|
from strix.config.settings import Settings
|
|
|
|
|
|
class StrixProvider(MultiProvider):
|
|
"""Route any non-OpenAI prefix through LiteLLM with the prefix preserved,
|
|
so users type ``deepseek/deepseek-chat`` rather than
|
|
``litellm/deepseek/deepseek-chat``.
|
|
"""
|
|
|
|
def _resolve_prefixed_model(
|
|
self,
|
|
*,
|
|
original_model_name: str,
|
|
prefix: str,
|
|
stripped_model_name: str | None,
|
|
) -> tuple[ModelProvider, str | None]:
|
|
if prefix in {"openai", "litellm", "any-llm"}:
|
|
return super()._resolve_prefixed_model(
|
|
original_model_name=original_model_name,
|
|
prefix=prefix,
|
|
stripped_model_name=stripped_model_name,
|
|
)
|
|
if prefix == "ollama" and stripped_model_name:
|
|
return self._get_fallback_provider("litellm"), f"ollama_chat/{stripped_model_name}"
|
|
return self._get_fallback_provider("litellm"), original_model_name
|
|
|
|
|
|
DEFAULT_MODEL_RETRY = ModelRetrySettings(
|
|
max_retries=5,
|
|
backoff=ModelRetryBackoffSettings(
|
|
initial_delay=2.0,
|
|
max_delay=90.0,
|
|
multiplier=2.0,
|
|
jitter=False,
|
|
),
|
|
policy=retry_policies.any(
|
|
retry_policies.provider_suggested(),
|
|
retry_policies.network_error(),
|
|
retry_policies.http_status((429, 500, 502, 503, 504)),
|
|
),
|
|
)
|
|
|
|
RECOMMENDED_MODEL_NAMES = (
|
|
"openai/gpt-5.4",
|
|
"anthropic/claude-sonnet-4-6",
|
|
"vertex_ai/gemini-3-pro-preview",
|
|
)
|
|
|
|
_RECOMMENDED_MODEL_NAME_SET = frozenset(name.lower() for name in RECOMMENDED_MODEL_NAMES)
|
|
|
|
FRONTIER_MODEL_PROVIDERS = (
|
|
"openai",
|
|
"anthropic",
|
|
"vertex_ai",
|
|
"gemini",
|
|
"google",
|
|
)
|
|
|
|
FRONTIER_MODEL_PREFIXES = (
|
|
"gpt-5",
|
|
"claude-opus-4",
|
|
"claude-sonnet-4",
|
|
"gemini-3",
|
|
)
|
|
|
|
|
|
def configure_sdk_model_defaults(settings: Settings) -> None:
|
|
"""Apply Strix config to SDK-native defaults."""
|
|
llm = settings.llm
|
|
set_tracing_disabled(True)
|
|
_configure_litellm_compatibility()
|
|
if llm.api_key:
|
|
set_default_openai_key(llm.api_key, use_for_tracing=False)
|
|
_configure_litellm_default("api_key", llm.api_key)
|
|
_mirror_api_key_to_provider_env(llm.model, llm.api_key)
|
|
if llm.api_base:
|
|
os.environ["OPENAI_BASE_URL"] = llm.api_base
|
|
_configure_litellm_default("api_base", llm.api_base)
|
|
set_default_openai_api("chat_completions")
|
|
else:
|
|
set_default_openai_api("responses")
|
|
|
|
|
|
def _mirror_api_key_to_provider_env(model_name: str | None, api_key: str) -> None:
|
|
if not model_name:
|
|
return
|
|
import litellm
|
|
|
|
name = model_name.strip()
|
|
for prefix in ("litellm/", "any-llm/"):
|
|
if name.lower().startswith(prefix):
|
|
name = name[len(prefix) :]
|
|
break
|
|
try:
|
|
report = litellm.validate_environment(model=name.lower())
|
|
except Exception: # noqa: BLE001
|
|
return
|
|
for env_key in report.get("missing_keys") or []:
|
|
if env_key.endswith("_API_KEY"):
|
|
os.environ.setdefault(env_key, api_key)
|
|
|
|
|
|
def _configure_litellm_compatibility() -> None:
|
|
"""Enable LiteLLM's permissive param handling and disable its callbacks."""
|
|
import litellm
|
|
|
|
litellm.drop_params = True
|
|
litellm.modify_params = True
|
|
litellm.turn_off_message_logging = True
|
|
litellm.disable_streaming_logging = True
|
|
litellm.suppress_debug_info = True
|
|
|
|
_register_litellm_cost_callback()
|
|
|
|
|
|
def _register_litellm_cost_callback() -> None:
|
|
import litellm
|
|
|
|
from strix.report.state import litellm_cost_callback
|
|
|
|
for bucket_name in ("success_callback", "_async_success_callback"):
|
|
bucket = getattr(litellm, bucket_name, None)
|
|
if not isinstance(bucket, list):
|
|
continue
|
|
if litellm_cost_callback in bucket:
|
|
continue
|
|
bucket.append(litellm_cost_callback)
|
|
|
|
|
|
def _configure_litellm_default(name: str, value: str) -> None:
|
|
"""Set LiteLLM's module-level defaults without adding a provider wrapper."""
|
|
import litellm
|
|
|
|
setattr(litellm, name, value)
|
|
|
|
|
|
def uses_chat_completions_tool_schema(model_name: str, settings: Settings) -> bool:
|
|
"""Return whether the resolved SDK route can only receive JSON function tools."""
|
|
model = model_name.strip().lower()
|
|
if "/" in model and not model.startswith("openai/"):
|
|
return True
|
|
if settings.llm.api_base:
|
|
return True
|
|
return not model_supports_reasoning(model_name)
|
|
|
|
|
|
def model_supports_reasoning(model_name: str) -> bool:
|
|
import litellm
|
|
|
|
name = model_name.strip().lower()
|
|
for prefix in ("litellm/", "any-llm/", "openai/"):
|
|
if name.startswith(prefix):
|
|
name = name[len(prefix) :]
|
|
break
|
|
entry = litellm.model_cost.get(name)
|
|
if entry is None and "/" in name:
|
|
entry = litellm.model_cost.get(name.rsplit("/", 1)[1])
|
|
return bool(entry and entry.get("supports_reasoning"))
|
|
|
|
|
|
def is_recommended_or_frontier_model(model_name: str) -> bool:
|
|
"""Return whether a model is recommended or in a frontier model family."""
|
|
name = _normalized_model_name(model_name)
|
|
if not name:
|
|
return False
|
|
if name in _RECOMMENDED_MODEL_NAME_SET:
|
|
return True
|
|
provider_name, bare_model_name = _split_model_provider(name)
|
|
if provider_name is None:
|
|
return _is_frontier_model_name(bare_model_name)
|
|
return provider_name in FRONTIER_MODEL_PROVIDERS and _is_frontier_model_name(bare_model_name)
|
|
|
|
|
|
def _normalized_model_name(model_name: str) -> str:
|
|
name = model_name.strip().lower()
|
|
for prefix in ("litellm/", "any-llm/"):
|
|
if name.startswith(prefix):
|
|
name = name[len(prefix) :]
|
|
break
|
|
return name
|
|
|
|
|
|
def _split_model_provider(model_name: str) -> tuple[str | None, str]:
|
|
if "/" not in model_name:
|
|
return None, model_name
|
|
provider_name, bare_model_name = model_name.rsplit("/", 1)
|
|
return provider_name, bare_model_name
|
|
|
|
|
|
def _is_frontier_model_name(model_name: str) -> bool:
|
|
return any(model_name.startswith(prefix) for prefix in FRONTIER_MODEL_PREFIXES)
|
|
|
|
|
|
def is_known_openai_bare_model(model_name: str) -> bool:
|
|
import litellm
|
|
|
|
name = model_name.strip().lower()
|
|
if not name or "/" in name:
|
|
return False
|
|
entry = litellm.model_cost.get(name)
|
|
return bool(entry and entry.get("litellm_provider") == "openai")
|