somew1nd
c2c90a6378
- Implement DevicePool and DeviceExecutionLock concurrency modes (per-device FIFO vs global serialization) - Extend ArtemisClient SDK and CLI with explicit device targeting and concurrency strategy configuration - Add server lifecycle management (artemis restart, stop, status) with cross-platform PID discovery - Update MCP server tools to support multi-device dispatch, trace tagging, and per-device execution - Enhance Showcase UI with step-by-step visual replay, playback speed controls, and device serial badges - Add comprehensive unit test coverage and edge-case verification scripts
438 lines
17 KiB
Python
438 lines
17 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:
|
|
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}")
|
|
|
|
|
|
class ModelRouter:
|
|
"""Dynamic router mapping agent roles to primary and fallback model endpoints."""
|
|
|
|
def __init__(self, default_model: str = "gemini-2.5-flash"):
|
|
self.default_endpoint = ModelEndpoint(
|
|
provider=ModelProvider.GOOGLE,
|
|
model_name=default_model,
|
|
)
|
|
self._role_routes: dict[str, ModelEndpoint] = {}
|
|
self._fallback_routes: dict[str, list[ModelEndpoint]] = {}
|
|
self._setup_default_routes()
|
|
|
|
def _setup_default_routes(self) -> None:
|
|
"""Initializes recommended role mappings."""
|
|
self._role_routes["planner"] = ModelEndpoint(
|
|
provider=ModelProvider.GOOGLE,
|
|
model_name="gemini-2.5-pro",
|
|
temperature=0.2,
|
|
)
|
|
self._role_routes["operator"] = ModelEndpoint(
|
|
provider=ModelProvider.GOOGLE,
|
|
model_name="gemini-2.5-flash",
|
|
temperature=0.0,
|
|
)
|
|
self._role_routes["validator"] = ModelEndpoint(
|
|
provider=ModelProvider.GOOGLE,
|
|
model_name="gemini-2.5-pro",
|
|
temperature=0.1,
|
|
)
|
|
self._role_routes["summarizer"] = ModelEndpoint(
|
|
provider=ModelProvider.GOOGLE,
|
|
model_name="gemini-2.5-flash",
|
|
temperature=0.0,
|
|
)
|
|
self._fallback_routes["planner"] = [
|
|
ModelEndpoint(provider=ModelProvider.GOOGLE, model_name="gemini-2.5-flash"),
|
|
]
|
|
|
|
def set_role_route(
|
|
self,
|
|
role: str,
|
|
endpoint: ModelEndpoint | str,
|
|
fallbacks: list[ModelEndpoint] | None = None,
|
|
) -> None:
|
|
"""Configures or overrides a route for a specific agent role."""
|
|
if isinstance(endpoint, str):
|
|
endpoint = ModelEndpoint(model_name=endpoint)
|
|
self._role_routes[role.lower()] = endpoint
|
|
if fallbacks is not None:
|
|
self._fallback_routes[role.lower()] = fallbacks
|
|
|
|
def get_endpoint(self, role: str) -> ModelEndpoint:
|
|
"""Retrieves the configured primary model endpoint for a role."""
|
|
return self._role_routes.get(role.lower(), self.default_endpoint)
|
|
|
|
def get_fallbacks(self, role: str) -> list[ModelEndpoint]:
|
|
"""Retrieves fallback endpoints for a role."""
|
|
return self._fallback_routes.get(role.lower(), [])
|
|
|
|
def instantiate_model(self, endpoint: ModelEndpoint) -> BaseChatModel:
|
|
"""Instantiates a concrete LangChain chat model via the unified ModelFactory."""
|
|
return ModelFactory.create_model(endpoint)
|
|
|
|
def create_chat_model(self, role: str) -> BaseChatModel:
|
|
"""Creates a chat model with configured fallbacks for the specified role."""
|
|
primary_endpoint = self.get_endpoint(role)
|
|
primary_model = self.instantiate_model(primary_endpoint)
|
|
|
|
fallbacks = self.get_fallbacks(role)
|
|
if fallbacks:
|
|
fallback_models = [self.instantiate_model(ep) for ep in fallbacks]
|
|
return primary_model.with_fallbacks(fallback_models)
|
|
|
|
return primary_model
|
|
|
|
|
|
# Global singleton router instance
|
|
default_router = ModelRouter()
|