237 lines
9.0 KiB
Python
Executable File
237 lines
9.0 KiB
Python
Executable File
import os
|
|
import time
|
|
from typing import AsyncIterator, Iterator
|
|
|
|
from .base import BaseLLMProvider, LLMMessage, LLMResponse
|
|
from .providers.deepseek import DeepSeekProvider
|
|
from .providers.glm import GLMProvider
|
|
from .providers.minimax import MiniMaxProvider
|
|
from .providers.ollama import OllamaProvider
|
|
|
|
|
|
class LLMManager:
|
|
PROVIDERS = {
|
|
"glm": GLMProvider,
|
|
"minimax": MiniMaxProvider,
|
|
"deepseek": DeepSeekProvider,
|
|
"ollama": OllamaProvider,
|
|
}
|
|
|
|
DEFAULT_MODELS = {
|
|
"glm": "glm-4-flash",
|
|
"minimax": "abab6.5s-chat",
|
|
"deepseek": "deepseek-v4-flash",
|
|
"ollama": "qwen3:0.6b",
|
|
}
|
|
|
|
def __init__(self, config: dict):
|
|
self.config = config
|
|
self._providers: dict[str, BaseLLMProvider] = {}
|
|
self._available_providers: list[str] | None = None
|
|
self._default_provider = config.get("default", {}).get("provider", "glm")
|
|
self._default_model = config.get("default", {}).get("model", "glm-4-flash")
|
|
self._default_temperature = config.get("default", {}).get("temperature", 0.7)
|
|
self._default_max_tokens = config.get("default", {}).get("max_tokens", 4096)
|
|
|
|
def _get_api_key(self, env_key: str) -> str | None:
|
|
provider_key = ""
|
|
if env_key:
|
|
provider_key = self.config.get(env_key.lower(), "") or os.environ.get(env_key, "")
|
|
if provider_key:
|
|
return provider_key
|
|
|
|
# Backward-compatible fallback for OpenAI-compatible providers.
|
|
openai_key = self.config.get("openai_api_key", "") or os.environ.get("OPENAI_API_KEY", "")
|
|
if openai_key:
|
|
return openai_key
|
|
|
|
return ""
|
|
|
|
def _check_ollama_available(self) -> bool:
|
|
try:
|
|
import requests
|
|
response = requests.get("http://localhost:11434/api/tags", timeout=3)
|
|
return response.status_code == 200
|
|
except Exception:
|
|
return False
|
|
|
|
def _get_available_providers(self) -> list[str]:
|
|
if self._available_providers is not None:
|
|
return self._available_providers
|
|
|
|
available = []
|
|
providers_config = self.config.get("providers", {})
|
|
|
|
# 检查是否有阿里云百炼的openai_api_key
|
|
openai_key = self.config.get("openai_api_key", "")
|
|
|
|
for name, cfg in providers_config.items():
|
|
if name == "ollama":
|
|
if self._check_ollama_available():
|
|
available.append(name)
|
|
else:
|
|
env_key = cfg.get("env_key", "")
|
|
if self._get_api_key(env_key) or openai_key:
|
|
available.append(name)
|
|
|
|
self._available_providers = available
|
|
return available
|
|
|
|
def get_provider(
|
|
self,
|
|
provider_name: str | None = None,
|
|
model: str | None = None,
|
|
temperature: float | None = None,
|
|
max_tokens: int | None = None,
|
|
allow_fallback: bool = True,
|
|
) -> BaseLLMProvider:
|
|
provider_name = provider_name or self._default_provider
|
|
model = model or self._default_model
|
|
temperature = temperature or self._default_temperature
|
|
max_tokens = max_tokens or self._default_max_tokens
|
|
|
|
cache_key = f"{provider_name}:{model}"
|
|
# 强制重新创建provider，避免缓存问题
|
|
if cache_key in self._providers:
|
|
del self._providers[cache_key]
|
|
|
|
provider_config = self.config.get("providers", {}).get(provider_name, {})
|
|
env_key = provider_config.get("env_key", "")
|
|
api_key = self._get_api_key(env_key)
|
|
|
|
if not api_key:
|
|
if allow_fallback:
|
|
available = self._get_available_providers()
|
|
if available:
|
|
provider_name = available[0]
|
|
model = self.DEFAULT_MODELS.get(provider_name, "")
|
|
provider_config = self.config.get("providers", {}).get(provider_name, {})
|
|
env_key = provider_config.get("env_key", "")
|
|
api_key = self._get_api_key(env_key)
|
|
cache_key = f"{provider_name}:{model}"
|
|
if cache_key in self._providers:
|
|
return self._providers[cache_key]
|
|
else:
|
|
raise ValueError(f"API key not found: {env_key}")
|
|
|
|
if not api_key and provider_name != "ollama":
|
|
raise ValueError(
|
|
"No LLM API key configured. Please set at least one of: "
|
|
"ZHIPUAI_API_KEY, MOONSHOT_API_KEY, MINIMAX_API_KEY, DEEPSEEK_API_KEY"
|
|
)
|
|
|
|
provider_class = self.PROVIDERS.get(provider_name)
|
|
if not provider_class:
|
|
raise ValueError(f"Unknown provider: {provider_name}")
|
|
|
|
provider_kwargs = {
|
|
"api_key": api_key,
|
|
"model": model,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
provider_base_url = provider_config.get("api_base") or self.config.get("openai_base_url")
|
|
if provider_base_url:
|
|
import inspect
|
|
|
|
signature = inspect.signature(provider_class.__init__)
|
|
if "api_base" in signature.parameters:
|
|
provider_kwargs["api_base"] = provider_base_url
|
|
|
|
provider = provider_class(**provider_kwargs)
|
|
self._providers[cache_key] = provider
|
|
return provider
|
|
|
|
def get_available_providers_info(self) -> list[dict]:
|
|
available = self._get_available_providers()
|
|
result = []
|
|
for name in available:
|
|
cfg = self.config.get("providers", {}).get(name, {})
|
|
result.append({
|
|
"name": name,
|
|
"display_name": cfg.get("name", name),
|
|
"default_model": self.DEFAULT_MODELS.get(name, ""),
|
|
"models": [m.get("name") for m in cfg.get("models", [])],
|
|
})
|
|
return result
|
|
|
|
def _track_llm_cost(self, response: LLMResponse, messages: list[LLMMessage], duration_ms: int, agent_name: str = "LLMManager") -> None:
|
|
try:
|
|
from gfa_core.mas.cost_tracker import track_cost
|
|
|
|
usage = response.usage or {}
|
|
input_tokens = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
|
|
output_tokens = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
|
|
if not input_tokens and messages:
|
|
input_tokens = sum(len(m.content or "") for m in messages) // 4
|
|
if not output_tokens and response.content:
|
|
output_tokens = len(response.content) // 4
|
|
track_cost(
|
|
agent_name=agent_name,
|
|
provider=response.provider or self._default_provider,
|
|
model=response.model or self._default_model,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
duration_ms=duration_ms,
|
|
query=(messages[-1].content if messages else "")[:100],
|
|
)
|
|
except Exception as exc:
|
|
print(f"[CostTracker] LLMManager track failed: {exc}")
|
|
|
|
def call(
|
|
self,
|
|
messages: list[LLMMessage],
|
|
provider: str | None = None,
|
|
model: str | None = None,
|
|
allow_fallback: bool = True,
|
|
**kwargs,
|
|
) -> LLMResponse:
|
|
track_usage = bool(kwargs.pop("track_cost", True))
|
|
track_agent_name = str(kwargs.pop("track_agent_name", "LLMManager") or "LLMManager")
|
|
p = self.get_provider(provider, model, allow_fallback=allow_fallback)
|
|
start_time = time.time()
|
|
response = p.call(messages, **kwargs)
|
|
if track_usage:
|
|
self._track_llm_cost(response, messages, int((time.time() - start_time) * 1000), track_agent_name)
|
|
return response
|
|
|
|
async def acall(
|
|
self,
|
|
messages: list[LLMMessage],
|
|
provider: str | None = None,
|
|
model: str | None = None,
|
|
allow_fallback: bool = True,
|
|
**kwargs,
|
|
) -> LLMResponse:
|
|
track_usage = bool(kwargs.pop("track_cost", True))
|
|
track_agent_name = str(kwargs.pop("track_agent_name", "LLMManager") or "LLMManager")
|
|
p = self.get_provider(provider, model, allow_fallback=allow_fallback)
|
|
start_time = time.time()
|
|
response = await p.acall(messages, **kwargs)
|
|
if track_usage:
|
|
self._track_llm_cost(response, messages, int((time.time() - start_time) * 1000), track_agent_name)
|
|
return response
|
|
|
|
def stream(
|
|
self,
|
|
messages: list[LLMMessage],
|
|
provider: str | None = None,
|
|
model: str | None = None,
|
|
allow_fallback: bool = True,
|
|
**kwargs,
|
|
) -> Iterator[str]:
|
|
p = self.get_provider(provider, model, allow_fallback=allow_fallback)
|
|
return p.stream(messages, **kwargs)
|
|
|
|
async def astream(
|
|
self,
|
|
messages: list[LLMMessage],
|
|
provider: str | None = None,
|
|
model: str | None = None,
|
|
allow_fallback: bool = True,
|
|
**kwargs,
|
|
) -> AsyncIterator[str]:
|
|
p = self.get_provider(provider, model, allow_fallback=allow_fallback)
|
|
async for chunk in p.astream(messages, **kwargs):
|
|
yield chunk
|