aria-code 0.52.0__py3-none-any.whl

aria_code/providers/llm/custom.py

@@ -0,0 +1,158 @@
+"""声明式自定义 provider —— 用户只写配置文件，不写 Python 代码。
+
+背景：registry.py 里内置了 12 个 provider，但它们的名字、base_url、环境变量
+都是**硬编码**的。用户要接一个未列出的服务（自建 vLLM、公司内网网关、
+Arthera 自己的云端 API、任何 OpenAI 兼容端点），此前唯一的办法是写一个
+BaseLLMProvider 子类再调 register_provider()——对"只想填个地址和 key"的
+用户门槛太高。
+
+设计对标 codex 的 `model_providers`（codex-rs/model-provider-info/src/lib.rs）：
+配置里**只写环境变量名，不写密钥值**。codex 那边同时提供了明文
+`experimental_bearer_token` 字段，但源码注释明确写着 "Use of this config is
+discouraged in favor of `env_key` for security reasons"——这里不提供明文字段，
+只认 env_key，理由见下面 build_custom_provider 的说明。
+
+用户配置示例（~/.aria/providers.yaml 或 ~/.arthera/providers.json）：
+
+    llm:
+      default: mycorp/qwen-72b
+
+    model_providers:
+      mycorp:
+        name: 公司内网推理网关
+        base_url: https://llm.internal.example.com/v1
+        env_key: MYCORP_LLM_TOKEN
+        env_key_instructions: 找基础架构组申请，写进 ~/.aria/.env
+        model: qwen-72b
+        http_headers:
+          X-Tenant: research
+        env_http_headers:
+          X-Trace-Id: MYCORP_TRACE_ID
+
+      arthera-cloud:
+        name: Arthera 云端（Cloud Run）
+        base_url: https://api.arthera.example/api/v2
+        env_key: ARTHERA_API_TOKEN
+
+声明后即可像内置 provider 一样使用：default / fallback / code_tasks 里都能写
+`mycorp/<model>`。
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, Optional
+
+from .base import ProviderConfig
+from .openai_compat import OpenAICompatProvider
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["CustomProviderSpec", "parse_custom_providers", "build_custom_provider"]
+
+# 配置里出现这些键一律拒绝：它们意味着用户把密钥明文写进了配置文件。
+# 配置文件常被连同项目一起备份、同步、甚至误提交，明文密钥的暴露面比环境
+# 变量大一个量级。codex 提供了明文字段但自己在注释里劝退；这里直接不提供，
+# 并在检测到时给出可操作的替代方案，而不是静默忽略——静默忽略会让用户以为
+# 配好了，实际请求全部 401，排查成本更高。
+_FORBIDDEN_INLINE_SECRET_KEYS = ("api_key", "apikey", "token", "bearer_token", "secret")
+
+
+class CustomProviderSpec:
+    """一个用户声明的 provider。"""
+
+    __slots__ = ("id", "name", "base_url", "env_key", "env_key_instructions",
+                 "model", "http_headers", "env_http_headers", "timeout")
+
+    def __init__(self, provider_id: str, raw: Dict[str, Any]):
+        self.id = provider_id
+        self.name = str(raw.get("name") or provider_id)
+        self.base_url = str(raw.get("base_url") or "").rstrip("/")
+        self.env_key = raw.get("env_key") or None
+        self.env_key_instructions = raw.get("env_key_instructions") or ""
+        self.model = raw.get("model") or None
+        self.http_headers = dict(raw.get("http_headers") or {})
+        self.env_http_headers = dict(raw.get("env_http_headers") or {})
+        self.timeout = int(raw.get("timeout") or 120)
+
+    def resolved_headers(self) -> Dict[str, str]:
+        """静态头 + 从环境变量取值的头。环境变量缺失时**不发**该头，
+        而不是发一个空值——空值头在很多网关那里会被当成"显式声明为空"，
+        比缺失更难排查。"""
+        headers = dict(self.http_headers)
+        for header_name, env_name in self.env_http_headers.items():
+            value = os.environ.get(str(env_name), "").strip()
+            if value:
+                headers[str(header_name)] = value
+        return headers
+
+    def api_key(self) -> str:
+        return os.environ.get(str(self.env_key), "").strip() if self.env_key else ""
+
+    def missing_key_hint(self) -> str:
+        base = f"provider '{self.id}' 需要环境变量 {self.env_key}，当前未设置或为空。"
+        return f"{base} {self.env_key_instructions}".strip()
+
+
+def parse_custom_providers(cfg: Dict[str, Any]) -> Dict[str, CustomProviderSpec]:
+    """从用户配置解析 model_providers 段。
+
+    配置有问题时记录 warning 并跳过该条，不抛异常——一个写错的自定义 provider
+    不该让整个 CLI 起不来，其余 provider 应当照常可用。
+    """
+    raw_section = cfg.get("model_providers") if isinstance(cfg, dict) else None
+    if not isinstance(raw_section, dict):
+        return {}
+
+    specs: Dict[str, CustomProviderSpec] = {}
+    for provider_id, raw in raw_section.items():
+        pid = str(provider_id).strip().lower()
+        if not pid or not isinstance(raw, dict):
+            logger.warning("model_providers 中跳过无效条目: %r", provider_id)
+            continue
+
+        inline = [k for k in raw if str(k).lower() in _FORBIDDEN_INLINE_SECRET_KEYS]
+        if inline:
+            logger.warning(
+                "model_providers['%s'] 含明文密钥字段 %s —— 已忽略该 provider。"
+                "请改用 env_key 指向环境变量名，把值写进 ~/.aria/.env："
+                "\n  %s:\n    base_url: ...\n    env_key: %s_API_KEY",
+                pid, inline, pid, pid.upper().replace("-", "_"),
+            )
+            continue
+
+        if not raw.get("base_url"):
+            logger.warning("model_providers['%s'] 缺少 base_url —— 已跳过", pid)
+            continue
+
+        specs[pid] = CustomProviderSpec(pid, raw)
+
+    if specs:
+        logger.info("✓ 已加载 %d 个自定义 provider: %s", len(specs), ", ".join(sorted(specs)))
+    return specs
+
+
+def build_custom_provider(spec: CustomProviderSpec, model: Optional[str] = None):
+    """按 spec 构造一个 OpenAI 兼容 provider 实例。
+
+    复用 OpenAICompatProvider 而不是新写一套 HTTP 客户端：绝大多数第三方与
+    自建推理服务（vLLM / LM Studio / Ollama 的 OpenAI 兼容层 / 各家云厂商网关）
+    都实现了 /chat/completions，重写一遍只会多一处要维护的重试与流式解析。
+    """
+    config = ProviderConfig(
+        name=spec.id,
+        api_key=spec.api_key() or None,
+        base_url=spec.base_url,
+        model=model or spec.model,
+        timeout=spec.timeout,
+        extra={"http_headers": spec.resolved_headers(), "custom": True},
+    )
+
+    class _DeclaredProvider(OpenAICompatProvider):
+        provider_name = spec.id
+        DEFAULT_BASE_URL = spec.base_url
+        DEFAULT_MODEL = spec.model or ""
+
+    _DeclaredProvider.__name__ = f"Declared_{spec.id.replace('-', '_')}"
+    return _DeclaredProvider(config)

aria_code/providers/llm/ollama.py

@@ -0,0 +1,129 @@
+"""
+providers/llm/ollama.py — Ollama 本地 LLM Provider
+====================================================
+连接本地运行的 Ollama 服务，支持工具调用（native + text-based fallback）。
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, AsyncIterator, Dict, List, Optional
+
+from .base import BaseLLMProvider, Message, ProviderConfig
+
+logger = logging.getLogger(__name__)
+
+
+class OllamaProvider(BaseLLMProvider):
+
+    provider_name    = "ollama"
+    supports_tools   = True
+    supports_thinking = False
+    local            = True
+
+    DEFAULT_URL = "http://localhost:11434"
+
+    def __init__(self, config: ProviderConfig):
+        super().__init__(config)
+        self.base_url = config.base_url or self.DEFAULT_URL
+        self.model    = config.model or "qwen2.5:7b"
+
+    async def is_available(self) -> bool:
+        """探测 Ollama 服务是否在线"""
+        import urllib.request
+        for url in [self.base_url,
+                    self.base_url.replace("localhost", "127.0.0.1")]:
+            try:
+                urllib.request.urlopen(f"{url}/api/tags", timeout=2).close()
+                return True
+            except Exception:
+                continue
+        return False
+
+    async def list_models(self) -> List[str]:
+        """返回已安装的模型列表"""
+        import urllib.request
+        import json as _json
+        try:
+            with urllib.request.urlopen(f"{self.base_url}/api/tags", timeout=3) as r:
+                data = _json.loads(r.read())
+                return [m["name"] for m in data.get("models", [])]
+        except Exception:
+            return []
+
+    async def stream(
+        self,
+        messages: List[Message],
+        tools: Optional[List[Dict]] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        cancel_event=None,
+    ) -> AsyncIterator[Dict[str, Any]]:
+        import aiohttp
+
+        temp     = temperature if temperature is not None else self.config.temperature
+        n_tokens = max_tokens  if max_tokens  is not None else self.config.max_tokens
+
+        payload = {
+            "model":   self.model,
+            "messages": [{"role": m.role, "content": m.content} for m in messages],
+            "stream":  True,
+            "options": {"temperature": temp, "num_predict": n_tokens, "num_ctx": 32768},
+        }
+        if tools:
+            # 只有支持 native tool call 的模型才注入（其他走 text-based 解析）
+            payload["tools"] = tools
+
+        url = f"{self.base_url}/api/chat"
+        try:
+            async with aiohttp.ClientSession() as sess:
+                async with sess.post(
+                    url, json=payload,
+                    timeout=aiohttp.ClientTimeout(total=self.config.timeout)
+                ) as resp:
+                    if resp.status != 200:
+                        body = await resp.text()
+                        yield {"type": "error",
+                               "message": f"Ollama HTTP {resp.status}: {body[:200]}"}
+                        return
+
+                    async for raw in resp.content:
+                        if cancel_event and cancel_event.is_set():
+                            return
+                        line = raw.decode("utf-8", errors="ignore").strip()
+                        if not line:
+                            continue
+                        try:
+                            data = json.loads(line)
+                        except json.JSONDecodeError:
+                            continue
+
+                        # Native tool calls
+                        tc_list = (data.get("message") or {}).get("tool_calls") or []
+                        for tc in tc_list:
+                            fn = tc.get("function") or {}
+                            args = fn.get("arguments", {})
+                            if isinstance(args, str):
+                                try:
+                                    args = json.loads(args)
+                                except json.JSONDecodeError:
+                                    args = {"_raw": args}
+                            yield {"type": "tool_call",
+                                   "name": fn.get("name", ""),
+                                   "arguments": args}
+
+                        # Text token
+                        token = (data.get("message") or {}).get("content", "")
+                        if token:
+                            yield {"type": "token", "text": token}
+
+                        if data.get("done"):
+                            yield {"type": "done"}
+                            return
+
+        except aiohttp.ClientConnectorError as e:
+            yield {"type": "error",
+                   "message": f"Ollama 连接失败 ({self.base_url}): {e}"}
+        except Exception as e:
+            yield {"type": "error", "message": f"Ollama 错误: {e}"}

