|
|
@@ -1,12 +1,16 @@
|
|
|
"""
|
|
|
lambdagent.providers.openai_compat_provider — OpenAI-compatible API provider.
|
|
|
|
|
|
-Covers: OpenAI, Ollama, DashScope, DeepSeek, Moonshot, Zhipu, and any
|
|
|
+Covers: OpenAI, Ollama, DashScope/Qwen, DeepSeek, Moonshot, Zhipu, and any
|
|
|
OpenAI-compatible endpoint.
|
|
|
|
|
|
L02: Added chat_typed() for ChatMessage/ChatResponse support,
|
|
|
overrides provider_name and default_model properties. Added PRESETS
|
|
|
class attribute for easy provider lookup.
|
|
|
+
|
|
|
+L03: Token-usage accumulation in chat() so get_usage() / reset_session()
|
|
|
+work the same way as ClaudeCodeProvider, enabling cost tracking in agentpaas.
|
|
|
+"qwen" added as a friendly alias for "dashscope".
|
|
|
"""
|
|
|
from __future__ import annotations
|
|
|
|
|
|
@@ -18,14 +22,16 @@ from typing import Dict, List
|
|
|
from .base import LLMProvider, ProviderConfig, ProviderError, ChatMessage, ChatResponse
|
|
|
|
|
|
|
|
|
-# Default base URLs and env keys per provider
|
|
|
+# Default base URLs and env keys per provider.
|
|
|
+# "qwen" is a friendly alias for DashScope (Alibaba Cloud).
|
|
|
_PROVIDER_DEFAULTS = {
|
|
|
- "openai": ("https://api.openai.com/v1", "OPENAI_API_KEY"),
|
|
|
- "ollama": ("http://127.0.0.1:11434/v1", ""),
|
|
|
- "dashscope": ("https://dashscope.aliyuncs.com/compatible-mode/v1", "DASHSCOPE_API_KEY"),
|
|
|
- "deepseek": ("https://api.deepseek.com/v1", "DEEPSEEK_API_KEY"),
|
|
|
- "moonshot": ("https://api.moonshot.cn/v1", "MOONSHOT_API_KEY"),
|
|
|
- "zhipu": ("https://open.bigmodel.cn/api/paas/v4", "ZHIPU_API_KEY"),
|
|
|
+ "openai": ("https://api.openai.com/v1", "OPENAI_API_KEY"),
|
|
|
+ "ollama": ("http://127.0.0.1:11434/v1", ""),
|
|
|
+ "dashscope": ("https://dashscope.aliyuncs.com/compatible-mode/v1", "DASHSCOPE_API_KEY"),
|
|
|
+ "qwen": ("https://dashscope.aliyuncs.com/compatible-mode/v1", "DASHSCOPE_API_KEY"),
|
|
|
+ "deepseek": ("https://api.deepseek.com/v1", "DEEPSEEK_API_KEY"),
|
|
|
+ "moonshot": ("https://api.moonshot.cn/v1", "MOONSHOT_API_KEY"),
|
|
|
+ "zhipu": ("https://open.bigmodel.cn/api/paas/v4", "ZHIPU_API_KEY"),
|
|
|
}
|
|
|
|
|
|
|
|
|
@@ -34,12 +40,13 @@ class OpenAICompatProvider(LLMProvider):
|
|
|
|
|
|
# L02: Provider configs: name -> (base_url, default_model, api_key_env)
|
|
|
PRESETS = {
|
|
|
- "openai": ("https://api.openai.com/v1", "gpt-4o", "OPENAI_API_KEY"),
|
|
|
- "dashscope": ("https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-max", "DASHSCOPE_API_KEY"),
|
|
|
- "deepseek": ("https://api.deepseek.com/v1", "deepseek-chat", "DEEPSEEK_API_KEY"),
|
|
|
- "moonshot": ("https://api.moonshot.cn/v1", "moonshot-v1-128k", "MOONSHOT_API_KEY"),
|
|
|
- "ollama": ("http://127.0.0.1:11434/v1", "qwen2.5:7b", ""),
|
|
|
- "zhipu": ("https://open.bigmodel.cn/api/paas/v4", "glm-4", "ZHIPU_API_KEY"),
|
|
|
+ "openai": ("https://api.openai.com/v1", "gpt-4o", "OPENAI_API_KEY"),
|
|
|
+ "dashscope": ("https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-max", "DASHSCOPE_API_KEY"),
|
|
|
+ "qwen": ("https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-max", "DASHSCOPE_API_KEY"),
|
|
|
+ "deepseek": ("https://api.deepseek.com/v1", "deepseek-chat", "DEEPSEEK_API_KEY"),
|
|
|
+ "moonshot": ("https://api.moonshot.cn/v1", "moonshot-v1-128k", "MOONSHOT_API_KEY"),
|
|
|
+ "ollama": ("http://127.0.0.1:11434/v1", "qwen2.5:7b", ""),
|
|
|
+ "zhipu": ("https://open.bigmodel.cn/api/paas/v4", "glm-4", "ZHIPU_API_KEY"),
|
|
|
}
|
|
|
|
|
|
def __init__(self, config: ProviderConfig, provider_name: str = "openai",
|
|
|
@@ -61,6 +68,11 @@ class OpenAICompatProvider(LLMProvider):
|
|
|
if model:
|
|
|
self.config.model = model
|
|
|
|
|
|
+ # L03: Token-usage accumulator — mirrors ClaudeCodeProvider.get_usage()
|
|
|
+ # so agentpaas cost tracking works for any provider.
|
|
|
+ self._usage_input = 0
|
|
|
+ self._usage_output = 0
|
|
|
+
|
|
|
def chat(self, messages: List[Dict[str, str]]) -> str:
|
|
|
url = f"{self.base_url.rstrip('/')}/chat/completions"
|
|
|
|
|
|
@@ -71,9 +83,7 @@ class OpenAICompatProvider(LLMProvider):
|
|
|
"max_tokens": self.config.max_tokens,
|
|
|
}, ensure_ascii=False).encode("utf-8")
|
|
|
|
|
|
- headers = {
|
|
|
- "Content-Type": "application/json",
|
|
|
- }
|
|
|
+ headers = {"Content-Type": "application/json"}
|
|
|
if self.api_key:
|
|
|
headers["Authorization"] = f"Bearer {self.api_key}"
|
|
|
|
|
|
@@ -82,6 +92,10 @@ class OpenAICompatProvider(LLMProvider):
|
|
|
try:
|
|
|
with urllib.request.urlopen(req, timeout=self.config.timeout) as resp:
|
|
|
data = json.loads(resp.read())
|
|
|
+ # L03: accumulate usage so get_usage() returns session totals
|
|
|
+ usage = data.get("usage", {})
|
|
|
+ self._usage_input += usage.get("prompt_tokens", 0)
|
|
|
+ self._usage_output += usage.get("completion_tokens", 0)
|
|
|
return data["choices"][0]["message"]["content"].strip()
|
|
|
except urllib.error.HTTPError as e:
|
|
|
error_body = e.read().decode()[:500]
|
|
|
@@ -93,6 +107,29 @@ class OpenAICompatProvider(LLMProvider):
|
|
|
except Exception as e:
|
|
|
raise ProviderError(f"{self._provider_name} error: {e}", self._provider_name)
|
|
|
|
|
|
+ # ── L03: Session-level usage API (matches ClaudeCodeProvider interface) ──
|
|
|
+
|
|
|
+ def get_usage(self) -> dict:
|
|
|
+ """Return accumulated token usage for this session.
|
|
|
+
|
|
|
+ Keys match ClaudeCodeProvider.get_usage() so _execute_agent's cost
|
|
|
+ tracking works without any changes. Cache fields are zero (not
|
|
|
+ applicable to OpenAI-compatible endpoints).
|
|
|
+ """
|
|
|
+ return {
|
|
|
+ "input_tokens": self._usage_input,
|
|
|
+ "output_tokens": self._usage_output,
|
|
|
+ "cache_creation_input_tokens": 0,
|
|
|
+ "cache_read_input_tokens": 0,
|
|
|
+ # Provider tag lets _compute_cost_usd apply the right pricing table
|
|
|
+ "provider": self._provider_name,
|
|
|
+ }
|
|
|
+
|
|
|
+ def reset_session(self) -> None:
|
|
|
+ """Reset per-session usage counters (called between runs)."""
|
|
|
+ self._usage_input = 0
|
|
|
+ self._usage_output = 0
|
|
|
+
|
|
|
# ── L02: Unified interface overrides ──
|
|
|
|
|
|
def chat_typed(self, messages: List[ChatMessage],
|