somew1nd
67248c887e
Add execution incidents, action bursts and shared Flash history. Expose verification and Explorer settings with streamed check results. Support graceful cancellation, recording recovery and bundled UI resources. Add CI quality checks and clean up comments and documentation.
368 lines
15 KiB
Python
368 lines
15 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Dynamic Multi-Provider Model Router for ARTEMIS.
|
|
|
|
Provides role-based LLM/VLM dispatching, dynamic fallback chains,
|
|
and unified provider configuration across Gemini, Vertex AI, OpenAI,
|
|
Anthropic, OpenRouter, and local Ollama/vLLM endpoints.
|
|
"""
|
|
|
|
from enum import StrEnum
|
|
import os
|
|
from typing import Any
|
|
|
|
from langchain_core.language_models.chat_models import BaseChatModel
|
|
from pydantic import BaseModel, Field
|
|
|
|
from artemis.config.settings import settings
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class ModelProvider(StrEnum):
|
|
GOOGLE = "google"
|
|
GEMINI = "google"
|
|
VERTEX_AI = "vertexai"
|
|
OPENAI = "openai"
|
|
ANTHROPIC = "anthropic"
|
|
OPENROUTER = "openrouter"
|
|
XAI = "xai"
|
|
OLLAMA = "ollama"
|
|
VLLM = "vllm"
|
|
CUSTOM = "custom"
|
|
|
|
@classmethod
|
|
def from_string(cls, val: Any) -> "ModelProvider":
|
|
"""Converts any string, enum, or mock representation into a canonical ModelProvider."""
|
|
if isinstance(val, cls):
|
|
return val
|
|
s = str(val).lower().replace("_", "").replace("-", "").strip()
|
|
mapping = {
|
|
"google": cls.GOOGLE,
|
|
"gemini": cls.GOOGLE,
|
|
"vertexai": cls.VERTEX_AI,
|
|
"vertex": cls.VERTEX_AI,
|
|
"openai": cls.OPENAI,
|
|
"anthropic": cls.ANTHROPIC,
|
|
"claude": cls.ANTHROPIC,
|
|
"openrouter": cls.OPENROUTER,
|
|
"xai": cls.XAI,
|
|
"grok": cls.XAI,
|
|
"ollama": cls.OLLAMA,
|
|
"vllm": cls.VLLM,
|
|
"custom": cls.CUSTOM,
|
|
}
|
|
return mapping.get(s, cls.GOOGLE)
|
|
|
|
|
|
class ModelEndpoint(BaseModel):
|
|
"""Configuration definition for an LLM/VLM model endpoint."""
|
|
|
|
provider: ModelProvider = Field(default=ModelProvider.GOOGLE, description="Model provider")
|
|
model_name: str = Field(default="gemini-2.5-flash", description="Model name identifier")
|
|
api_key: str | None = Field(default=None, description="API Key or secret")
|
|
api_base: str | None = Field(default=None, description="Custom API endpoint base URL")
|
|
temperature: float = Field(default=0.0, description="Sampling temperature")
|
|
max_tokens: int | None = Field(default=None, description="Maximum completion tokens")
|
|
timeout_seconds: float = Field(default=60.0, description="Request timeout in seconds")
|
|
is_multimodal: bool = Field(
|
|
default=True, description="Whether endpoint accepts image/video inputs"
|
|
)
|
|
reasoning_effort: str | None = Field(
|
|
default=None, description="Reasoning effort: 'low', 'medium', 'high'"
|
|
)
|
|
thinking_budget: int | None = Field(default=None, description="Thinking budget token count")
|
|
thinking_level: str | None = Field(
|
|
default=None, description="Thinking level: 'minimal', 'low', 'medium', 'high'"
|
|
)
|
|
include_thoughts: bool | None = Field(
|
|
default=None, description="Whether to include thought traces in response"
|
|
)
|
|
enable_grounding: bool = Field(
|
|
default=False,
|
|
description="Whether to enable Google Search grounding for Gemini endpoints",
|
|
)
|
|
|
|
def cache_key(self) -> tuple:
|
|
"""Returns a deterministic hashable cache key for connection pooling."""
|
|
return (
|
|
self.provider.value,
|
|
self.model_name,
|
|
self.temperature,
|
|
self.timeout_seconds,
|
|
self.thinking_budget,
|
|
self.thinking_level,
|
|
self.include_thoughts,
|
|
self.reasoning_effort,
|
|
self.enable_grounding,
|
|
self.api_base,
|
|
)
|
|
|
|
|
|
def _patch_langchain_google_genai():
|
|
"""Patches ChatGoogleGenerativeAI._process_tool_config to preserve and ensure include_server_side_tool_invocations."""
|
|
try:
|
|
from langchain_google_genai.chat_models import ChatGoogleGenerativeAI
|
|
from google.genai.types import ToolConfig
|
|
|
|
if getattr(ChatGoogleGenerativeAI, "_is_artemis_patched", False):
|
|
return
|
|
|
|
original_process = ChatGoogleGenerativeAI._process_tool_config
|
|
|
|
def patched_process_tool_config(self, tool_choice, tool_config, formatted_tools):
|
|
config = original_process(self, tool_choice, tool_config, formatted_tools)
|
|
|
|
has_builtin = False
|
|
has_functions = False
|
|
if formatted_tools:
|
|
for tool in formatted_tools:
|
|
tool_dict = tool if isinstance(tool, dict) else getattr(tool, "__dict__", {})
|
|
if any(
|
|
tool_dict.get(k) is not None
|
|
for k in (
|
|
"google_search",
|
|
"code_execution",
|
|
"google_maps",
|
|
"google_search_retrieval",
|
|
)
|
|
):
|
|
has_builtin = True
|
|
if tool_dict.get("function_declarations"):
|
|
has_functions = True
|
|
|
|
if tool_config:
|
|
normalized = (
|
|
ToolConfig.model_validate(tool_config)
|
|
if isinstance(tool_config, dict)
|
|
else tool_config
|
|
)
|
|
if getattr(normalized, "include_server_side_tool_invocations", None) is not None:
|
|
if config is None:
|
|
config = ToolConfig(
|
|
include_server_side_tool_invocations=normalized.include_server_side_tool_invocations
|
|
)
|
|
else:
|
|
config.include_server_side_tool_invocations = (
|
|
normalized.include_server_side_tool_invocations
|
|
)
|
|
|
|
if has_builtin and has_functions:
|
|
if config is None:
|
|
config = ToolConfig(include_server_side_tool_invocations=True)
|
|
else:
|
|
config.include_server_side_tool_invocations = True
|
|
|
|
return config
|
|
|
|
ChatGoogleGenerativeAI._process_tool_config = patched_process_tool_config
|
|
ChatGoogleGenerativeAI._is_artemis_patched = True
|
|
except Exception as e:
|
|
logger.warning(f"Could not patch ChatGoogleGenerativeAI._process_tool_config: {e}")
|
|
|
|
|
|
class ModelFactory:
|
|
"""Unified factory for instantiating and caching LangChain chat models across providers."""
|
|
|
|
_cache: dict[tuple, BaseChatModel] = {}
|
|
|
|
@classmethod
|
|
def get_model(cls, endpoint: ModelEndpoint) -> BaseChatModel:
|
|
key = endpoint.cache_key()
|
|
if key not in cls._cache:
|
|
cls._cache[key] = cls.create_model(endpoint)
|
|
return cls._cache[key]
|
|
|
|
@classmethod
|
|
def create_model(cls, endpoint: ModelEndpoint) -> BaseChatModel:
|
|
if os.environ.get("ARTEMIS_FAKE_LLM") == "1":
|
|
from artemis.llm.fake_model import FakeChatModel
|
|
|
|
delay = float(os.environ.get("ARTEMIS_FAKE_LLM_DELAY_S", "0") or 0)
|
|
logger.warning(
|
|
f"ARTEMIS_FAKE_LLM=1 — returning FakeChatModel (delay={delay}s) "
|
|
f"instead of {endpoint.provider}/{endpoint.model_name}"
|
|
)
|
|
return FakeChatModel(delay_s=delay)
|
|
|
|
provider = ModelProvider.from_string(endpoint.provider)
|
|
|
|
if provider == ModelProvider.GOOGLE:
|
|
_patch_langchain_google_genai()
|
|
from langchain_google_genai import (
|
|
ChatGoogleGenerativeAI,
|
|
HarmBlockThreshold,
|
|
HarmCategory,
|
|
)
|
|
|
|
api_key = (
|
|
endpoint.api_key
|
|
or (settings.GOOGLE_API_KEY.get_secret_value() if settings.GOOGLE_API_KEY else None)
|
|
or os.environ.get("GOOGLE_API_KEY")
|
|
or os.environ.get("GEMINI_API_KEY")
|
|
)
|
|
thinking_level = endpoint.thinking_level
|
|
if endpoint.model_name and any(v in endpoint.model_name for v in ("2.5", "2.0", "1.5")):
|
|
thinking_level = None
|
|
|
|
kwargs: dict[str, Any] = {
|
|
"model": endpoint.model_name,
|
|
"temperature": endpoint.temperature,
|
|
"max_output_tokens": endpoint.max_tokens,
|
|
"api_key": api_key,
|
|
"timeout": endpoint.timeout_seconds,
|
|
"thinking_budget": endpoint.thinking_budget,
|
|
"thinking_level": thinking_level,
|
|
"include_thoughts": endpoint.include_thoughts,
|
|
"safety_settings": {
|
|
HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
|
|
HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
|
|
HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
|
|
HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
|
|
},
|
|
}
|
|
return ChatGoogleGenerativeAI(**{k: v for k, v in kwargs.items() if v is not None})
|
|
|
|
elif provider == ModelProvider.VERTEX_AI:
|
|
from langchain_google_vertexai import (
|
|
ChatVertexAI,
|
|
HarmBlockThreshold,
|
|
HarmCategory,
|
|
)
|
|
|
|
kwargs = {
|
|
"model_name": endpoint.model_name,
|
|
"temperature": endpoint.temperature,
|
|
"max_output_tokens": endpoint.max_tokens,
|
|
"timeout": endpoint.timeout_seconds,
|
|
"thinking_budget": endpoint.thinking_budget,
|
|
"safety_settings": {
|
|
HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
|
|
HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
|
|
HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
|
|
HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
|
|
},
|
|
}
|
|
return ChatVertexAI(**{k: v for k, v in kwargs.items() if v is not None})
|
|
|
|
elif provider == ModelProvider.OPENAI:
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
api_key = (
|
|
endpoint.api_key
|
|
or (settings.OPENAI_API_KEY.get_secret_value() if settings.OPENAI_API_KEY else None)
|
|
or os.environ.get("OPENAI_API_KEY", "EMPTY")
|
|
)
|
|
base_url = endpoint.api_base or (
|
|
str(settings.OPENAI_BASE_URL) if settings.OPENAI_BASE_URL else None
|
|
)
|
|
kwargs = {
|
|
"model": endpoint.model_name,
|
|
"temperature": endpoint.temperature,
|
|
"max_tokens": endpoint.max_tokens,
|
|
"api_key": api_key,
|
|
"base_url": base_url,
|
|
"timeout": endpoint.timeout_seconds,
|
|
}
|
|
if endpoint.reasoning_effort:
|
|
kwargs["reasoning_effort"] = endpoint.reasoning_effort
|
|
return ChatOpenAI(**{k: v for k, v in kwargs.items() if v is not None})
|
|
|
|
elif provider == ModelProvider.ANTHROPIC:
|
|
from langchain_anthropic import ChatAnthropic
|
|
|
|
api_key = (
|
|
endpoint.api_key
|
|
or (
|
|
settings.ANTHROPIC_API_KEY.get_secret_value()
|
|
if settings.ANTHROPIC_API_KEY
|
|
else None
|
|
)
|
|
or os.environ.get("ANTHROPIC_API_KEY")
|
|
)
|
|
kwargs = {
|
|
"model": endpoint.model_name,
|
|
"temperature": endpoint.temperature,
|
|
"api_key": api_key,
|
|
"timeout": endpoint.timeout_seconds,
|
|
}
|
|
budget = endpoint.thinking_budget
|
|
if not budget and endpoint.reasoning_effort:
|
|
effort_map = {"low": 2048, "medium": 8192, "high": 32768}
|
|
budget = effort_map.get(endpoint.reasoning_effort.lower())
|
|
if budget:
|
|
kwargs["thinking"] = {"type": "enabled", "budget_tokens": budget}
|
|
kwargs["temperature"] = 1.0
|
|
return ChatAnthropic(**{k: v for k, v in kwargs.items() if v is not None})
|
|
|
|
elif provider == ModelProvider.OPENROUTER:
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
api_key = (
|
|
endpoint.api_key
|
|
or (
|
|
settings.OPEN_ROUTER_API_KEY.get_secret_value()
|
|
if settings.OPEN_ROUTER_API_KEY
|
|
else None
|
|
)
|
|
or os.environ.get("OPEN_ROUTER_API_KEY")
|
|
)
|
|
return ChatOpenAI(
|
|
model=endpoint.model_name,
|
|
temperature=endpoint.temperature,
|
|
api_key=api_key,
|
|
base_url=endpoint.api_base or "https://openrouter.ai/api/v1",
|
|
timeout=endpoint.timeout_seconds,
|
|
)
|
|
|
|
elif provider == ModelProvider.XAI:
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
api_key = (
|
|
endpoint.api_key
|
|
or (settings.XAI_API_KEY.get_secret_value() if settings.XAI_API_KEY else None)
|
|
or os.environ.get("XAI_API_KEY")
|
|
)
|
|
return ChatOpenAI(
|
|
model=endpoint.model_name,
|
|
temperature=endpoint.temperature,
|
|
api_key=api_key,
|
|
base_url=endpoint.api_base or "https://api.x.ai/v1",
|
|
timeout=endpoint.timeout_seconds,
|
|
)
|
|
|
|
elif provider in (ModelProvider.OLLAMA, ModelProvider.VLLM, ModelProvider.CUSTOM):
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
api_key = endpoint.api_key or os.environ.get("OPENAI_API_KEY", "EMPTY")
|
|
base_url = endpoint.api_base or os.environ.get(
|
|
"OPENAI_BASE_URL", "http://localhost:8000/v1"
|
|
)
|
|
kwargs = {
|
|
"model": endpoint.model_name,
|
|
"temperature": endpoint.temperature,
|
|
"max_tokens": endpoint.max_tokens,
|
|
"api_key": api_key,
|
|
"base_url": base_url,
|
|
"timeout": endpoint.timeout_seconds,
|
|
}
|
|
if endpoint.reasoning_effort:
|
|
kwargs["reasoning_effort"] = endpoint.reasoning_effort
|
|
return ChatOpenAI(**{k: v for k, v in kwargs.items() if v is not None})
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported model provider: {provider}")
|