114 lines
4.2 KiB
Python
114 lines
4.2 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.
|
|
|
|
"""Single transport for all Artemis models: OpenAI Responses API."""
|
|
|
|
from enum import StrEnum
|
|
import hashlib
|
|
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
|
|
|
|
|
|
class ModelProvider(StrEnum):
|
|
OPENAI = "openai"
|
|
|
|
@classmethod
|
|
def from_string(cls, val: Any) -> "ModelProvider":
|
|
if val is None or str(val).strip().lower() in ("", "openai"):
|
|
return cls.OPENAI
|
|
raise ValueError(f"Unsupported LLM provider {val!r}; only OpenAI Responses is supported")
|
|
|
|
|
|
class ModelEndpoint(BaseModel):
|
|
provider: ModelProvider = ModelProvider.OPENAI
|
|
model_name: str = Field(default_factory=lambda: settings.OPENAI_MODEL)
|
|
api_key: str | None = Field(default=None, repr=False)
|
|
api_base: str | None = None
|
|
temperature: float | None = None
|
|
max_tokens: int | None = Field(default=None, gt=0)
|
|
timeout_seconds: float = Field(default_factory=lambda: settings.OPENAI_TIMEOUT_SECONDS, gt=0)
|
|
max_retries: int = Field(default_factory=lambda: settings.OPENAI_MAX_RETRIES, ge=0)
|
|
is_multimodal: bool = True
|
|
reasoning_effort: str | None = Field(default_factory=lambda: settings.OPENAI_REASONING_EFFORT)
|
|
enable_grounding: bool = False
|
|
|
|
def cache_key(self) -> tuple:
|
|
key = self.api_key or (
|
|
settings.OPENAI_API_KEY.get_secret_value() if settings.OPENAI_API_KEY else ""
|
|
)
|
|
return (
|
|
self.provider.value,
|
|
self.model_name,
|
|
self.temperature,
|
|
self.max_tokens,
|
|
self.timeout_seconds,
|
|
self.max_retries,
|
|
self.reasoning_effort,
|
|
self.enable_grounding,
|
|
self.api_base or settings.OPENAI_BASE_URL,
|
|
hashlib.sha256(key.encode()).hexdigest(),
|
|
)
|
|
|
|
|
|
class ModelFactory:
|
|
_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:
|
|
ModelProvider.from_string(endpoint.provider)
|
|
if os.environ.get("ARTEMIS_FAKE_LLM") == "1":
|
|
from artemis.llm.fake_model import FakeChatModel
|
|
|
|
return FakeChatModel(
|
|
delay_s=float(os.environ.get("ARTEMIS_FAKE_LLM_DELAY_S", "0") or 0)
|
|
)
|
|
|
|
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
|
|
)
|
|
if not api_key:
|
|
raise ValueError("请在环境变量或 .env 中配置 OPENAI_API_KEY")
|
|
kwargs = {
|
|
"model": endpoint.model_name,
|
|
"api_key": api_key,
|
|
"base_url": endpoint.api_base or settings.OPENAI_BASE_URL or None,
|
|
"timeout": endpoint.timeout_seconds,
|
|
"stream_chunk_timeout": endpoint.timeout_seconds,
|
|
# Artemis owns the retry budget; SDK retries would multiply it.
|
|
"max_retries": 0,
|
|
"temperature": endpoint.temperature,
|
|
"max_tokens": endpoint.max_tokens,
|
|
"reasoning_effort": None
|
|
if endpoint.reasoning_effort == "default"
|
|
else endpoint.reasoning_effort,
|
|
"use_responses_api": True,
|
|
"output_version": "responses/v1",
|
|
"store": False,
|
|
}
|
|
return ChatOpenAI(**{k: v for k, v in kwargs.items() if v is not None})
|