aria_code/providers/llm/openai_compat.py

@@ -0,0 +1,313 @@
+"""
+providers/llm/openai_compat.py — OpenAI 兼容协议通用 Provider
+==============================================================
+DeepSeek / OpenAI / Groq / Together / LM Studio / vLLM / llama.cpp
+全部走同一套 /v1/chat/completions SSE 协议。
+
+各 provider 只需继承并声明 DEFAULT_BASE_URL 和 DEFAULT_MODEL。
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+from typing import Any, AsyncIterator, Dict, List, Optional
+
+from .base import BaseLLMProvider, Message, ProviderConfig
+
+logger = logging.getLogger(__name__)
+
+
+def chat_completions_url(base_url: str) -> str:
+    """Join an OpenAI-compatible base URL without duplicating its version."""
+    base = base_url.rstrip("/")
+    if re.search(r"/v\d+(?:beta)?$", base):
+        return f"{base}/chat/completions"
+    return f"{base}/v1/chat/completions"
+
+
+class OpenAICompatProvider(BaseLLMProvider):
+    """
+    通用 OpenAI 兼容 Provider。
+    子类覆盖 DEFAULT_BASE_URL / DEFAULT_MODEL / provider_name。
+    """
+
+    provider_name    = "openai_compat"
+    supports_tools   = True
+    supports_thinking = False
+    local            = False
+
+    DEFAULT_BASE_URL = "https://api.openai.com"
+    DEFAULT_MODEL    = "gpt-4o-mini"
+
+    def __init__(self, config: ProviderConfig):
+        super().__init__(config)
+        self.base_url = (config.base_url or self.DEFAULT_BASE_URL).rstrip("/")
+        self.model    = config.model or self.DEFAULT_MODEL
+        self.api_key  = config.api_key or ""
+
+    async def is_available(self) -> bool:
+        return bool(self.api_key)
+
+    async def stream(
+        self,
+        messages: List[Message],
+        tools: Optional[List[Dict]] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        cancel_event=None,
+    ) -> AsyncIterator[Dict[str, Any]]:
+        import aiohttp
+
+        temp     = temperature if temperature is not None else self.config.temperature
+        n_tokens = max_tokens  if max_tokens  is not None else self.config.max_tokens
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type":  "application/json",
+        }
+        payload: Dict[str, Any] = {
+            "model":       self.model,
+            "messages":    [{"role": m.role, "content": m.content} for m in messages],
+            "temperature": temp,
+            "max_tokens":  n_tokens,
+            "stream":      True,
+        }
+        if tools:
+            payload["tools"] = [
+                {"type": "function", "function": t} for t in tools
+            ]
+            payload["tool_choice"] = "auto"
+
+        # Some compatible endpoints already include their API version
+        # (Moonshot /v1, ZhiPu /v4). Avoid generating /v1/v1 or /v4/v1.
+        url = chat_completions_url(self.base_url)
+        # aiohttp 不自动读系统代理，需显式传入
+        proxy = (os.getenv("HTTPS_PROXY") or os.getenv("https_proxy")
+                 or os.getenv("HTTP_PROXY")  or os.getenv("http_proxy"))
+
+        try:
+            async with aiohttp.ClientSession() as sess:
+                async with sess.post(
+                    url, json=payload, headers=headers,
+                    proxy=proxy,
+                    timeout=aiohttp.ClientTimeout(total=self.config.timeout)
+                ) as resp:
+                    if resp.status != 200:
+                        body = await resp.text()
+                        yield {"type": "error",
+                               "message": f"HTTP {resp.status}: {body[:300]}"}
+                        return
+
+                    pending_tool: Dict = {}
+                    async for raw in resp.content:
+                        if cancel_event and cancel_event.is_set():
+                            return
+
+                        line = raw.decode("utf-8", errors="ignore").strip()
+                        if not line or line == "data: [DONE]":
+                            continue
+                        if not line.startswith("data: "):
+                            continue
+
+                        try:
+                            data = json.loads(line[6:])
+                        except json.JSONDecodeError:
+                            continue
+
+                        # 处理 thinking tokens（DeepSeek-R1 等）
+                        reasoning = (
+                            ((data.get("choices") or [{}])[0]
+                             .get("delta") or {})
+                            .get("reasoning_content")
+                        )
+                        if reasoning:
+                            yield {"type": "thinking", "text": reasoning}
+
+                        # 普通 token
+                        delta = ((data.get("choices") or [{}])[0].get("delta") or {})
+                        token = delta.get("content") or ""
+                        if token:
+                            yield {"type": "token", "text": token}
+
+                        # 工具调用（SSE 增量拼接）
+                        for tc in delta.get("tool_calls") or []:
+                            fn   = tc.get("function") or {}
+                            idx  = tc.get("index", 0)
+                            name = fn.get("name", "")
+                            args_chunk = fn.get("arguments", "")
+
+                            if name:
+                                if pending_tool:
+                                    # emit previous tool
+                                    try:
+                                        args = json.loads(pending_tool["args"])
+                                    except Exception:
+                                        args = {"_raw": pending_tool["args"]}
+                                    yield {"type": "tool_call",
+                                           "name": pending_tool["name"],
+                                           "arguments": args}
+                                pending_tool = {"name": name, "args": "", "idx": idx}
+                            if args_chunk:
+                                pending_tool["args"] = pending_tool.get("args", "") + args_chunk
+
+                        # finish reason
+                        finish = ((data.get("choices") or [{}])[0].get("finish_reason"))
+                        if finish in ("stop", "tool_calls", "length"):
+                            if pending_tool:
+                                try:
+                                    args = json.loads(pending_tool["args"])
+                                except Exception:
+                                    args = {"_raw": pending_tool["args"]}
+                                yield {"type": "tool_call",
+                                       "name": pending_tool["name"],
+                                       "arguments": args}
+                                pending_tool = {}
+                            yield {"type": "done"}
+                            return
+
+        except aiohttp.ClientConnectorError as e:
+            yield {"type": "error", "message": f"连接 {self.base_url} 失败: {e}"}
+        except Exception as e:
+            yield {"type": "error", "message": f"Provider 错误: {e}"}
+
+
+# ── 具体 Provider 子类（只需声明几个属性）────────────────────────────────────
+
+class DeepSeekProvider(OpenAICompatProvider):
+    provider_name    = "deepseek"
+    DEFAULT_BASE_URL = "https://api.deepseek.com"
+    DEFAULT_MODEL    = "deepseek-chat"
+    supports_thinking = True   # deepseek-reasoner 支持
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("DEEPSEEK_API_KEY", "")
+        super().__init__(config)
+        # 思考模型别名
+        if self.model in ("deepseek-reasoner", "deepseek-r1"):
+            self.supports_thinking = True
+
+
+class OpenAIProvider(OpenAICompatProvider):
+    provider_name    = "openai"
+    DEFAULT_BASE_URL = "https://api.openai.com"
+    DEFAULT_MODEL    = "gpt-4o-mini"
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("OPENAI_API_KEY", "")
+        super().__init__(config)
+
+
+class GroqProvider(OpenAICompatProvider):
+    provider_name    = "groq"
+    DEFAULT_BASE_URL = "https://api.groq.com/openai"
+    DEFAULT_MODEL    = "llama-3.3-70b-versatile"
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("GROQ_API_KEY", "")
+        super().__init__(config)
+
+
+class GoogleProvider(OpenAICompatProvider):
+    """Gemini through Google's OpenAI-compatible endpoint.
+
+    The registry previously had no Google entry at all, so Gemini — the default
+    model in the shipped config — could not be reached through the provider
+    chain and never appeared in fallback.  Vertex AI (service-account / ADC) is
+    handled separately by apps/cli/providers/vertexai_stream.py; this covers the
+    API-key path.
+    """
+
+    provider_name    = "google"
+    DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai"
+    DEFAULT_MODEL    = "gemini-2.5-flash"
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("GEMINI_API_KEY", "") or os.getenv("GOOGLE_API_KEY", "")
+        super().__init__(config)
+
+
+class TogetherProvider(OpenAICompatProvider):
+    provider_name    = "together"
+    DEFAULT_BASE_URL = "https://api.together.xyz"
+    DEFAULT_MODEL    = "meta-llama/Llama-3.3-70B-Instruct-Turbo"
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("TOGETHER_API_KEY", "")
+        super().__init__(config)
+
+
+class DashScopeProvider(OpenAICompatProvider):
+    """阿里云通义（OpenAI 兼容端点）"""
+    provider_name    = "dashscope"
+    DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode"
+    DEFAULT_MODEL    = "qwen-plus"
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("DASHSCOPE_API_KEY", "")
+        super().__init__(config)
+
+
+class LMStudioProvider(OpenAICompatProvider):
+    """LM Studio 本地服务"""
+    provider_name    = "lmstudio"
+    DEFAULT_BASE_URL = "http://localhost:1234"
+    DEFAULT_MODEL    = "loaded-model"
+    local            = True
+
+    async def is_available(self) -> bool:
+        import urllib.request
+        try:
+            urllib.request.urlopen(
+                f"{self.base_url}/v1/models", timeout=2
+            ).close()
+            return True
+        except Exception:
+            return False
+
+
+# ── 国内可访问 Provider（OpenAI 兼容协议）────────────────────────────────────
+
+class SiliconFlowProvider(OpenAICompatProvider):
+    """硅基流动 — 中国大陆可直连，支持 DeepSeek-V3/R1、Qwen 等主流模型"""
+    provider_name    = "siliconflow"
+    DEFAULT_BASE_URL = "https://api.siliconflow.cn"
+    DEFAULT_MODEL    = "deepseek-ai/DeepSeek-V3"
+    supports_thinking = True   # DeepSeek-R1 在此运行
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("SILICONFLOW_API_KEY", "")
+        super().__init__(config)
+
+
+class MoonshotProvider(OpenAICompatProvider):
+    """Moonshot / Kimi — 中国大陆可直连"""
+    provider_name    = "moonshot"
+    DEFAULT_BASE_URL = "https://api.moonshot.cn/v1"
+    DEFAULT_MODEL    = "moonshot-v1-8k"
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("MOONSHOT_API_KEY", "")
+        super().__init__(config)
+
+
+class ZhiPuProvider(OpenAICompatProvider):
+    """智谱 GLM — 中国大陆可直连"""
+    provider_name    = "zhipu"
+    DEFAULT_BASE_URL = "https://open.bigmodel.cn/api/paas/v4"
+    DEFAULT_MODEL    = "glm-4-flash"
+
+    def __init__(self, config: ProviderConfig):
+        if not config.api_key:
+            config.api_key = os.getenv("ZHIPUAI_API_KEY", "")
+        super().__init__(config)