ace/cli/client.py

"""HTTP client for the Kayba hosted API."""

from __future__ import annotations

import json
import os
import re
from typing import Any, Dict, List, Optional


class KaybaAPIError(Exception):
    """Structured error from the Kayba API."""

    def __init__(self, code: str, message: str, status_code: int = 0):
        self.code = code
        self.message = message
        self.status_code = status_code
        super().__init__(f"[{code}] {message}")


DEFAULT_BASE_URL = "https://use.kayba.ai/api"
MAX_TRACE_UPLOAD_BODY_BYTES = 900_000


def _chunk_trace_uploads(
    traces: List[Dict[str, Any]],
) -> List[List[Dict[str, Any]]]:
    """Split uploads into request-sized batches under the body size cap."""
    batches: List[List[Dict[str, Any]]] = []
    current: List[Dict[str, Any]] = []
    current_size = len('{"traces":[]}')

    for trace in traces:
        trace_size = len(
            json.dumps(trace, ensure_ascii=False, separators=(",", ":")).encode(
                "utf-8"
            )
        )
        separator_size = 1 if current else 0
        candidate_size = current_size + separator_size + trace_size

        if current and candidate_size > MAX_TRACE_UPLOAD_BODY_BYTES:
            batches.append(current)
            current = [trace]
            current_size = len('{"traces":[]}') + trace_size
            continue

        current.append(trace)
        current_size = candidate_size

    if current:
        batches.append(current)

    return batches


