agentengine-sdk-python 0.8.0__py3-none-any.whl → 0.8.2__py3-none-any.whl

ksadk/studio/model_profile_service.py

@@ -0,0 +1,214 @@
+"""Model Profile connection diagnostics."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any
+from urllib.parse import urlparse
+
+import httpx
+
+from ksadk.studio.contracts import ModelSpec, NetworkPolicy
+from ksadk.studio.errors import StudioError
+
+
+async def test_model_profile_connection(
+    *,
+    catalog: Any,
+    model_client: Any,
+    resource_id: str,
+) -> dict:
+    descriptor = catalog.get(resource_id)
+    if descriptor.kind != "model":
+        raise StudioError(
+            "RESOURCE_KIND_INVALID",
+            "连接测试只能用于 Model Profile",
+            status_code=422,
+            details={"resourceId": resource_id},
+        )
+    spec = ModelSpec.model_validate(descriptor.contract)
+    resolved = catalog.resolver.resolve_model(spec)
+    resolved.parameters = resolved.parameters.model_copy(
+        update={
+            "temperature": 0,
+            "max_tokens": min(resolved.parameters.max_tokens or 2048, 64),
+        }
+    )
+    host = (urlparse(resolved.endpoint_url).hostname or "").lower().rstrip(".")
+    started = time.monotonic()
+    response = await model_client.complete(
+        resolved,
+        messages=[{"role": "user", "content": "这是连接测试。请只回复 OK。"}],
+        network_policy=NetworkPolicy(
+            mode="restricted",
+            allowed_hosts=[host] if host else [],
+            allow_private_network=False,
+        ),
+        timeout_seconds=20,
+        max_attempts=1,
+        backoff_seconds=0,
+        allow_empty=True,
+    )
+    return {
+        "ok": True,
+        "resourceId": resource_id,
+        "model": resolved.model,
+        "finishReason": response.finish_reason,
+        "latencyMs": int((time.monotonic() - started) * 1000),
+    }
+
+
+def _probe_candidates(raw_url: str) -> tuple[str, list[tuple[str, str]]]:
+    """把用户粘贴的任意形态地址归一化为 (base_url, [(protocol, endpoint_url), ...])。
+
+    支持三种输入形态：
+    - https://host/v1/chat/completions → 直接候选
+    - https://host/v1/responses → 直接候选
+    - https://host 或 https://host/v1 → 推导 chat / responses 候选
+    """
+    url = raw_url.strip().rstrip("/")
+    if url.endswith("/chat/completions"):
+        base = url[: -len("/chat/completions")]
+        return base, [("chat", url)]
+    if url.endswith("/responses"):
+        base = url[: -len("/responses")]
+        return base, [("responses", url)]
+    if url.endswith("/v1"):
+        return url, [("chat", f"{url}/chat/completions"), ("responses", f"{url}/responses")]
+    return url, [
+        ("chat", f"{url}/v1/chat/completions"),
+        ("responses", f"{url}/v1/responses"),
+        ("chat", f"{url}/chat/completions"),
+        ("responses", f"{url}/responses"),
+    ]
+
+
+def _classify_probe(status_code: int) -> str:
+    if status_code == 200:
+        return "ok"
+    if status_code in {401, 403}:
+        return "auth_required"
+    if status_code in {400, 422, 429}:
+        return "recognized"
+    if status_code in {404, 405, 406}:
+        return "unavailable"
+    return "error"
+
+
+async def probe_model_endpoint(
+    *,
+    url: str,
+    credential: str | None,
+    network_guard: Any = None,
+    timeout_seconds: float = 8.0,
+) -> dict[str, Any]:
+    """智能探测模型端点：识别 chat / responses 协议、可用模型清单并给出推荐配置。"""
+    base, candidates = _probe_candidates(url)
+    host = (urlparse(base).hostname or "").lower().rstrip(".")
+    headers = {"Content-Type": "application/json"}
+    if credential:
+        headers["Authorization"] = f"Bearer {credential}"
+
+    async def _post(
+        client: httpx.AsyncClient,
+        protocol: str,
+        endpoint: str,
+        probe_model: str,
+    ) -> dict[str, Any]:
+        if network_guard is not None:
+            await network_guard.check(
+                endpoint,
+                NetworkPolicy(
+                    mode="restricted",
+                    allowed_hosts=[host] if host else [],
+                    allow_private_network=False,
+                ),
+            )
+        payload: dict[str, Any] = (
+            {
+                "model": probe_model,
+                "messages": [{"role": "user", "content": "ping"}],
+                "max_tokens": 1,
+                "stream": False,
+            }
+            if protocol == "chat"
+            else {"model": probe_model, "input": "ping", "max_output_tokens": 16}
+        )
+        started = time.monotonic()
+        try:
+            response = await client.post(endpoint, headers=headers, json=payload)
+        except (httpx.TimeoutException, httpx.TransportError) as exc:
+            return {
+                "protocol": protocol,
+                "endpointUrl": endpoint,
+                "status": "unreachable",
+                "errorType": type(exc).__name__,
+            }
+        return {
+            "protocol": protocol,
+            "endpointUrl": endpoint,
+            "status": _classify_probe(response.status_code),
+            "httpStatus": response.status_code,
+            "latencyMs": int((time.monotonic() - started) * 1000),
+        }
+
+    async def _models(client: httpx.AsyncClient) -> list[str]:
+        models_urls = (
+            (f"{base}/models", f"{base}/v1/models")
+            if not base.endswith("/v1")
+            else (f"{base}/models",)
+        )
+        for models_url in models_urls:
+            try:
+                response = await client.get(models_url, headers=headers)
+            except (httpx.TimeoutException, httpx.TransportError):
+                continue
+            if response.status_code != 200:
+                continue
+            try:
+                data = response.json()
+            except ValueError:
+                continue
+            items = data.get("data") or data.get("models") or []
+            names = [
+                str(item.get("id") or item.get("name") or "")
+                for item in items
+                if isinstance(item, dict)
+            ]
+            names = [name for name in names if name]
+            if names:
+                return sorted(set(names))[:100]
+        return []
+
+    timeout = httpx.Timeout(timeout_seconds, connect=min(5.0, timeout_seconds))
+    async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client:
+        # 先拿模型清单：部分网关对不存在的模型名直接 403，用真实模型探测才能区分协议
+        models = await _models(client)
+        probe_model = models[0] if models else "__probe__"
+        attempts = await asyncio.gather(
+            *(_post(client, protocol, endpoint, probe_model) for protocol, endpoint in candidates)
+        )
+
+    recommended: dict[str, Any] | None = None
+    for wanted in ("ok", "auth_required", "recognized"):
+        hit = next((a for a in attempts if a["status"] == wanted), None)
+        if hit is not None:
+            recommended = {
+                "protocol": hit["protocol"],
+                "wireApi": "responses" if hit["protocol"] == "responses" else "chat",
+                "endpointUrl": hit["endpointUrl"],
+                "status": hit["status"],
+            }
+            break
+
+    return {
+        "input": url,
+        "baseUrl": base,
+        "attempts": list(attempts),
+        "recommended": recommended,
+        "models": models,
+    }
+
+
+__all__ = ["test_model_profile_connection", "probe_model_endpoint"]

ksadk/studio/operations.py

@@ -0,0 +1,275 @@
+"""Persistent asynchronous operations with idempotency and resumable events."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Awaitable, Callable, cast
+from uuid import uuid4
+
+from pydantic import ValidationError
+
+from ksadk.studio.contracts import (
+    Operation,
+    OperationEvent,
+    OperationKind,
+    OperationStatus,
+)
+from ksadk.studio.errors import StudioError, not_found
+from ksadk.studio.workspace import Workspace
+
+logger = logging.getLogger(__name__)
+
+
+class OperationManager:
+    def __init__(self, workspace: Workspace) -> None:
+        self.workspace = workspace
+        self._tasks: dict[str, asyncio.Task] = {}
+        self._recover_interrupted()
+
+    def submit(
+        self,
+        *,
+        kind: OperationKind,
+        resource_id: str,
+        idempotency_key: str,
+        metadata: dict | None = None,
+        runner: Callable[[str], Awaitable[object]],
+    ) -> Operation:
+        existing = self._find_by_idempotency_key(idempotency_key)
+        if existing is not None:
+            return existing
+        operation = Operation(
+            id=f"op_{uuid4().hex}",
+            kind=kind,
+            resource_id=resource_id,
+            metadata=metadata or {},
+        )
+        self._write(operation, [], idempotency_key)
+        self.append(operation.id, "operation.queued", {"kind": kind})
+        task = asyncio.create_task(self._run(operation.id, runner))
+        self._tasks[operation.id] = task
+
+        def remove_completed(_task: asyncio.Task, op_id: str = operation.id) -> None:
+            self._tasks.pop(op_id, None)
+
+        task.add_done_callback(remove_completed)
+        return operation
+
+    async def _run(
+        self,
+        operation_id: str,
+        runner: Callable[[str], Awaitable[object]],
+    ) -> None:
+        operation = self.get(operation_id)
+        operation.status = OperationStatus.RUNNING
+        self._save_record(operation)
+        self.append(operation_id, "operation.started", {})
+        try:
+            result = await runner(operation_id)
+            result_id = getattr(result, "id", None)
+            if result_id:
+                operation.resource_id = str(result_id)
+            operation.status = OperationStatus.SUCCEEDED
+            operation.completed_at = datetime.now(timezone.utc)
+            self._save_record(operation)
+            self.append(
+                operation_id,
+                "operation.succeeded",
+                {"resourceId": operation.resource_id},
+            )
+        except asyncio.CancelledError:
+            operation.status = OperationStatus.CANCELLED
+            operation.completed_at = datetime.now(timezone.utc)
+            self._save_record(operation)
+            self.append(operation_id, "operation.cancelled", {})
+        except StudioError as exc:
+            operation.status = OperationStatus.FAILED
+            operation.error = {
+                "code": exc.code,
+                "message": exc.message,
+                "field": exc.field,
+            }
+            operation.completed_at = datetime.now(timezone.utc)
+            self._save_record(operation)
+            self.append(operation_id, "operation.failed", operation.error)
+        except Exception as exc:
+            # Keep the browser response generic so an exception cannot leak a
+            # credential, but retain the traceback in the local Studio log for
+            # an operator to diagnose a failed deployment.
+            logger.exception("Studio operation failed: operation_id=%s", operation_id)
+            operation.status = OperationStatus.FAILED
+            operation.error = {
+                "code": "INTERNAL_ERROR",
+                "message": "本地操作执行失败",
+                "exceptionType": type(exc).__name__,
+            }
+            # TypeError carries only Python call-shape information and is safe
+            # to surface to the local operator.  Do not expose arbitrary
+            # exception text: it may include provider request data.
+            if isinstance(exc, TypeError):
+                operation.error["exceptionMessage"] = str(exc)
+            operation.completed_at = datetime.now(timezone.utc)
+            self._save_record(operation)
+            self.append(operation_id, "operation.failed", operation.error)
+
+    def get(self, operation_id: str) -> Operation:
+        operation, _, _ = self._read(operation_id)
+        return operation
+
+    def list(self, *, kind: OperationKind | None = None) -> list[Operation]:
+        directory = self.workspace.resolve(".agentkit/operations")
+        operations: list[Operation] = []
+        for path in directory.glob("op_*.json"):
+            try:
+                operation = self.get(path.stem)
+            except StudioError:
+                continue
+            if kind is None or operation.kind == kind:
+                operations.append(operation)
+        return sorted(operations, key=lambda item: item.created_at, reverse=True)
+
+    def events(self, operation_id: str, *, after: int = 0) -> list[OperationEvent]:
+        _, events, _ = self._read(operation_id)
+        return [event for event in events if event.id > after]
+
+    def append(self, operation_id: str, event_type: str, data: dict) -> OperationEvent:
+        operation, events, key = self._read(operation_id)
+        event = OperationEvent(
+            id=len(events) + 1,
+            operation_id=operation_id,
+            type=event_type,
+            data=data,
+        )
+        events.append(event)
+        self._write(operation, events, key)
+        return event
+
+    def cancel(self, operation_id: str) -> Operation:
+        operation = self.get(operation_id)
+        if operation.status in {
+            OperationStatus.SUCCEEDED,
+            OperationStatus.FAILED,
+            OperationStatus.CANCELLED,
+            OperationStatus.INTERRUPTED,
+        }:
+            return operation
+        task = self._tasks.get(operation_id)
+        if task is not None:
+            task.cancel()
+        if operation.status == OperationStatus.QUEUED:
+            operation.status = OperationStatus.CANCELLED
+            operation.completed_at = datetime.now(timezone.utc)
+            self._save_record(operation)
+            self.append(operation_id, "operation.cancelled", {})
+        return self.get(operation_id)
+
+    async def wait(self, operation_id: str, *, timeout: float = 30) -> Operation:
+        deadline = asyncio.get_running_loop().time() + timeout
+        while True:
+            operation = self.get(operation_id)
+            if operation.status in {
+                OperationStatus.SUCCEEDED,
+                OperationStatus.FAILED,
+                OperationStatus.CANCELLED,
+                OperationStatus.INTERRUPTED,
+            }:
+                return operation
+            if asyncio.get_running_loop().time() >= deadline:
+                raise TimeoutError(operation_id)
+            await asyncio.sleep(0.01)
+
+    def _path(self, operation_id: str) -> Path:
+        return self.workspace.resolve(
+            Path(".agentkit/operations") / f"{operation_id}.json"
+        )
+
+    def _save_record(self, operation: Operation) -> None:
+        _, events, key = self._read(operation.id)
+        self._write(operation, events, key)
+
+    def _write(
+        self,
+        operation: Operation,
+        events: list[OperationEvent],
+        idempotency_key: str,
+    ) -> None:
+        payload = {
+            "operation": operation.model_dump(
+                by_alias=True, exclude_none=True, mode="json"
+            ),
+            "events": [
+                event.model_dump(by_alias=True, exclude_none=True, mode="json")
+                for event in events
+            ],
+            "idempotencyKey": idempotency_key,
+        }
+        self.workspace.atomic_write_text(
+            self._path(operation.id),
+            json.dumps(payload, ensure_ascii=False, sort_keys=True, indent=2) + "\n",
+        )
+
+    def _read(
+        self, operation_id: str
+    ) -> tuple[Operation, list[OperationEvent], str]:
+        path = self._path(operation_id)
+        if not path.is_file():
+            raise not_found("operation", operation_id)
+        try:
+            payload = json.loads(path.read_text(encoding="utf-8"))
+            operation = Operation.model_validate(payload["operation"])
+            events = [
+                OperationEvent.model_validate(item) for item in payload["events"]
+            ]
+            return operation, events, str(payload["idempotencyKey"])
+        except (OSError, ValueError, KeyError, ValidationError) as exc:
+            raise StudioError(
+                "OPERATION_RECORD_INVALID",
+                "Operation 记录损坏",
+                status_code=500,
+                details={"id": operation_id},
+            ) from exc
+
+    def _find_by_idempotency_key(self, key: str) -> Operation | None:
+        directory = self.workspace.resolve(".agentkit/operations")
+        for path in directory.glob("op_*.json"):
+            try:
+                payload = json.loads(path.read_text(encoding="utf-8"))
+                if payload.get("idempotencyKey") == key:
+                    return cast(Operation, Operation.model_validate(payload["operation"]))
+            except (OSError, ValueError, KeyError, ValidationError):
+                continue
+        return None
+
+    def _recover_interrupted(self) -> None:
+        directory = self.workspace.resolve(".agentkit/operations")
+        for path in directory.glob("op_*.json"):
+            try:
+                payload = json.loads(path.read_text(encoding="utf-8"))
+                operation = Operation.model_validate(payload["operation"])
+                events = [
+                    OperationEvent.model_validate(item)
+                    for item in payload.get("events", [])
+                ]
+                key = str(payload["idempotencyKey"])
+            except (OSError, ValueError, KeyError, ValidationError):
+                continue
+            if operation.status not in {
+                OperationStatus.QUEUED,
+                OperationStatus.RUNNING,
+            }:
+                continue
+            operation.status = OperationStatus.INTERRUPTED
+            operation.completed_at = datetime.now(timezone.utc)
+            events.append(
+                OperationEvent(
+                    id=len(events) + 1,
+                    operation_id=operation.id,
+                    type="operation.interrupted",
+                    data={"reason": "daemon_restart"},
+                )
+            )
+            self._write(operation, events, key)