cheso-cli 0.2.0__py3-none-any.whl

cheso_cli/client.py

@@ -0,0 +1,534 @@
+# pylint: disable=broad-except
+"""精简的 Cheso Manager HTTP 客户端，使用 X-Cheso-Token 认证。"""
+
+from __future__ import annotations
+
+import contextlib
+import queue
+import threading
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import Any, Protocol
+from urllib.parse import quote
+
+import httpx
+
+from cheso_cli.config import USER_AGENT
+
+# 查询忙碌项目中运行任务时的默认分页大小；列表按 updated_at 降序排列。
+PLAN_LIST_PAGE_SIZE = 20
+# 发起任务需要分配运行环境，比其他接口调用更慢。
+CREATE_PLAN_TIMEOUT = 60.0
+
+
+class ApiError(Exception):
+    """Manager API 返回的业务错误或传输错误。"""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        code: int | None = None,
+        http_status: int | None = None,
+        span_id: str | None = None,
+        body: Any = None,
+        retry_after: int | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.code = code
+        self.http_status = http_status
+        self.span_id = span_id
+        self.body = body
+        self.retry_after = retry_after
+
+    def __str__(self) -> str:
+        parts = [self.message]
+        if self.code is not None:
+            parts.append(f"code={self.code}")
+        if self.http_status is not None:
+            parts.append(f"http={self.http_status}")
+        if self.span_id:
+            parts.append(f"span_id={self.span_id}")
+        return " ".join(parts)
+
+    @property
+    def details(self) -> dict[str, Any]:
+        """错误响应中的结构化详情，即响应封装里的 ``data`` 字段；没有时返回空字典。"""
+        if isinstance(self.body, dict) and isinstance(self.body.get("data"), dict):
+            return self.body["data"]
+        return {}
+
+
+class TimeLimitReached(httpx.TimeoutException):
+    """命令的 ``--timeout`` 已到，不再发送下一个请求。"""
+
+    def __init__(self) -> None:
+        super().__init__("--timeout reached")
+
+
+@dataclass(frozen=True)
+class Deadline:
+    """命令 ``--timeout`` 对应的截止时刻，使用 ``clock`` 计时。"""
+
+    at: float
+    clock: Callable[[], float]
+
+    def left(self) -> float:
+        """返回剩余秒数；没有剩余时间时抛出 TimeLimitReached。"""
+        left = self.at - self.clock()
+        if left <= 0:
+            raise TimeLimitReached()
+        return left
+
+    def expired(self) -> bool:
+        return self.clock() >= self.at
+
+    def cap(self, timeout: httpx.Timeout) -> httpx.Timeout:
+        """将 ``timeout`` 的每个阶段都限制在剩余时间内。"""
+        left = self.left()
+        return httpx.Timeout(
+            connect=min(timeout.connect or left, left),
+            read=min(timeout.read or left, left),
+            write=min(timeout.write or left, left),
+            pool=min(timeout.pool or left, left),
+        )
+
+
+_END = object()
+
+
+def iterate_until(deadline: Deadline, items: Callable[[], Iterator[Any]]) -> Iterator[Any]:
+    """依次返回 ``items()`` 产生的数据，每次最多等到 ``deadline``。
+
+    httpx 的超时限制每次套接字读取，而非整个响应；服务端持续少量发送数据时，请求仍可能超过截止时间。
+    因此，将 ``items()`` 中阻塞的网络 I/O 放到守护线程执行，当前线程只等到截止时间。
+    超时后抛出 TimeLimitReached，并通知工作线程停止；工作线程自行关闭响应，后续产出的数据直接丢弃。
+    """
+    deadline.left()  # 超过截止时间后不再启动操作。
+    handoff: queue.Queue[tuple[bool, Any]] = queue.Queue(maxsize=4)
+    stop = threading.Event()
+
+    def hand_over(entry: tuple[bool, Any]) -> bool:
+        while not stop.is_set():
+            try:
+                handoff.put(entry, timeout=0.1)
+                return True
+            except queue.Full:
+                continue
+        return False
+
+    def produce() -> None:
+        source: Iterator[Any] | None = None
+        try:
+            source = items()
+            for item in source:
+                if not hand_over((True, item)):
+                    return
+            hand_over((True, _END))
+        except Exception as e:  # pylint: disable=broad-except  # 异常会在等待线程中重新抛出
+            hand_over((False, e))
+        finally:
+            close = getattr(source, "close", None)
+            if close is not None:
+                with contextlib.suppress(Exception):
+                    close()
+
+    threading.Thread(target=produce, name="cheso-io", daemon=True).start()
+    try:
+        while True:
+            try:
+                ok, value = handoff.get(timeout=deadline.left())
+            except queue.Empty:
+                raise TimeLimitReached() from None
+            if not ok:
+                raise value
+            if value is _END:
+                return
+            yield value
+    finally:
+        stop.set()
+
+
+def call_until(deadline: Deadline, fn: Callable[[], Any]) -> Any:
+    """等待阻塞网络操作 ``fn()``，最多等到 ``deadline``；实现见 iterate_until。"""
+    with contextlib.closing(iterate_until(deadline, lambda: iter((fn(),)))) as results:
+        return next(results)
+
+
+class TokenSource(Protocol):
+    """提供 X-Cheso-Token；OAuth 登录收到 401 后，可提供刷新后的令牌。"""
+
+    def token(self) -> str:
+        """返回下一个请求要发送的令牌。"""
+        ...
+
+    def refresh_after_unauthorized(self, rejected: str) -> str | None:
+        """返回新令牌供调用方重试一次；无法续期时返回 None。"""
+        ...
+
+
+class StaticToken:
+    """固定令牌，来自 ``--token`` 或 ``CHESO_API_TOKEN``，无需刷新。"""
+
+    def __init__(self, value: str) -> None:
+        self._value = value
+
+    def token(self) -> str:
+        return self._value
+
+    def refresh_after_unauthorized(self, rejected: str) -> str | None:
+        return None
+
+
+def _parse_retry_after(value: str | None) -> int | None:
+    try:
+        seconds = int((value or "").strip())
+    except ValueError:
+        return None
+    return seconds if seconds >= 0 else None
+
+
+def _path_id(value: str) -> str:
+    return quote(value, safe="")
+
+
+class ChesoClient:
+    """封装 Manager 公开 API 请求和响应的精简客户端。"""
+
+    def __init__(
+        self,
+        base_url: str,
+        token: str | None = None,
+        *,
+        auth: TokenSource | None = None,
+        timeout: float = 30.0,
+        transport: httpx.BaseTransport | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        if auth is None:
+            if not token:
+                raise ValueError("token is required (CHESO_API_TOKEN / --token)")
+            auth = StaticToken(token)
+        self._auth = auth
+        self.base_url = base_url.rstrip("/")
+        default_headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+            "User-Agent": USER_AGENT,
+        }
+        default_headers.update(headers or {})
+        self._timeout = timeout
+        self._deadline: Deadline | None = None
+        # API 调用不跟随重定向，避免将令牌发送到其他地址。
+        self._client = httpx.Client(
+            base_url=self.base_url,
+            headers=default_headers,
+            timeout=timeout,
+            transport=transport,
+            follow_redirects=False,
+        )
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> ChesoClient:
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()
+
+    @contextlib.contextmanager
+    def time_limit(self, deadline: Deadline) -> Iterator[None]:
+        """作用域内不在 ``deadline`` 之后发起请求，等待请求也不超过该时刻，见 call_until。
+
+        令牌刷新不会被截断：丢失刷新响应后就无法知道轮换后的刷新令牌，下一次刷新会使登录失效。
+        """
+        previous, self._deadline = self._deadline, deadline
+        try:
+            yield
+        finally:
+            self._deadline = previous
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        *,
+        params: dict[str, Any] | None = None,
+        json_body: dict[str, Any] | None = None,
+        timeout: float | None = None,
+    ) -> Any:
+        if self._deadline is not None:
+            self._deadline.left()
+        token = self._auth.token()
+        resp = self._send(method, path, token, params=params, json_body=json_body, timeout=timeout)
+        if resp.status_code == 401:
+            fresh = self._auth.refresh_after_unauthorized(token)
+            if fresh and fresh != token:
+                resp = self._send(method, path, fresh, params=params, json_body=json_body, timeout=timeout)
+        return self._decode(resp)
+
+    def _send(
+        self,
+        method: str,
+        path: str,
+        token: str,
+        *,
+        params: dict[str, Any] | None,
+        json_body: dict[str, Any] | None,
+        timeout: float | None,
+    ) -> httpx.Response:
+        extra: dict[str, Any] = {} if timeout is None else {"timeout": timeout}
+        deadline = self._deadline
+        if deadline is not None:
+            # 同时限制各阶段超时，避免已放弃的请求在后台滞留过久。
+            extra["timeout"] = deadline.cap(httpx.Timeout(self._timeout if timeout is None else timeout))
+
+        def send() -> httpx.Response:
+            return self._client.request(
+                method, path, params=params, json=json_body, headers={"X-Cheso-Token": token}, **extra
+            )
+
+        return send() if deadline is None else call_until(deadline, send)
+
+    @staticmethod
+    def _decode(resp: httpx.Response) -> Any:
+        retry_after = _parse_retry_after(resp.headers.get("Retry-After"))
+        try:
+            payload = resp.json()
+        except ValueError as e:
+            span_hdr = resp.headers.get("X-Span-Id")
+            raise ApiError(
+                "request failed",
+                http_status=resp.status_code,
+                span_id=span_hdr if isinstance(span_hdr, str) and span_hdr else None,
+                body=resp.text[:500] if resp.text else None,
+                retry_after=retry_after,
+            ) from e
+
+        if not isinstance(payload, dict):
+            raise ApiError(
+                "unexpected response shape",
+                http_status=resp.status_code,
+                body=payload,
+            )
+
+        code = payload.get("code")
+        message = str(payload.get("message") or "")
+        span_id = payload.get("span_id")
+        if resp.status_code >= 400 or (isinstance(code, int) and code != 0):
+            raise ApiError(
+                message or "request failed",
+                code=code if isinstance(code, int) else None,
+                http_status=resp.status_code,
+                span_id=span_id if isinstance(span_id, str) else None,
+                body=payload,
+                retry_after=retry_after,
+            )
+        return payload.get("data")
+
+    # ---- 账号 ----
+
+    def profile(self) -> Any:
+        return self.request("GET", "/api/v1/profile")
+
+    def whoami(self) -> Any:
+        return self.request("GET", "/api/v1/whoami")
+
+    def credit_grants(self) -> Any:
+        """当前可用的积分授权：``{total, grants: [{amount_available, valid_from, expires_at, ...}]}``。"""
+        return self.request("GET", "/api/v1/billing/credits")
+
+    # ---- 项目 ----
+
+    def project_create(self, *, description: str | None = None) -> Any:
+        body: dict[str, Any] = {}
+        if description is not None:
+            body["description"] = description
+        return self.request("POST", "/api/v1/projects/create", json_body=body)
+
+    def project_list(
+        self,
+        *,
+        status: str | None = None,
+        starred: bool | None = None,
+        sort_by: str | None = None,
+        page: int | None = None,
+        page_size: int | None = None,
+    ) -> Any:
+        params: dict[str, Any] = {}
+        if status is not None:
+            params["status"] = status
+        if starred is not None:
+            params["starred"] = "true" if starred else "false"
+        if sort_by is not None:
+            params["sort_by"] = sort_by
+        if page is not None:
+            params["page"] = page
+        if page_size is not None:
+            params["page_size"] = page_size
+        return self.request("GET", "/api/v1/projects", params=params or None)
+
+    def project_get(self, project_id: str) -> Any:
+        return self.request("GET", f"/api/v1/projects/{project_id}")
+
+    def project_update(
+        self,
+        project_id: str,
+        *,
+        description: str | None = None,
+        starred: bool | None = None,
+        cover_file_id: str | None = None,
+        cover_file_version: int | None = None,
+    ) -> Any:
+        body: dict[str, Any] = {}
+        if description is not None:
+            body["description"] = description
+        if starred is not None:
+            body["starred"] = starred
+        if cover_file_id is not None:
+            body["cover_file_id"] = cover_file_id
+        if cover_file_version is not None:
+            body["cover_file_version"] = cover_file_version
+        return self.request("PUT", f"/api/v1/projects/{project_id}", json_body=body)
+
+    def project_delete(self, project_id: str) -> Any:
+        return self.request("DELETE", f"/api/v1/projects/{project_id}")
+
+    def project_fork(self, project_id: str, *, description: str | None = None) -> Any:
+        body: dict[str, Any] = {}
+        if description is not None:
+            body["description"] = description
+        return self.request(
+            "POST",
+            f"/api/v1/projects/{project_id}/fork",
+            json_body=body,
+        )
+
+    def project_list_all(
+        self,
+        *,
+        status: str | None = "active",
+        page_size: int = 100,
+    ) -> list[dict[str, Any]]:
+        """逐页读取项目列表，直到没有更多结果。"""
+        items: list[dict[str, Any]] = []
+        page = 1
+        while True:
+            data = self.project_list(status=status, page=page, page_size=page_size)
+            batch = list((data or {}).get("projects") or [])
+            items.extend(p for p in batch if isinstance(p, dict))
+            total = int((data or {}).get("total") or 0)
+            if len(items) >= total or not batch:
+                break
+            page += 1
+        return items
+
+    def session_create(self, *, project_id: str, title: str | None = None) -> Any:
+        body: dict[str, Any] = {"project_id": project_id}
+        if title is not None:
+            body["title"] = title
+        return self.request("POST", "/api/v1/sessions/create", json_body=body)
+
+    # ---- Plan（任务）----
+
+    def plan_create(
+        self,
+        *,
+        prompt: str,
+        project_id: str | None = None,
+        session_id: str | None = None,
+        contexts: list[dict[str, Any]] | None = None,
+        agent_type: str | None = None,
+    ) -> Any:
+        """发起 Plan。始终不发送 ``on_running``；项目忙时返回 2614，不排队。"""
+        body: dict[str, Any] = {"prompt": prompt}
+        if project_id is not None:
+            body["project_id"] = project_id
+        if session_id is not None:
+            body["session_id"] = session_id
+        if contexts:
+            body["contexts"] = contexts
+        if agent_type:
+            body["agent_type"] = agent_type
+        return self.request("POST", "/api/v1/plans/create", json_body=body, timeout=CREATE_PLAN_TIMEOUT)
+
+    def plan_get(self, plan_id: str) -> Any:
+        return self.request("GET", f"/api/v1/plans/{plan_id}")
+
+    def plan_list(self, *, project_id: str, page_size: int = PLAN_LIST_PAGE_SIZE) -> Any:
+        return self.request("GET", "/api/v1/plans", params={"project_id": project_id, "page_size": page_size})
+
+    def plan_result(self, plan_id: str) -> Any:
+        """返回字段：``{plan_id, status, result{text,text_truncated,files[]}, interaction, next_plan_id, reason}``。"""
+        return self.request("GET", f"/api/v1/plans/{_path_id(plan_id)}/result")
+
+    def plan_interact(self, plan_id: str, *, args: dict[str, Any], skip: bool) -> Any:
+        return self.request(
+            "POST",
+            f"/api/v1/plans/{_path_id(plan_id)}/interact",
+            json_body={"args": args, "skip": skip},
+        )
+
+    def plan_terminate(self, plan_id: str) -> Any:
+        return self.request("POST", f"/api/v1/plans/{_path_id(plan_id)}/terminate", json_body={})
+
+    def plan_credit_get(self, *, project_id: str, plan_id: str) -> Any:
+        return self.request(
+            "GET",
+            "/api/v1/billing/credit-transactions/plans",
+            params={"project_id": project_id, "plan_id": plan_id},
+        )
+
+    # ---- 文件 ----
+
+    def file_presign(self, body: dict[str, Any]) -> Any:
+        """返回 ``{file_id, version, upload_url, instant}``；``instant`` 表示内容已经存储，可秒传。"""
+        return self.request("POST", "/api/v1/files/presign", json_body=body)
+
+    def file_confirm(self, *, file_id: str, version: int) -> Any:
+        return self.request("POST", "/api/v1/files/confirm", json_body={"file_id": file_id, "version": version})
+
+    def file_download_url(self, file_id: str, *, version: int | None = None) -> Any:
+        """返回 ``{file_id, download_url, attachment_url}``；这些 URL 带预签名，6 小时后过期。"""
+        body: dict[str, Any] = {"file_id": file_id}
+        if version:
+            body["version"] = version
+        return self.request("POST", "/api/v1/files/download_url", json_body=body)
+
+    def file_get(self, file_id: str) -> Any:
+        return self.request("GET", f"/api/v1/files/{_path_id(file_id)}")
+
+    def project_file_tree(self, project_id: str) -> Any:
+        return self.request("GET", "/api/v1/project/files/tree", params={"project_id": project_id})
+
+    # ---- 转换与翻译 ----
+
+    def convert_supported(self) -> Any:
+        return self.request("GET", "/api/v1/converts/supported")
+
+    def convert_create(self, *, file_id: str, source_type: str, target_type: str) -> Any:
+        return self.request(
+            "POST",
+            "/api/v1/converts/create",
+            json_body={"file_id": file_id, "source_type": source_type, "target_type": target_type},
+        )
+
+    def convert_get(self, job_id: str) -> Any:
+        return self.request("GET", f"/api/v1/converts/{_path_id(job_id)}")
+
+    def translation_create(
+        self,
+        *,
+        file_id: str,
+        target_language: str,
+        source_language: str | None = None,
+    ) -> Any:
+        body: dict[str, Any] = {"file_id": file_id, "target_language": target_language}
+        if source_language:
+            body["source_language"] = source_language
+        return self.request("POST", "/api/v1/translations/create", json_body=body)
+
+    def translation_get(self, job_id: str) -> Any:
+        return self.request("GET", f"/api/v1/translations/{_path_id(job_id)}")

