266 lines
8.3 KiB
Python
266 lines
8.3 KiB
Python
import os
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from langchain_anthropic import ChatAnthropic
|
|
from langchain_core.language_models import BaseChatModel
|
|
from langchain_core.messages import BaseMessage
|
|
from langchain_google_genai import ChatGoogleGenerativeAI
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
from ra_aid.chat_models.deepseek_chat import ChatDeepseekReasoner
|
|
from ra_aid.logging_config import get_logger
|
|
|
|
from .models_params import models_params
|
|
|
|
known_temp_providers = {
|
|
"openai",
|
|
"anthropic",
|
|
"openrouter",
|
|
"openai-compatible",
|
|
"gemini",
|
|
"deepseek",
|
|
}
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def get_env_var(name: str, expert: bool = False) -> Optional[str]:
|
|
"""Get environment variable with optional expert prefix and fallback."""
|
|
prefix = "EXPERT_" if expert else ""
|
|
value = os.getenv(f"{prefix}{name}")
|
|
|
|
# If expert mode and no expert value, fall back to base value
|
|
if expert and not value:
|
|
value = os.getenv(name)
|
|
|
|
return value
|
|
|
|
|
|
def create_deepseek_client(
|
|
model_name: str,
|
|
api_key: str,
|
|
base_url: str,
|
|
temperature: Optional[float] = None,
|
|
is_expert: bool = False,
|
|
) -> BaseChatModel:
|
|
"""Create DeepSeek client with appropriate configuration."""
|
|
if model_name.lower() == "deepseek-reasoner":
|
|
return ChatDeepseekReasoner(
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
temperature=(
|
|
0 if is_expert else (temperature if temperature is not None else 1)
|
|
),
|
|
model=model_name,
|
|
)
|
|
|
|
return ChatOpenAI(
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
temperature=0 if is_expert else (temperature if temperature is not None else 1),
|
|
model=model_name,
|
|
)
|
|
|
|
|
|
def create_openrouter_client(
|
|
model_name: str,
|
|
api_key: str,
|
|
temperature: Optional[float] = None,
|
|
is_expert: bool = False,
|
|
) -> BaseChatModel:
|
|
"""Create OpenRouter client with appropriate configuration."""
|
|
if model_name.startswith("deepseek/") and "deepseek-r1" in model_name.lower():
|
|
return ChatDeepseekReasoner(
|
|
api_key=api_key,
|
|
base_url="https://openrouter.ai/api/v1",
|
|
temperature=(
|
|
0 if is_expert else (temperature if temperature is not None else 1)
|
|
),
|
|
model=model_name,
|
|
)
|
|
|
|
return ChatOpenAI(
|
|
api_key=api_key,
|
|
base_url="https://openrouter.ai/api/v1",
|
|
model=model_name,
|
|
**({"temperature": temperature} if temperature is not None else {}),
|
|
)
|
|
|
|
|
|
def get_provider_config(provider: str, is_expert: bool = False) -> Dict[str, Any]:
|
|
"""Get provider-specific configuration."""
|
|
configs = {
|
|
"openai": {
|
|
"api_key": get_env_var("OPENAI_API_KEY", is_expert),
|
|
"base_url": None,
|
|
},
|
|
"anthropic": {
|
|
"api_key": get_env_var("ANTHROPIC_API_KEY", is_expert),
|
|
"base_url": None,
|
|
},
|
|
"openrouter": {
|
|
"api_key": get_env_var("OPENROUTER_API_KEY", is_expert),
|
|
"base_url": "https://openrouter.ai/api/v1",
|
|
},
|
|
"openai-compatible": {
|
|
"api_key": get_env_var("OPENAI_API_KEY", is_expert),
|
|
"base_url": get_env_var("OPENAI_API_BASE", is_expert),
|
|
},
|
|
"gemini": {
|
|
"api_key": get_env_var("GEMINI_API_KEY", is_expert),
|
|
"base_url": None,
|
|
},
|
|
"deepseek": {
|
|
"api_key": get_env_var("DEEPSEEK_API_KEY", is_expert),
|
|
"base_url": "https://api.deepseek.com",
|
|
},
|
|
}
|
|
config = configs.get(provider, {})
|
|
if not config or not config.get("api_key"):
|
|
raise ValueError(
|
|
f"Missing required environment variable for provider: {provider}"
|
|
)
|
|
return config
|
|
|
|
|
|
def create_llm_client(
|
|
provider: str,
|
|
model_name: str,
|
|
temperature: Optional[float] = None,
|
|
is_expert: bool = False,
|
|
) -> BaseChatModel:
|
|
"""Create a language model client with appropriate configuration.
|
|
|
|
Args:
|
|
provider: The LLM provider to use
|
|
model_name: Name of the model to use
|
|
temperature: Optional temperature setting (0.0-2.0)
|
|
is_expert: Whether this is an expert model (uses deterministic output)
|
|
|
|
Returns:
|
|
Configured language model client
|
|
"""
|
|
config = get_provider_config(provider, is_expert)
|
|
if not config:
|
|
raise ValueError(f"Unsupported provider: {provider}")
|
|
|
|
logger.debug(
|
|
"Creating LLM client with provider=%s, model=%s, temperature=%s, expert=%s",
|
|
provider,
|
|
model_name,
|
|
temperature,
|
|
is_expert,
|
|
)
|
|
|
|
# Get model configuration
|
|
model_config = models_params.get(provider, {}).get(model_name, {})
|
|
|
|
# Default to True for known providers that support temperature if not specified
|
|
if "supports_temperature" not in model_config:
|
|
model_config["supports_temperature"] = provider in known_temp_providers
|
|
|
|
supports_temperature = model_config["supports_temperature"]
|
|
|
|
# Handle temperature settings
|
|
if is_expert:
|
|
temp_kwargs = {"temperature": 0} if supports_temperature else {}
|
|
elif temperature is not None and supports_temperature:
|
|
temp_kwargs = {"temperature": temperature}
|
|
elif provider == "openai-compatible" and supports_temperature:
|
|
temp_kwargs = {"temperature": 0.3}
|
|
else:
|
|
temp_kwargs = {}
|
|
|
|
if provider == "deepseek":
|
|
return create_deepseek_client(
|
|
model_name=model_name,
|
|
api_key=config["api_key"],
|
|
base_url=config["base_url"],
|
|
**temp_kwargs,
|
|
is_expert=is_expert,
|
|
)
|
|
elif provider == "openrouter":
|
|
return create_openrouter_client(
|
|
model_name=model_name,
|
|
api_key=config["api_key"],
|
|
**temp_kwargs,
|
|
is_expert=is_expert,
|
|
)
|
|
elif provider == "openai":
|
|
openai_kwargs = {
|
|
"api_key": config["api_key"],
|
|
"model": model_name,
|
|
**temp_kwargs,
|
|
}
|
|
if is_expert:
|
|
openai_kwargs["reasoning_effort"] = "high"
|
|
return ChatOpenAI(**openai_kwargs)
|
|
elif provider == "anthropic":
|
|
return ChatAnthropic(
|
|
api_key=config["api_key"],
|
|
model_name=model_name,
|
|
**temp_kwargs,
|
|
)
|
|
elif provider == "openai-compatible":
|
|
return ChatOpenAI(
|
|
api_key=config["api_key"],
|
|
base_url=config["base_url"],
|
|
model=model_name,
|
|
**temp_kwargs,
|
|
)
|
|
elif provider == "gemini":
|
|
return ChatGoogleGenerativeAI(
|
|
api_key=config["api_key"],
|
|
model=model_name,
|
|
**temp_kwargs,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported provider: {provider}")
|
|
|
|
|
|
def initialize_llm(
|
|
provider: str, model_name: str, temperature: float | None = None
|
|
) -> BaseChatModel:
|
|
"""Initialize a language model client based on the specified provider and model."""
|
|
return create_llm_client(provider, model_name, temperature, is_expert=False)
|
|
|
|
|
|
def initialize_expert_llm(provider: str, model_name: str) -> BaseChatModel:
|
|
"""Initialize an expert language model client based on the specified provider and model."""
|
|
return create_llm_client(provider, model_name, temperature=None, is_expert=True)
|
|
|
|
|
|
def validate_provider_env(provider: str) -> bool:
|
|
"""Check if the required environment variables for a provider are set."""
|
|
required_vars = {
|
|
"openai": "OPENAI_API_KEY",
|
|
"anthropic": "ANTHROPIC_API_KEY",
|
|
"openrouter": "OPENROUTER_API_KEY",
|
|
"openai-compatible": "OPENAI_API_KEY",
|
|
"gemini": "GEMINI_API_KEY",
|
|
"deepseek": "DEEPSEEK_API_KEY",
|
|
}
|
|
key = required_vars.get(provider.lower())
|
|
if key:
|
|
return bool(os.getenv(key))
|
|
return False
|
|
|
|
|
|
def merge_chat_history(
|
|
original_history: List[BaseMessage], fallback_history: List[BaseMessage]
|
|
) -> List[BaseMessage]:
|
|
"""Merge original and fallback chat histories while preserving order.
|
|
|
|
Args:
|
|
original_history: The original chat message history
|
|
fallback_history: Additional messages from fallback attempts
|
|
|
|
Returns:
|
|
List[BaseMessage]: Combined message history preserving chronological order
|
|
|
|
Note:
|
|
The function appends fallback messages to maintain context for future
|
|
interactions while preserving the original conversation flow.
|
|
"""
|
|
return original_history + fallback_history
|