my-pi-agent 0.1.0

package/src/my_agent_llm/auth/manager.py

@@ -0,0 +1,259 @@
+"""AuthManager: 管理 ~/.my-pi-agent/auth.json 的线程与进程安全凭据中心。"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+from pathlib import Path
+import time
+from typing import Any
+
+from filelock import FileLock  # pyright: ignore[reportMissingImports]
+import httpx
+
+from my_agent_llm.auth.schema import (
+    ApiKeyCredential,
+    AuthStore,
+    OAuthCredential,
+)
+
+__all__ = [
+    "DEFAULT_ANTIGRAVITY_CLIENT_ID",
+    "DEFAULT_ANTIGRAVITY_CLIENT_SECRET",
+    "GOOGLE_TOKEN_URL",
+    "AuthManager",
+]
+
+DEFAULT_ANTIGRAVITY_CLIENT_ID = os.environ.get(
+    "ANTIGRAVITY_CLIENT_ID",
+    "my-pi-agent-desktop-client-id",
+)
+DEFAULT_ANTIGRAVITY_CLIENT_SECRET = os.environ.get(
+    "ANTIGRAVITY_CLIENT_SECRET",
+    "my-pi-agent-desktop-client-secret",
+)
+GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
+
+
+class AuthManager:
+    """产品级全局凭据管理器。
+
+    提供线程安全与跨进程文件锁保护的 auth.json 读写、Profile 调度与 OAuth 静默续期。
+    """
+
+    def __init__(
+        self,
+        auth_path: Path | str | None = None,
+        lock_path: Path | str | None = None,
+        lock_timeout: float = 5.0,
+    ) -> None:
+        if auth_path:
+            self.auth_path = Path(auth_path).resolve()
+        else:
+            home = Path(
+                os.environ.get("MY_AGENT_HOME")
+                or Path(os.environ.get("USERPROFILE") or os.environ.get("HOME") or "~").expanduser() / ".my-pi-agent"
+            ).resolve()
+            self.auth_path = home / "auth.json"
+
+        if lock_path:
+            self.lock_path = Path(lock_path).resolve()
+        else:
+            self.lock_path = self.auth_path.with_name(f"{self.auth_path.name}.lock")
+        self.lock_timeout = lock_timeout
+
+    def load_store(self) -> AuthStore:
+        """从 auth.json 加载并自动归一化格式。
+
+        支持标准的 AuthStore 结构以及简写平铺格式。如果文件不存在或损坏，优雅降级为空 AuthStore。
+        """
+        if not self.auth_path.exists():
+            return AuthStore()
+
+        try:
+            content = self.auth_path.read_text(encoding="utf-8")
+            data = json.loads(content)
+            if not isinstance(data, dict):
+                return AuthStore()
+
+            # 兼容扁平平铺简写格式：{"deepseek": {"type": "api_key", ...}}
+            if "providers" not in data:
+                normalized_providers: dict[str, Any] = {}
+                for prov, item in data.items():
+                    if isinstance(item, dict) and "type" in item:
+                        normalized_providers[prov] = {"default": item}
+                active_profiles = data.get("active_profiles")
+                if not isinstance(active_profiles, dict):
+                    active_profiles = {}
+                version = data.get("version", 1)
+                return AuthStore(
+                    version=version,
+                    active_profiles=active_profiles,
+                    providers=normalized_providers,
+                )
+            return AuthStore.model_validate(data)
+        except Exception:
+            return AuthStore()
+
+    def save_store(self, store: AuthStore) -> None:
+        """以 0o600 权限与临时文件原子替换保存凭证。"""
+        self.auth_path.parent.mkdir(parents=True, exist_ok=True)
+        tmp_path = self.auth_path.with_name(f"{self.auth_path.name}.tmp")
+        tmp_path.write_text(store.model_dump_json(indent=2), encoding="utf-8")
+        with contextlib.suppress(Exception):
+            tmp_path.chmod(0o600)
+        tmp_path.replace(self.auth_path)
+
+    def set_api_key(
+        self,
+        provider: str,
+        key: str,
+        profile: str = "default",
+        base_url: str | None = None,
+        env: dict[str, str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """安全设置指定提供商的 API Key 凭据。"""
+        self.lock_path.parent.mkdir(parents=True, exist_ok=True)
+        with FileLock(str(self.lock_path), timeout=self.lock_timeout):
+            store = self.load_store()
+            prov_dict = store.providers.setdefault(provider, {})
+            prov_dict[profile] = ApiKeyCredential(
+                key=key,
+                base_url=base_url,
+                env=env or {},
+            )
+            store.active_profiles.setdefault(provider, profile)
+            self.save_store(store)
+
+    def set_oauth(
+        self,
+        provider: str,
+        access: str,
+        refresh: str,
+        expires: int,
+        profile: str = "default",
+        project_id: str = "aicode-consumers",
+        email: str | None = None,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """安全设置指定提供商的 OAuth 凭据。"""
+        self.lock_path.parent.mkdir(parents=True, exist_ok=True)
+        with FileLock(str(self.lock_path), timeout=self.lock_timeout):
+            store = self.load_store()
+            prov_dict = store.providers.setdefault(provider, {})
+            prov_dict[profile] = OAuthCredential(
+                access=access,
+                refresh=refresh,
+                expires=expires,
+                project_id=project_id,
+                email=email,
+                client_id=client_id,
+                client_secret=client_secret,
+            )
+            store.active_profiles.setdefault(provider, profile)
+            self.save_store(store)
+
+    def get_credential(self, provider: str, profile: str | None = None) -> ApiKeyCredential | OAuthCredential | None:
+        """获取指定提供商在指定 Profile 下的凭证（缺省为 active profile 或 default）。"""
+        store = self.load_store()
+        profile_name = profile or store.active_profiles.get(provider, "default")
+        return store.providers.get(provider, {}).get(profile_name)
+
+    def remove_credential(self, provider: str, profile: str | None = None) -> bool:
+        """安全删除指定提供商的凭据。
+
+        若未指定 profile，则删除该 provider 下的所有凭据及 active_profiles 映射；
+        若指定了 profile，则删除对应 Profile 的凭据；若删除后该 provider 无其他 Profile，
+        则同时清理其映射。
+        返回 True 表示成功删除，False 表示目标凭据不存在。
+        """
+        self.lock_path.parent.mkdir(parents=True, exist_ok=True)
+        with FileLock(str(self.lock_path), timeout=self.lock_timeout):
+            store = self.load_store()
+            removed = False
+            if profile is None:
+                if provider in store.providers:
+                    del store.providers[provider]
+                    removed = True
+                if provider in store.active_profiles:
+                    del store.active_profiles[provider]
+                    removed = True
+            else:
+                if provider in store.providers and profile in store.providers[provider]:
+                    del store.providers[provider][profile]
+                    removed = True
+                    if not store.providers[provider]:
+                        del store.providers[provider]
+                if store.active_profiles.get(provider) == profile:
+                    if provider in store.providers and store.providers[provider]:
+                        store.active_profiles[provider] = next(iter(store.providers[provider]))
+                    else:
+                        store.active_profiles.pop(provider, None)
+            if removed:
+                self.save_store(store)
+            return removed
+
+    async def get_valid_token(self, provider: str, profile: str | None = None) -> str:
+        """获取有效的调用 Token，自动完成 API Key 变量解析或 OAuth 静默刷新。"""
+        self.lock_path.parent.mkdir(parents=True, exist_ok=True)
+        with FileLock(str(self.lock_path), timeout=max(self.lock_timeout, 10.0)):
+            store = self.load_store()
+            profile_name = profile or store.active_profiles.get(provider, "default")
+            prov_dict = store.providers.get(provider, {})
+            cred = prov_dict.get(profile_name)
+
+            if not cred:
+                raise RuntimeError(f"未找到提供商 '{provider}' (Profile: {profile_name}) 的有效凭据，请先配置或登录。")
+
+            if isinstance(cred, ApiKeyCredential):
+                return cred.resolve_key()
+
+            if isinstance(cred, OAuthCredential):
+                if not cred.is_expired():
+                    return cred.access
+
+                # 执行 OAuth 静默刷新
+                client_id = cred.client_id or DEFAULT_ANTIGRAVITY_CLIENT_ID
+                client_secret = cred.client_secret or DEFAULT_ANTIGRAVITY_CLIENT_SECRET
+                try:
+                    async with httpx.AsyncClient(timeout=15.0) as client:
+                        resp = await client.post(
+                            GOOGLE_TOKEN_URL,
+                            data={
+                                "client_id": client_id,
+                                "client_secret": client_secret,
+                                "refresh_token": cred.refresh,
+                                "grant_type": "refresh_token",
+                            },
+                        )
+                        if resp.status_code != 200:
+                            raise RuntimeError(f"OAuth Token 自动刷新失败 ({resp.status_code}): {resp.text}")
+                        data = resp.json()
+                except httpx.RequestError as exc:
+                    raise RuntimeError(f"OAuth Token 自动刷新网络连接失败: {exc}") from exc
+
+                new_access = data["access_token"]
+                expires_in = data.get("expires_in", 3600)
+                try:
+                    exp_seconds = int(expires_in)
+                    now_ms = int(time.time() * 1000)
+                except (ValueError, TypeError):
+                    exp_seconds = 3600
+                    now_ms = 0
+                update_kwargs: dict[str, Any] = {
+                    "access": new_access,
+                    "expires": now_ms + (exp_seconds - 300) * 1000,
+                }
+                if "refresh_token" in data and data["refresh_token"]:
+                    update_kwargs["refresh"] = data["refresh_token"]
+
+                updated_cred = cred.model_copy(update=update_kwargs)
+                prov_dict[profile_name] = updated_cred
+                self.save_store(store)
+                return new_access
+
+            raise RuntimeError(f"未知的凭据类型: {type(cred)}")

package/src/my_agent_llm/auth/quota.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import httpx
+
+from .antigravity import (
+    ANTIGRAVITY_USER_AGENT,
+    DEFAULT_ANTIGRAVITY_ENDPOINT,
+    AntigravityCredentials,
+)
+
+
+@dataclass
+class QuotaBucket:
+    """单个模型或操作类型的配额桶。"""
+
+    bucket_id: str
+    display_name: str
+    remaining_fraction: float
+    reset_time: str | None = None
+
+    @property
+    def remaining_percent(self) -> int:
+        """配额剩余百分比（0-100 整数）。"""
+        return int(round(max(0.0, min(1.0, self.remaining_fraction)) * 100))
+
+
+def retrieve_user_quota_summary(creds: AntigravityCredentials) -> list[QuotaBucket]:
+    """向 Google Cloud Code 网关查询当前用户各模型配额余量。"""
+    url = f"{DEFAULT_ANTIGRAVITY_ENDPOINT}/v1internal:retrieveUserQuotaSummary"
+    headers = {
+        "Authorization": f"Bearer {creds.access_token}",
+        "x-goog-user-project": creds.project_id,
+        "User-Agent": ANTIGRAVITY_USER_AGENT,
+        "Content-Type": "application/json",
+    }
+    try:
+        res = httpx.post(url, headers=headers, json={}, timeout=10.0)
+        if res.status_code != 200:
+            return []
+        data = res.json()
+        buckets: list[QuotaBucket] = []
+        for group in data.get("groups", []):
+            for b in group.get("buckets", []):
+                buckets.append(
+                    QuotaBucket(
+                        bucket_id=b.get("bucketId", ""),
+                        display_name=b.get("displayName", ""),
+                        remaining_fraction=float(b.get("remainingFraction", 0.0)),
+                        reset_time=b.get("resetTime"),
+                    )
+                )
+        return buckets
+    except Exception:
+        return []

package/src/my_agent_llm/auth/schema.py

@@ -0,0 +1,94 @@
+"""AuthStore 强类型凭证实体模型与模式定义。"""
+
+from __future__ import annotations
+
+from enum import Enum
+import os
+import shlex
+import subprocess
+import time
+from typing import Annotated, Literal, Union
+from pydantic import BaseModel, Field
+
+__all__ = [
+    "CredentialType",
+    "ApiKeyCredential",
+    "OAuthCredential",
+    "Credential",
+    "AuthStore",
+]
+
+
+class CredentialType(str, Enum):
+    """凭证类型枚举。"""
+
+    API_KEY = "api_key"
+    OAUTH = "oauth"
+
+
+class ApiKeyCredential(BaseModel):
+    """API Key 凭据模型。"""
+
+    type: Literal[CredentialType.API_KEY, "api_key"] = CredentialType.API_KEY
+    key: str
+    base_url: str | None = None
+    env: dict[str, str] = Field(default_factory=dict)
+
+    def resolve_key(self) -> str:
+        """解析 API Key：处理字面量、!command 子进程输出与 $ENV_VAR/${ENV_VAR} 语法。"""
+        raw = self.key.strip()
+        if raw.startswith("!"):
+            cmd = raw[1:].strip()
+            if os.name == "nt":
+                parts = [
+                    p.strip('"') if (p.startswith('"') and p.endswith('"')) else p
+                    for p in shlex.split(cmd, posix=False)
+                ]
+            else:
+                parts = shlex.split(cmd)
+            if not parts:
+                return ""
+            res = subprocess.run(parts, capture_output=True, text=True, check=True)
+            return res.stdout.strip()
+        if raw.startswith("$"):
+            var_name = raw[1:].strip("{}")
+            return os.environ.get(var_name, "")
+        return raw
+
+
+class OAuthCredential(BaseModel):
+    """OAuth 凭据模型（支持 Google Cloud Code / Antigravity）。"""
+
+    type: Literal[CredentialType.OAUTH, "oauth"] = CredentialType.OAUTH
+    access: str
+    refresh: str
+    expires: int = 0  # 毫秒时间戳，<=0 表示永久或不计算过期
+    project_id: str = "aicode-consumers"
+    email: str | None = None
+    client_id: str | None = None
+    client_secret: str | None = None
+
+    def is_expired(self, buffer_ms: int = 300_000) -> bool:
+        """判定凭证是否已过期。
+
+        默认包含 5 分钟 (300,000ms) 缓冲期，保证在真正过期前完成自动刷新。
+        当 expires <= 0 时视为不计算过期。
+        """
+        if self.expires <= 0:
+            return False
+        return self.expires <= int(time.time() * 1000) + buffer_ms
+
+
+# 使用 Pydantic 区分联合体，基于 `type` 字段精准反序列化
+Credential = Annotated[
+    Union[ApiKeyCredential, OAuthCredential],
+    Field(discriminator="type"),
+]
+
+
+class AuthStore(BaseModel):
+    """全局统一凭证存储实体（映射 ~/.my-pi-agent/auth.json）。"""
+
+    version: int = 1
+    active_profiles: dict[str, str] = Field(default_factory=dict)
+    providers: dict[str, dict[str, Credential]] = Field(default_factory=dict)

package/src/my_agent_llm/client.py

@@ -0,0 +1,116 @@
+# pyright: reportUnreachable=false
+"""LLM 门面：按 provider 路由到对应实现，对外一套 API，只透传不碰 SDK。"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from typing import Any
+
+from .config import Config  # pyright: ignore[reportMissingImports]
+from .events import StreamEvent  # pyright: ignore[reportMissingImports]
+from .models import (  # pyright: ignore[reportMissingImports]
+    Message,
+    Response,
+    StreamChunk,
+)
+from .providers import Provider  # pyright: ignore[reportMissingImports]
+from .providers.registry import (
+    PROVIDER_REGISTRY,  # pyright: ignore[reportMissingImports]
+)
+
+
+class LLM:
+    """统一 LLM 客户端门面：严格接收强类型 Config 配置，多态路由至具体 Provider。"""
+
+    def __init__(self, config: Config) -> None:
+        """标准工程构造：严格接收 Config 实例。"""
+        if not isinstance(config, Config):  # pyright: ignore[reportUnreachable]
+            raise TypeError(f"LLM expects a Config instance, got {type(config).__name__}")  # pyright: ignore[reportUnreachable]
+        if config.provider not in PROVIDER_REGISTRY:
+            raise ValueError(f"Unknown provider '{config.provider}'. Available: {', '.join(sorted(PROVIDER_REGISTRY))}")
+        if not config.api_key and config.provider != "antigravity":
+            raise ValueError(f"No API key for provider: {config.provider}")
+        provider_cls = PROVIDER_REGISTRY[config.provider]
+        self._provider: Provider = provider_cls(config)
+        self.config = config
+
+    @property
+    def model(self) -> str:
+        """当前配置的模型名。"""
+        if not self.config.model:
+            raise ValueError("No model specified. Pass model=... or set Config.model.")
+        return self.config.model
+
+    def _resolve_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """统一合并调用参数与全局 Config 中的默认采样配置。"""
+        opts = dict(kwargs)
+        if self.config.temperature is not None:
+            opts.setdefault("temperature", self.config.temperature)
+        if self.config.max_tokens is not None:
+            opts.setdefault("max_tokens", self.config.max_tokens)
+        return opts
+
+    def chat(
+        self,
+        messages: list[Message],
+        *,
+        tools: list[dict] | None = None,
+        model: str | None = None,
+        **kwargs: Any,
+    ) -> Response:
+        """同步对话：完整历史 + 可选工具。"""
+        opts = self._resolve_kwargs(kwargs)
+        return self._provider.chat(messages, model=model or self.model, tools=tools, **opts)
+
+    def stream(
+        self,
+        messages: list[Message],
+        *,
+        tools: list[dict] | None = None,
+        model: str | None = None,
+        **kwargs: Any,
+    ) -> Iterator[StreamChunk]:
+        """同步流式。"""
+        opts = self._resolve_kwargs(kwargs)
+        return self._provider.stream(messages, model=model or self.model, tools=tools, **opts)
+
+    async def achat(
+        self,
+        messages: list[Message],
+        *,
+        tools: list[dict] | None = None,
+        model: str | None = None,
+        **kwargs: Any,
+    ) -> Response:
+        """异步对话。"""
+        opts = self._resolve_kwargs(kwargs)
+        return await self._provider.achat(messages, model=model or self.model, tools=tools, **opts)
+
+    async def achat_stream(
+        self,
+        messages: list[Message],
+        *,
+        tools: list[dict] | None = None,
+        model: str | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[StreamChunk]:
+        """异步流式。调用方直接 `async for chunk in llm.achat_stream(...)` 迭代，不 await。"""
+        opts = self._resolve_kwargs(kwargs)
+        async for chunk in self._provider.achat_stream(messages, model=model or self.model, tools=tools, **opts):
+            yield chunk
+
+    async def astream_events(
+        self,
+        messages: list[Message],
+        *,
+        tools: list[dict] | None = None,
+        model: str | None = None,
+        signal: Any | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[StreamEvent]:
+        """异步高阶流式事件流：直接产出 StreamStartEvent/TextDeltaEvent/StreamDoneEvent/StreamErrorEvent。"""
+        opts = self._resolve_kwargs(kwargs)
+        async for ev in self._provider.astream_events(
+            messages, model=model or self.model, tools=tools, signal=signal, **opts
+        ):
+            yield ev

package/src/my_agent_llm/config.py

@@ -0,0 +1,17 @@
+"""LLM 客户端配置：集中校验、可复用、不可变。"""
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Config(BaseModel):
+    """统一配置：provider/model/api_key/采样参数/网络参数。"""
+
+    provider: str = "openai"
+    model: str | None = None
+    api_key: str | None = None
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    max_tokens: int | None = Field(default=None, ge=1)
+    timeout: int = Field(default=30, ge=1)
+    max_retries: int = Field(default=3, ge=0)
+    base_url: str | None = None
+
+    model_config = ConfigDict(frozen=True)

package/src/my_agent_llm/events.py

@@ -0,0 +1,84 @@
+"""模型边界层高阶流式事件模型（对标 Tau provider_events）。
+
+提供供应商中立的高阶流式生命周期事件，携带实时累积的 partial: Message 快照。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from my_agent_llm.models import Message, ToolCall
+
+__all__ = [
+    "StreamDoneEvent",
+    "StreamErrorEvent",
+    "StreamEvent",
+    "StreamStartEvent",
+    "TextDeltaEvent",
+    "ThinkingDeltaEvent",
+    "ToolCallDeltaEvent",
+    "ToolCallDoneEvent",
+]
+
+
+@dataclass(frozen=True)
+class StreamEvent:
+    """高阶流式事件基类。"""
+
+
+@dataclass(frozen=True)
+class StreamStartEvent(StreamEvent):
+    """首个 Token 生成时发射，携带初始 partial Message 快照。"""
+
+    partial: Message
+
+
+@dataclass(frozen=True)
+class TextDeltaEvent(StreamEvent):
+    """正文文本增量事件。"""
+
+    delta: str
+    partial: Message
+
+
+@dataclass(frozen=True)
+class ThinkingDeltaEvent(StreamEvent):
+    """思维链推理增量事件（对标 DeepSeek-R1 / Claude 3.7 Thinking）。"""
+
+    delta: str
+    partial: Message
+
+
+@dataclass(frozen=True)
+class ToolCallDeltaEvent(StreamEvent):
+    """工具参数增量事件。"""
+
+    index: int
+    delta: str
+    partial: Message
+
+
+@dataclass(frozen=True)
+class ToolCallDoneEvent(StreamEvent):
+    """单个工具调用参数组装完毕事件。"""
+
+    index: int
+    tool_call: ToolCall
+    partial: Message
+
+
+@dataclass(frozen=True)
+class StreamDoneEvent(StreamEvent):
+    """流式正常结束事件，携带完整完型的 Message 实体与 Usage。"""
+
+    message: Message
+    usage: dict[str, int] | None = None
+
+
+@dataclass(frozen=True)
+class StreamErrorEvent(StreamEvent):
+    """流式异常或中途取消事件（Never-Throw 保证）。"""
+
+    error: Message
+    stop_reason: str = "error"  # "error" | "cancelled"
+    exc: Exception | None = None