cheso_cli/config.py

@@ -0,0 +1,108 @@
+"""站点、基础 URL 的解析及其他进程级配置。"""
+
+from __future__ import annotations
+
+import os
+import shlex
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from importlib.metadata import version
+from urllib.parse import urlsplit
+
+from cheso_cli.errors import UsageError
+
+# ``pyproject.toml`` 是唯一版本来源；安装包和可编辑安装均通过发行包 metadata 读取。
+VERSION = version("cheso-cli")
+USER_AGENT = f"cheso-cli/{VERSION}"
+
+# 公开站点：API 网关、OAuth issuer 和网站使用同一源。
+SITES: dict[str, str] = {
+    "global": "https://cheso.ai",
+    "cn": "https://zhishi.qq.com",
+}
+DEFAULT_SITE = "global"
+
+ENV_BASE_URL = "CHESO_BASE_URL"
+ENV_TOKEN = "CHESO_API_TOKEN"
+# 仅供测试：将请求路由到测试环境中的开发容器。
+ENV_DEV_SET = "CHESO_DEV_SET"
+HEADER_DEV_SET = "X-Dev-Set"
+
+# 网站页面路径，见 services/frontend/src/router/routes.ts。
+PROJECT_PAGE = "/agent/"
+SUBSCRIBE_PAGE = "/agent/subscribe"
+
+
+def normalize_base_url(raw: str) -> str:
+    """校验基础 URL 并去掉末尾斜杠；只接受 http(s) 地址，可以带路径。"""
+    value = raw.strip().rstrip("/")
+    parts = urlsplit(value)
+    if parts.scheme not in ("https", "http") or not parts.hostname:
+        raise UsageError(f"invalid base URL {raw!r}: expected https://host")
+    if parts.query or parts.fragment or parts.username or parts.password:
+        raise UsageError(f"invalid base URL {raw!r}: must not contain credentials, query or fragment")
+    return value
+
+
+@dataclass(frozen=True)
+class Settings:
+    """目标服务地址，以及可选的测试路由请求头。"""
+
+    base_url: str
+    dev_set: str = ""
+    # 记录本次选择站点的 --site / --base-url；来自 CHESO_BASE_URL、上次登录或默认值时为空。
+    # 给调用方建议的后续命令会重复这些选项，确保仍然访问同一站点。
+    option: str = ""
+
+    def command(self, args: str) -> str:
+        """生成访问当前站点的 ``cheso`` 命令行。"""
+        return f"cheso {self.option} {args}" if self.option else f"cheso {args}"
+
+    @property
+    def web_url(self) -> str:
+        return self.base_url
+
+    @property
+    def topup_url(self) -> str:
+        return self.web_url + SUBSCRIBE_PAGE
+
+    def project_url(self, project_id: str | None) -> str:
+        return f"{self.web_url}{PROJECT_PAGE}{project_id}" if project_id else ""
+
+    def extra_headers(self) -> dict[str, str]:
+        """向 Cheso 发送请求时附加的请求头；预签名存储 URL 不使用这些请求头。"""
+        headers = {"User-Agent": USER_AGENT}
+        if self.dev_set:
+            headers[HEADER_DEV_SET] = self.dev_set
+        return headers
+
+
+def resolve_settings(
+    *,
+    site: str | None,
+    base_url: str | None,
+    env: Mapping[str, str] | None = None,
+    logged_in_default: Callable[[], str] | None = None,
+) -> Settings:
+    """优先级：``--base-url`` > ``--site`` > ``CHESO_BASE_URL`` > 上次 ``cheso login`` 的站点 > 国际站。"""
+    env = os.environ if env is None else env
+    option = ""
+    if base_url:
+        resolved = normalize_base_url(base_url)
+        option = f"--base-url {shlex.quote(resolved)}"
+    elif site:
+        if site not in SITES:
+            raise UsageError(f"unknown site {site!r}; choose one of: {', '.join(SITES)}")
+        resolved = SITES[site]
+        option = f"--site {site}"
+    elif env.get(ENV_BASE_URL, "").strip():
+        resolved = normalize_base_url(env[ENV_BASE_URL])
+    else:
+        resolved = SITES[DEFAULT_SITE]
+        remembered = logged_in_default() if logged_in_default else ""
+        if remembered:
+            try:
+                resolved = normalize_base_url(remembered)
+            except UsageError:
+                pass
+    return Settings(base_url=resolved, dev_set=env.get(ENV_DEV_SET, "").strip(), option=option)