class KaybaClient:
    """HTTP client for the Kayba hosted API.

    Args:
        api_key: Kayba API key. Falls back to KAYBA_API_KEY env var.
        base_url: API base URL. Falls back to KAYBA_API_URL env var,
                  then to https://use.kayba.ai/api.
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
    ):
        try:
            import requests
        except ImportError as exc:
            raise KaybaAPIError(
                "DEPENDENCY_MISSING",
                "The hosted Kayba CLI requires the cloud extra. Install with "
                "`uv add \"ace-framework[cloud]\"` or "
                "`pip install 'ace-framework[cloud]'`.",
            ) from exc

        self.api_key = api_key or os.environ.get("KAYBA_API_KEY", "")
        if not self.api_key:
            raise KaybaAPIError(
                "AUTH_MISSING",
                "No API key provided. Set KAYBA_API_KEY or pass --api-key.",
            )
        self.base_url = (
            base_url or os.environ.get("KAYBA_API_URL") or DEFAULT_BASE_URL
        ).rstrip("/")
        self.session: Any = requests.Session()
        self.session.headers["Authorization"] = f"Bearer {self.api_key}"

    @staticmethod
    def _summarize_http_body(body: str, limit: int = 240) -> str:
        """Collapse whitespace so raw HTML and proxy errors stay readable."""
        snippet = re.sub(r"\s+", " ", body or "").strip()
        if not snippet:
            return "Unexpected non-JSON error from the Kayba API."
        if len(snippet) <= limit:
            return snippet
        return snippet[: limit - 3] + "..."

    def _request(
        self,
        method: str,
        path: str,
        *,
        json: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, str]] = None,
    ) -> Any:
        """Send a request and return parsed JSON, raising on API errors."""
        url = f"{self.base_url}{path}"
        resp = self.session.request(method, url, json=json, params=params)

        if resp.status_code >= 400:
            try:
                body = resp.json()
                err = body.get("error", {})
                if isinstance(err, str):
                    raise KaybaAPIError(
                        code="API_ERROR",
                        message=err,
                        status_code=resp.status_code,
                    )
                message = err.get("message", resp.text)
                if (
                    resp.status_code == 413
                    or "maximum content size" in message.lower()
                    or "too large" in message.lower()
                ):
                    raise KaybaAPIError(
                        code="PAYLOAD_TOO_LARGE",
                        message=message,
                        status_code=resp.status_code,
                    )
                raise KaybaAPIError(
                    code=err.get("code", "UNKNOWN"),
                    message=message,
                    status_code=resp.status_code,
                )
            except (ValueError, KeyError, AttributeError):
                message = self._summarize_http_body(resp.text)
                if resp.status_code == 413:
                    message = (
                        "Upload rejected because the request body is too large. "
                        "Try smaller traces or upload fewer files at once."
                    )
                elif resp.status_code in (401, 403):
                    message = "Authentication failed; check KAYBA_API_KEY"
                else:
                    message = f"HTTP {resp.status_code} from Kayba API: {message}"
                raise KaybaAPIError(
                    code="HTTP_ERROR",
                    message=message,
                    status_code=resp.status_code,
                )

        if resp.status_code == 204:
            return {}
        return resp.json()

    # -- Traces --

    def upload_traces(self, traces: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Upload trace files.

        Args:
            traces: List of dicts with keys: filename, content, fileType.
        """
        batches = _chunk_trace_uploads(traces)
        if len(batches) == 1:
            return self._request("POST", "/traces", json={"traces": traces})

        combined: Dict[str, Any] = {"count": 0, "traces": []}
        for batch in batches:
            result = self._request("POST", "/traces", json={"traces": batch})
            uploaded = result.get("traces", [])
            combined["count"] += result.get("count", len(uploaded) or len(batch))
            combined["traces"].extend(uploaded)
            for key, value in result.items():
                if key not in {"count", "traces"} and key not in combined:
                    combined[key] = value
        return combined

    def list_traces(self) -> Dict[str, Any]:
        """List all traces (metadata only, no content)."""
        return self._request("GET", "/traces")

    def get_trace(self, trace_id: str) -> Dict[str, Any]:
        """Get a single trace with full content."""
        return self._request("GET", f"/traces/{trace_id}")

    def get_traces(self, trace_ids: List[str]) -> Dict[str, Any]:
        """Batch get traces by IDs (with content)."""
        return self._request("POST", "/traces/batch", json={"ids": trace_ids})

    def delete_trace(self, trace_id: str) -> Dict[str, Any]:
        """Delete a single trace."""
        return self._request("DELETE", f"/traces/{trace_id}")

    def delete_traces(self, trace_ids: List[str]) -> Dict[str, Any]:
        """Delete multiple traces."""
        results = []
        errors = []
        for tid in trace_ids:
            try:
                self.delete_trace(tid)
                results.append(tid)
            except KaybaAPIError as e:
                errors.append({"id": tid, "error": str(e)})
        return {"deleted": results, "errors": errors}

    # -- Insights --

    def generate_insights(
        self,
        *,
        trace_ids: Optional[List[str]] = None,
        model: Optional[str] = None,
        epochs: Optional[int] = None,
        reflector_mode: Optional[str] = None,
        anthropic_key: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Start async insight generation."""
        body: Dict[str, Any] = {}
        if trace_ids:
            body["traceIds"] = trace_ids
        if model:
            body["model"] = model
        if epochs is not None:
            body["epochs"] = epochs
        if reflector_mode:
            body["reflectorMode"] = reflector_mode
        if anthropic_key:
            body["anthropicApiKey"] = anthropic_key
        return self._request("POST", "/insights/generate", json=body)

    def list_insights(
        self,
        *,
        status: Optional[str] = None,
        section: Optional[str] = None,
    ) -> Dict[str, Any]:
        """List insights, optionally filtered."""
        params: Dict[str, str] = {}
        if status:
            params["status"] = status
        if section:
            params["section"] = section
        return self._request("GET", "/insights", params=params or None)

    def triage_insight(
        self,
        insight_id: str,
        status: str,
        note: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Accept or reject a single insight."""
        body: Dict[str, Any] = {"status": status}
        if note:
            body["note"] = note
        return self._request("PATCH", f"/insights/{insight_id}", json=body)

    # -- Jobs --

    def get_job(self, job_id: str) -> Dict[str, Any]:
        """Get job status."""
        return self._request("GET", f"/jobs/{job_id}")

    def materialize_job(self, job_id: str) -> Dict[str, Any]:
        """Materialize completed job results into the skillbook."""
        return self._request("POST", f"/jobs/{job_id}")

    # -- Prompts --

    def generate_prompt(
        self,
        *,
        insight_ids: Optional[List[str]] = None,
        label: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Generate a prompt from accepted insights."""
        body: Dict[str, Any] = {}
        if insight_ids:
            body["insightIds"] = insight_ids
        if label:
            body["label"] = label
        return self._request("POST", "/prompts/generate", json=body)

    def list_prompts(self) -> Dict[str, Any]:
        """List all prompt versions."""
        return self._request("GET", "/prompts")

    def get_prompt(self, prompt_id: str) -> Dict[str, Any]:
        """Get a specific prompt by ID."""
        return self._request("GET", f"/prompts/{prompt_id}")

    # -- Integrations --

    def get_integrations(self) -> Dict[str, Any]:
        """Get current integration settings."""
        return self._request("GET", "/integrations")

    def update_integration(self, name: str, config: Dict[str, Any]) -> Dict[str, Any]:
        """Update an integration's config."""
        return self._request("PUT", f"/integrations/{name}", json=config)

    def test_integration(self, name: str) -> Dict[str, Any]:
        """Test an integration connection."""
        return self._request("POST", f"/integrations/{name}/test")