my-pi-agent 0.1.0

package/src/my_agent_core/tools/core.py

@@ -0,0 +1,239 @@
+"""工具声明与分发 —— 核心模型与装饰器（Pydantic 驱动）。"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, get_type_hints, overload
+
+from pydantic import BaseModel, ConfigDict, ValidationError, create_model
+
+
+@dataclass
+class ToolResult:
+    """工具执行结果：成功/失败 + 数据或错误消息 + 结构化元数据。"""
+
+    ok: bool
+    data: Any = None
+    error: str | None = None
+    meta: dict[str, Any] = field(default_factory=dict)
+    terminate: bool = False
+
+    def serialize(self) -> str:
+        """转成写入 messages 的字符串。失败时返回错误文本。"""
+        if self.ok:
+            return str(self.data)
+        return self.error or "Unknown error"
+
+
+_FRAMEWORK_RESERVED_PARAMS: frozenset[str] = frozenset(
+    {"on_update", "signal", "tool_call_id"}
+)
+
+
+class Tool:
+    """一个可被模型调用的工具：函数本体 + 参数模型/原始Schema + 协议转换 + 超时与并发配置。"""
+
+    def __init__(
+        self,
+        func: Callable[..., Any],
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        params_model: type[BaseModel] | None = None,
+        raw_schema: dict[str, Any] | None = None,
+        timeout: float | None = None,
+        is_parallel_safe: bool = False,
+    ):
+        self.func = func
+        self.name = name or func.__name__
+        self.description = description or (inspect.getdoc(func) or "")
+        self.raw_schema = raw_schema
+        self.timeout = timeout
+        self.is_parallel_safe = is_parallel_safe
+        self.is_async = inspect.iscoroutinefunction(func)
+        sig = inspect.signature(func)
+        self._accepts_on_update = "on_update" in sig.parameters
+        self._accepts_signal = "signal" in sig.parameters
+        self._accepts_tool_call_id = "tool_call_id" in sig.parameters
+        self.params_model = params_model or (
+            None if raw_schema else self._create_params_model(func)
+        )
+
+    def _create_params_model(self, func: Callable[..., Any]) -> type[BaseModel]:
+        """从函数签名动态建模（pydantic create_model）。"""
+        hints = get_type_hints(func)
+        fields: dict[str, Any] = {}
+        for param_name, param in inspect.signature(func).parameters.items():
+            if param_name in _FRAMEWORK_RESERVED_PARAMS:
+                continue
+            if param.kind in (
+                inspect.Parameter.VAR_POSITIONAL,
+                inspect.Parameter.VAR_KEYWORD,
+            ):
+                raise TypeError(
+                    f"tool '{func.__name__}': parameter '{param_name}' is "
+                    "*args/**kwargs, which is not supported（不支持）"
+                )
+            if param_name not in hints:
+                raise TypeError(
+                    f"tool '{func.__name__}': parameter '{param_name}' "
+                    "has no type annotation（没有类型标注）"
+                )
+            default = ... if param.default is inspect.Parameter.empty else param.default
+            fields[param_name] = (hints[param_name], default)
+
+        try:
+            model = create_model(
+                f"{func.__name__}_Args",
+                __config__=ConfigDict(extra="forbid"),  # 多余参数 → 校验错误
+                **fields,
+            )
+            return model
+        except Exception as exc:  # 无法建模的类型 → 装饰时明确失败
+            raise TypeError(
+                f"tool '{func.__name__}': cannot build parameter schema: {exc}"
+            ) from exc
+
+    def to_openai_schema(self) -> dict[str, Any]:
+        """生成 OpenAI tools 参数。"""
+        if self.raw_schema is not None:
+            params = self.raw_schema
+        elif self.params_model is not None:
+            params = self.params_model.model_json_schema()
+        else:
+            params = {}
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": params,
+            },
+        }
+
+    async def execute(
+        self,
+        args: dict[str, Any],
+        signal: Any | None = None,
+        on_update: Callable[[Any], None] | None = None,
+        tool_call_id: str | None = None,
+    ) -> ToolResult:
+        """校验 + 执行，永不抛（错误全部转 ToolResult）。自动适配 sync/async 函数。"""
+        accepting_updates = True
+
+        def guarded_on_update(partial: Any) -> None:
+            if accepting_updates and on_update is not None:
+                on_update(partial)
+
+        extra_kwargs: dict[str, Any] = {}
+        if self._accepts_on_update:
+            extra_kwargs["on_update"] = guarded_on_update
+        if self._accepts_signal:
+            extra_kwargs["signal"] = signal
+        if self._accepts_tool_call_id and tool_call_id is not None:
+            extra_kwargs["tool_call_id"] = tool_call_id
+
+        if self.params_model is not None:
+            try:
+                validated = self.params_model.model_validate(args)
+            except ValidationError as exc:
+                return ToolResult(ok=False, error=str(exc))
+            kwargs = {**validated.model_dump(), **extra_kwargs}
+
+            def func_call() -> Any:
+                return self.func(**kwargs)
+
+            async def async_func_call() -> Any:
+                return await self.func(**kwargs)
+        else:
+
+            def func_call() -> Any:
+                if extra_kwargs:
+                    return self.func(args, **extra_kwargs)
+                return self.func(args)
+
+            async def async_func_call() -> Any:
+                if extra_kwargs:
+                    return await self.func(args, **extra_kwargs)
+                return await self.func(args)
+
+        try:
+            if self.is_async:
+                result = await async_func_call()
+            else:
+                result = await asyncio.to_thread(func_call)
+        except Exception as exc:  # 工具错误 → 消息，喂回模型
+            return ToolResult(
+                ok=False, error=f"Error executing tool '{self.name}': {exc}"
+            )
+        finally:
+            accepting_updates = False
+
+        if isinstance(result, ToolResult):
+            return result
+        return ToolResult(ok=True, data=result)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """直接调用工具函数本体。"""
+        return self.func(*args, **kwargs)
+
+
+@overload
+def tool(
+    func: Callable[..., Any],
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    params_model: type[BaseModel] | None = None,
+    raw_schema: dict[str, Any] | None = None,
+    timeout: float | None = None,
+    is_parallel_safe: bool = False,
+) -> Tool: ...
+
+
+@overload
+def tool(
+    func: None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    params_model: type[BaseModel] | None = None,
+    raw_schema: dict[str, Any] | None = None,
+    timeout: float | None = None,
+    is_parallel_safe: bool = False,
+) -> Callable[[Callable[..., Any]], Tool]: ...
+
+
+def tool(
+    func: Callable[..., Any] | None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    params_model: type[BaseModel] | None = None,
+    raw_schema: dict[str, Any] | None = None,
+    timeout: float | None = None,
+    is_parallel_safe: bool = False,
+) -> Tool | Callable[[Callable[..., Any]], Tool]:
+    """@tool 装饰器：支持 @tool 与 @tool(name=..., description=..., params_model=...)。
+
+    schema 生成由 pydantic 驱动：参数类型支持 pydantic 全集（list/dict/Optional/嵌套等），
+    允许默认值；无标注参数与 *args/**kwargs 在装饰时抛 TypeError（明确失败）。
+    """
+
+    def decorator(f: Callable[..., Any]) -> Tool:
+        return Tool(
+            func=f,
+            name=name,
+            description=description,
+            params_model=params_model,
+            raw_schema=raw_schema,
+            timeout=timeout,
+            is_parallel_safe=is_parallel_safe,
+        )
+
+    if func is None:
+        return decorator
+    return decorator(func)

package/src/my_agent_llm/__init__.py

@@ -0,0 +1,45 @@
+"""统一 LLM 客户端包（模型边界层）。"""
+
+from .client import LLM
+from .config import Config
+from .events import (  # pyright: ignore[reportMissingImports]
+    StreamDoneEvent,
+    StreamErrorEvent,
+    StreamEvent,
+    StreamStartEvent,
+    TextDeltaEvent,
+    ThinkingDeltaEvent,
+    ToolCallDeltaEvent,
+    ToolCallDoneEvent,
+)
+from .models import (
+    Message,
+    Response,
+    StreamChunk,
+    ToolCall,
+    TurnOutcome,
+    normalize_finish_reason,
+)
+from .providers import Provider
+from .stream import StreamAccumulator  # pyright: ignore[reportMissingImports]
+
+__all__ = [
+    "LLM",
+    "Config",
+    "Message",
+    "Response",
+    "StreamAccumulator",
+    "StreamChunk",
+    "StreamDoneEvent",
+    "StreamErrorEvent",
+    "StreamEvent",
+    "StreamStartEvent",
+    "TextDeltaEvent",
+    "ThinkingDeltaEvent",
+    "ToolCallDeltaEvent",
+    "ToolCallDoneEvent",
+    "ToolCall",
+    "TurnOutcome",
+    "normalize_finish_reason",
+    "Provider",
+]

package/src/my_agent_llm/auth/__init__.py

@@ -0,0 +1,46 @@
+"""凭据管理与鉴权模块。"""
+
+from .antigravity import (  # pyright: ignore[reportMissingImports]
+    ANTIGRAVITY_USER_AGENT,
+    DEFAULT_ANTIGRAVITY_ENDPOINT,
+    GOOGLE_OAUTH_TOKEN_URL,
+    AntigravityAuthResolver,
+    AntigravityCredentials,
+)
+from .manager import (  # pyright: ignore[reportMissingImports]
+    DEFAULT_ANTIGRAVITY_CLIENT_ID,
+    DEFAULT_ANTIGRAVITY_CLIENT_SECRET,
+    GOOGLE_TOKEN_URL,
+    AuthManager,
+)
+from .quota import (  # pyright: ignore[reportMissingImports]
+    QuotaBucket,
+    retrieve_user_quota_summary,
+)
+from .schema import (  # pyright: ignore[reportMissingImports]
+    ApiKeyCredential,
+    AuthStore,
+    Credential,
+    CredentialType,
+    OAuthCredential,
+)
+
+__all__ = [
+    "ANTIGRAVITY_USER_AGENT",
+    "DEFAULT_ANTIGRAVITY_ENDPOINT",
+    "GOOGLE_OAUTH_TOKEN_URL",
+    "AntigravityAuthResolver",
+    "AntigravityCredentials",
+    "QuotaBucket",
+    "retrieve_user_quota_summary",
+    # AuthManager & 凭据模型
+    "DEFAULT_ANTIGRAVITY_CLIENT_ID",
+    "DEFAULT_ANTIGRAVITY_CLIENT_SECRET",
+    "GOOGLE_TOKEN_URL",
+    "AuthManager",
+    "ApiKeyCredential",
+    "AuthStore",
+    "Credential",
+    "CredentialType",
+    "OAuthCredential",
+]

package/src/my_agent_llm/auth/antigravity.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import logging
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_ANTIGRAVITY_ENDPOINT = "https://cloudcode-pa.googleapis.com"
+ANTIGRAVITY_USER_AGENT = "antigravity/cli/1.1.23 (aidev_client; os_type=windows; arch=amd64; auth_method=consumer)"
+GOOGLE_OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token"  # noqa: S105
+
+
+@dataclass
+class AntigravityCredentials:
+    """Antigravity / Google Cloud Code 鉴权凭据实体。"""
+
+    access_token: str
+    refresh_token: str | None = None
+    expires_at: int = 0
+    project_id: str = "aicode-consumers"
+    email: str | None = None
+    auth_file_path: Path | None = None
+    client_id: str | None = None
+    client_secret: str | None = None
+
+
+class AntigravityAuthResolver:
+    """Antigravity / Google Cloud Code Assist 凭据解析与自动刷新器（对标 pi-antigravity）。"""
+
+    def __init__(
+        self,
+        pi_auth_path: Path | None = None,
+        credentials_path: Path | None = None,
+        workspace: Path | None = None,
+    ) -> None:
+        home = Path(os.environ.get("USERPROFILE") or os.environ.get("HOME") or "~").expanduser()
+        self.workspace_auth_path = (Path(workspace).resolve() / ".my_agent" / "auth.json") if workspace else None
+        if pi_auth_path is not None:
+            self.pi_auth_path = Path(pi_auth_path).resolve()
+        elif "MY_AGENT_HOME" not in os.environ:
+            default_pi_auth = home / ".pi" / "agent" / "auth.json"
+            self.pi_auth_path = default_pi_auth if default_pi_auth.exists() else None
+        else:
+            self.pi_auth_path = None
+
+        if credentials_path is not None:
+            self.credentials_path = Path(credentials_path).resolve()
+        else:
+            self.credentials_path = home / ".my_agent" / "credentials.json"
+
+    def resolve_credentials_raw(self) -> AntigravityCredentials | None:
+        """从工作区/项目凭据或环境变量动态获取凭据，零硬编码，隔离外部环境。"""
+        # 1. 环境变量优先
+        env_token = os.environ.get("ANTIGRAVITY_ACCESS_TOKEN") or os.environ.get("ANTIGRAVITY_API_KEY")
+        if env_token:
+            return AntigravityCredentials(
+                access_token=env_token,
+                project_id=os.environ.get("ANTIGRAVITY_PROJECT_ID", "aicode-consumers"),
+                client_id=os.environ.get("GOOGLE_OAUTH_CLIENT_ID"),
+                client_secret=os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET"),
+            )
+
+        # 2. 依次读取凭据文件
+        candidate_paths: list[Path] = []
+        if self.pi_auth_path:
+            candidate_paths.append(self.pi_auth_path)
+        if self.workspace_auth_path:
+            candidate_paths.append(self.workspace_auth_path)
+        if self.credentials_path:
+            candidate_paths.append(self.credentials_path)
+
+        for target_path in candidate_paths:
+            if target_path and target_path.exists():
+                try:
+                    text_content = target_path.read_text(encoding="utf-8")
+                    data = json.loads(text_content)
+                    if not isinstance(data, dict):
+                        continue
+
+                    entry = data.get("antigravity") or data.get("google-antigravity")
+                    if entry and isinstance(entry, dict):
+                        access_token = entry.get("access") or entry.get("access_token")
+                        if access_token:
+                            expires_raw = entry.get("expires") or entry.get("expires_at", 0)
+                            try:
+                                exp_val = int(expires_raw)
+                            except (ValueError, TypeError):
+                                exp_val = 0
+                            return AntigravityCredentials(
+                                access_token=access_token,
+                                refresh_token=entry.get("refresh") or entry.get("refresh_token"),
+                                expires_at=exp_val,
+                                project_id=entry.get("projectId") or entry.get("project_id", "aicode-consumers"),
+                                email=entry.get("email"),
+                                auth_file_path=target_path,
+                                client_id=entry.get("client_id") or os.environ.get("GOOGLE_OAUTH_CLIENT_ID"),
+                                client_secret=entry.get("client_secret")
+                                or os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET"),
+                            )
+                except Exception as exc:
+                    logger.debug("Failed to read credentials from %s: %s", target_path, exc)
+                    continue
+
+        return None
+
+    def resolve_credentials(self) -> AntigravityCredentials | None:
+        """获取凭据别名，对齐规范。"""
+        return self.resolve_credentials_raw()
+
+    def is_expired(self, creds: AntigravityCredentials | None) -> bool:
+        """检查凭据是否已过期或即将过期（剩余不足 5 分钟）。"""
+        if not creds:
+            return True
+        if creds.expires_at <= 0:
+            return False
+        try:
+            return creds.expires_at <= int(time.time() * 1000) + 300000
+        except Exception:
+            return True
+
+    def refresh(self, creds: AntigravityCredentials) -> AntigravityCredentials:
+        """向 Google OAuth 端点发起刷新请求换取新 token 并持久化写回。"""
+        if not creds.refresh_token:
+            raise RuntimeError("Cannot refresh Antigravity token: missing refresh_token.")
+
+        client_id = creds.client_id or os.environ.get("GOOGLE_OAUTH_CLIENT_ID")
+        client_secret = creds.client_secret or os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET")
+        if not client_id or not client_secret:
+            raise RuntimeError("Cannot refresh Antigravity token: missing client_id/client_secret in auth.json or env.")
+
+        payload = {
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "refresh_token": creds.refresh_token,
+            "grant_type": "refresh_token",
+        }
+        res = httpx.post(GOOGLE_OAUTH_TOKEN_URL, data=payload, timeout=15.0)
+        if res.status_code != 200:
+            raise RuntimeError(f"Failed to refresh Antigravity token: {res.status_code} {res.text}")
+
+        data = res.json()
+        new_access = data["access_token"]
+        expires_in = data.get("expires_in", 3600)
+        try:
+            new_expires_at = int(time.time() * 1000) + (expires_in * 1000) - (300 * 1000)
+        except Exception:
+            new_expires_at = 0
+
+        updated_creds = AntigravityCredentials(
+            access_token=new_access,
+            refresh_token=creds.refresh_token,
+            expires_at=new_expires_at,
+            project_id=creds.project_id,
+            email=creds.email,
+            auth_file_path=creds.auth_file_path,
+            client_id=client_id,
+            client_secret=client_secret,
+        )
+
+        # 写回持久化文件
+        if creds.auth_file_path and creds.auth_file_path.exists():
+            try:
+                raw_text = creds.auth_file_path.read_text(encoding="utf-8")
+                raw_data = json.loads(raw_text)
+                target_key = (
+                    "antigravity"
+                    if "antigravity" in raw_data
+                    else ("google-antigravity" if "google-antigravity" in raw_data else "antigravity")
+                )
+                if target_key not in raw_data:
+                    raw_data[target_key] = {}
+                if "access" in raw_data[target_key] or "access_token" not in raw_data[target_key]:
+                    raw_data[target_key]["access"] = new_access
+                if "access_token" in raw_data[target_key]:
+                    raw_data[target_key]["access_token"] = new_access
+                raw_data[target_key]["expires"] = new_expires_at
+                # 原子写入
+                tmp_file = creds.auth_file_path.with_suffix(".tmp")
+                tmp_file.write_text(json.dumps(raw_data, indent=2), encoding="utf-8")
+                tmp_file.replace(creds.auth_file_path)
+            except Exception as exc:
+                logger.debug("Failed to write back refreshed credentials: %s", exc)
+
+        return updated_creds
+
+    def refresh_token(self, creds: AntigravityCredentials) -> AntigravityCredentials:
+        """refresh 的别名方法。"""
+        return self.refresh(creds)
+
+    def get_valid_credentials(self) -> AntigravityCredentials:
+        """获取有效凭据。优先直接使用 auth.json 凭据；若过期且配置了密钥则自动刷新。"""
+        creds = self.resolve_credentials_raw()
+        if not creds:
+            raise RuntimeError(
+                f"No Antigravity credentials found. Please ensure {self.pi_auth_path} exists or set ANTIGRAVITY_ACCESS_TOKEN."
+            )
+        client_id = creds.client_id or os.environ.get("GOOGLE_OAUTH_CLIENT_ID")
+        client_secret = creds.client_secret or os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET")
+        if self.is_expired(creds) and creds.refresh_token and client_id and client_secret:
+            with contextlib.suppress(Exception):
+                creds = self.refresh(creds)
+        return creds