kaibo 0.1.0__py3-none-any.whl

kaibo/__init__.py

@@ -0,0 +1,22 @@
+"""kaibo: memory for AI agents. Four operations: remember, recall, explain, forget."""
+
+from ._core import SDK_VERSION
+from .async_client import AsyncKaibo
+from .client import Kaibo
+from .errors import ApiError, KaiboError, NetworkError, RateLimitError, TimeoutError, WebhookSignatureError
+from .render import render_context
+from .webhooks import verify_webhook
+
+__version__ = SDK_VERSION
+__all__ = [
+    "ApiError",
+    "AsyncKaibo",
+    "Kaibo",
+    "KaiboError",
+    "NetworkError",
+    "RateLimitError",
+    "TimeoutError",
+    "WebhookSignatureError",
+    "render_context",
+    "verify_webhook",
+]

kaibo/_core.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import json
+import random
+import uuid
+from dataclasses import dataclass
+from typing import Any, Iterable, Mapping, Optional
+
+import httpx
+
+from .errors import ApiError, KaiboError, RateLimitError
+
+SDK_VERSION = "0.1.0"
+
+
+def env_default(name: str) -> Optional[str]:
+    """An environment variable, or None when unset or blank."""
+    import os
+
+    v = os.environ.get(name, "").strip()
+    return v or None
+RETRYABLE_STATUS = frozenset({408, 425, 429, 500, 502, 503, 504})
+
+
+def backoff_seconds(attempt: int, rand: float, base: float = 0.25, cap: float = 8.0) -> float:
+    """Exponential backoff with full jitter: uniform in [0, min(cap, base * 2**attempt)]."""
+    return float(rand * min(cap, base * (2**attempt)))
+
+
+def new_idempotency_key() -> str:
+    return f"sdk-{uuid.uuid4()}"
+
+
+@dataclass(frozen=True)
+class Prepared:
+    method: str
+    url: str
+    headers: dict[str, str]
+    json: Optional[Any]
+
+
+class Config:
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str,
+        project: Optional[str],
+        retries: int,
+        timeout: float,
+    ) -> None:
+        if not api_key:
+            raise KaiboError(
+                "no api key. Pass api_key=... or set KAIBO_KEY "
+                "(keys look like kaibo_<id>.<secret>; create one in the dashboard, under api keys)."
+            )
+        if not base_url:
+            raise KaiboError(
+                "no address. Pass base_url=... or set KAIBO_URL to the address you open "
+                "the kaibo dashboard at (for example http://localhost:3000)."
+            )
+        self.api_key = api_key
+        self.base_url = base_url.rstrip("/")
+        self.project = project
+        self.retries = retries
+        self.timeout = timeout
+
+    def prepare(
+        self,
+        method: str,
+        path: str,
+        *,
+        query: Optional[Mapping[str, Any]] = None,
+        body: Optional[Any] = None,
+        accept: str = "application/json",
+    ) -> Prepared:
+        q = {k: str(v) for k, v in (query or {}).items() if v is not None}
+        url = f"{self.base_url}{path}"
+        if q:
+            url += "?" + str(httpx.QueryParams(q))
+        headers = {
+            "accept": accept,
+            "authorization": f"Bearer {self.api_key}",
+            "x-sdk": f"kaibo-sdk-py/{SDK_VERSION}",
+        }
+        if self.project:
+            headers["x-project-id"] = self.project
+        if body is not None:
+            headers["content-type"] = "application/json"
+        return Prepared(method, url, headers, body)
+
+
+def api_error(res: httpx.Response) -> ApiError:
+    code, message = "unknown", f"request failed ({res.status_code})"
+    request_id = res.headers.get("x-request-id")
+    try:
+        err = res.json().get("error")
+        if isinstance(err, dict):
+            code = str(err.get("code", code))
+            message = str(err.get("message", message))
+            request_id = err.get("requestId", request_id)
+    except Exception:  # non-JSON error body
+        pass
+    if res.status_code == 429:
+        try:
+            ra = float(res.headers.get("retry-after", ""))
+        except ValueError:
+            ra = 0.0
+        return RateLimitError(message, request_id=request_id, retry_after=ra if ra > 0 else None)
+    return ApiError(message, code=code, status=res.status_code, request_id=request_id)
+
+
+def retry_wait(err: ApiError, attempt: int, rand: float) -> float:
+    if isinstance(err, RateLimitError) and err.retry_after is not None:
+        return err.retry_after
+    return backoff_seconds(attempt, rand)
+
+
+def jitter() -> float:
+    return random.random()
+
+
+def parse_sse(lines: Iterable[str]) -> Iterable[tuple[str, str]]:
+    """Yields (event, data) for each server-sent event in a stream of text lines."""
+    event: str = "message"
+    data: list[str] = []
+    for line in lines:
+        if line == "":
+            if data:
+                yield event, "\n".join(data)
+            event, data = "message", []
+        elif line.startswith(":"):
+            continue
+        elif line.startswith("event:"):
+            event = line[6:].strip()
+        elif line.startswith("data:"):
+            data.append(line[5:].lstrip(" "))
+    if data:
+        yield event, "\n".join(data)
+
+
+def loads(text: str) -> Any:
+    return json.loads(text)

kaibo/async_client.py

@@ -0,0 +1,453 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, AsyncIterator, Awaitable, Callable, List, Optional, Sequence, Union
+
+import httpx
+
+from ._core import (
+    RETRYABLE_STATUS,
+    env_default,
+    Config,
+    Prepared,
+    api_error,
+    jitter,
+    loads,
+    parse_sse,
+    retry_wait,
+)
+from .client import (
+    _RECALL_EVENTS,
+    _export_body,
+    _export_record,
+    _forget_body,
+    _import_body,
+    _key_body,
+    _q,
+    _recall_body,
+    _remember_body,
+)
+from .render import render_context
+from .errors import ApiError, KaiboError, NetworkError, TimeoutError
+from .models import (
+    ApiKey,
+    MemoryRef,
+    Project,
+    ContextPack,
+    CreatedApiKey,
+    CreatedWebhook,
+    DeletionReceipt,
+    Explanation,
+    ExportRecord,
+    ForgetDryRun,
+    ImportResult,
+    ImportSource,
+    JobStatus,
+    KeyScope,
+    Memory,
+    MemoryStatus,
+    MemoryType,
+    RecallEvent,
+    ReceiptVerification,
+    RememberResult,
+    StoredEvent,
+    Usage,
+    Webhook,
+    WebhookDelivery,
+    WebhookEventType,
+    WhoAmI,
+)
+
+ASleep = Callable[[float], Awaitable[None]]
+
+
+class AsyncKaibo:
+    """asyncio client with the same surface as `Kaibo`.
+
+    >>> async with AsyncKaibo(api_key=os.environ["KAIBO_KEY"]) as kaibo:
+    ...     await kaibo.remember(text="We ship on July 17.", wait=True)
+    ...     pack = await kaibo.recall("ship date")
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        *,
+        base_url: Optional[str] = None,
+        project: Optional[str] = None,
+        retries: int = 2,
+        timeout: float = 30.0,
+        http_client: Optional[httpx.AsyncClient] = None,
+        sleep: ASleep = asyncio.sleep,
+    ) -> None:
+        self._cfg = Config(
+            api_key or env_default("KAIBO_KEY") or "",
+            base_url or env_default("KAIBO_URL") or "",
+            project or env_default("KAIBO_PROJECT"),
+            retries,
+            timeout,
+        )
+        self._http = http_client or httpx.AsyncClient(timeout=timeout)
+        self._owns_http = http_client is None
+        self._sleep = sleep
+        self.events = _AEvents(self)
+        self.memories = _AMemories(self)
+        self.jobs = _AJobs(self)
+        self.webhooks = _AWebhooks(self)
+        self.imports = _AImports(self)
+        self.exports = _AExports(self)
+        self.keys = _AKeys(self)
+        self.projects = _AProjects(self)
+
+    async def aclose(self) -> None:
+        if self._owns_http:
+            await self._http.aclose()
+
+    async def __aenter__(self) -> "AsyncKaibo":
+        return self
+
+    async def __aexit__(self, *_: object) -> None:
+        await self.aclose()
+
+    async def _send(self, p: Prepared, *, retry: bool, stream: bool = False) -> httpx.Response:
+        attempts = (self._cfg.retries if retry else 0) + 1
+        last: Optional[Exception] = None
+        for attempt in range(attempts):
+            try:
+                req = self._http.build_request(
+                    p.method, p.url, headers=p.headers, json=p.json, timeout=self._cfg.timeout
+                )
+                res = await self._http.send(req, stream=stream)
+                if res.is_success:
+                    return res
+                if stream:
+                    await res.aread()
+                err = api_error(res)
+                await res.aclose()
+                if res.status_code in RETRYABLE_STATUS and attempt < attempts - 1:
+                    last = err
+                    await self._sleep(retry_wait(err, attempt, jitter()))
+                    continue
+                raise err
+            except ApiError:
+                raise
+            except httpx.TimeoutException as e:
+                last = TimeoutError(f"request timed out after {self._cfg.timeout}s")
+                last.__cause__ = e
+            except httpx.HTTPError as e:
+                last = NetworkError("network request failed")
+                last.__cause__ = e
+            if attempt >= attempts - 1:
+                raise last
+            await self._sleep(retry_wait(ApiError("", code="x", status=0, request_id=None), attempt, jitter()))
+        raise last or KaiboError("request failed")
+
+    async def _request(self, method: str, path: str, *, retry: bool, query: Any = None, body: Any = None) -> Any:
+        res = await self._send(self._cfg.prepare(method, path, query=query, body=body), retry=retry)
+        return None if res.status_code == 204 else res.json()
+
+    async def remember(
+        self,
+        text: Optional[str] = None,
+        *,
+        events: Optional[list[dict[str, Any]]] = None,
+        stream: Optional[str] = None,
+        idempotency_key: Optional[str] = None,
+        wait: Union[bool, float] = False,
+    ) -> RememberResult:
+        body = _remember_body(text, events, stream, idempotency_key)
+        result = RememberResult.model_validate(await self._request("POST", "/v1/remember", body=body, retry=True))
+        if wait:
+            await self.jobs.wait(result.job_id, timeout=10.0 if wait is True else float(wait))
+        return result
+
+    async def recall(
+        self,
+        query: str,
+        *,
+        as_of: Optional[str] = None,
+        budget_tokens: Optional[int] = None,
+        types: Optional[Sequence[MemoryType]] = None,
+        scope: Optional[str] = None,
+        mode: Optional[str] = None,
+        min_trust: Optional[str] = None,
+    ) -> ContextPack:
+        body = _recall_body(query, as_of, budget_tokens, types, scope, mode, min_trust)
+        return ContextPack.model_validate(await self._request("POST", "/v1/recall", body=body, retry=True))
+
+    async def context(self, query: str, **options: Any) -> str:
+        """Recall, rendered as text for a prompt (see `Kaibo.context`)."""
+        return render_context(await self.recall(query, **options))
+
+    async def verify_pack(self, pack: ContextPack) -> dict[str, Any]:
+        body = {"pack": pack.model_dump(by_alias=True, mode="json", exclude_none=True)}
+        return await self._request("POST", "/v1/attest/verify", body=body, retry=True)
+
+    async def diff(self, from_: str, to: str) -> dict[str, Any]:
+        return await self._request("POST", "/v1/diff", body={"from": from_, "to": to}, retry=True)
+
+    async def verify_memory(self, memory_id: MemoryRef) -> dict[str, Any]:
+        return await self._request("POST", f"/v1/memories/{_q(memory_id)}/verify", retry=False)
+
+    async def recall_stream(
+        self,
+        query: str,
+        *,
+        as_of: Optional[str] = None,
+        budget_tokens: Optional[int] = None,
+        types: Optional[Sequence[MemoryType]] = None,
+        scope: Optional[str] = None,
+    ) -> AsyncIterator[RecallEvent]:
+        body = _recall_body(query, as_of, budget_tokens, types, scope)
+        p = self._cfg.prepare("POST", "/v1/recall", body=body, accept="text/event-stream")
+        res = await self._send(p, retry=False, stream=True)
+        try:
+            lines: list[str] = []
+            async for line in res.aiter_lines():
+                lines.append(line)
+                if line == "":
+                    for name, data in parse_sse(lines):
+                        model = _RECALL_EVENTS.get(name)
+                        if model:
+                            yield model.model_validate(loads(data))
+                    lines = []
+        finally:
+            await res.aclose()
+
+    async def explain(self, memory_id: MemoryRef, *, depth: Optional[int] = None) -> Explanation:
+        return Explanation.model_validate(
+            await self._request("GET", f"/v1/explain/{_q(memory_id)}", query={"depth": depth}, retry=True)
+        )
+
+    async def forget(
+        self, *, memory_id: Optional[str] = None, query: Optional[str] = None, reason: str
+    ) -> DeletionReceipt:
+        body = _forget_body(memory_id, query, reason, False)
+        out = await self._request("POST", "/v1/forget", body=body, retry=False)
+        return DeletionReceipt.model_validate(out["receipt"])
+
+    async def forget_dry_run(
+        self, *, memory_id: Optional[str] = None, query: Optional[str] = None, reason: str
+    ) -> ForgetDryRun:
+        body = _forget_body(memory_id, query, reason, True)
+        return ForgetDryRun.model_validate(await self._request("POST", "/v1/forget", body=body, retry=False))
+
+    async def verify_receipt(self, receipt_id: str) -> ReceiptVerification:
+        return ReceiptVerification.model_validate(
+            await self._request("GET", f"/v1/forget/receipts/{_q(receipt_id)}/verify", retry=True)
+        )
+
+    async def get_receipt(self, receipt_id: str) -> DeletionReceipt:
+        return DeletionReceipt.model_validate(
+            await self._request("GET", f"/v1/forget/receipts/{_q(receipt_id)}", retry=True)
+        )
+
+    async def whoami(self) -> WhoAmI:
+        return WhoAmI.model_validate(await self._request("GET", "/v1/whoami", retry=True))
+
+    async def usage(self) -> Usage:
+        return Usage.model_validate(await self._request("GET", "/v1/usage", retry=True))
+
+
+class _AEvents:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def get(self, event_id: str) -> StoredEvent:
+        return StoredEvent.model_validate(await self._c._request("GET", f"/v1/events/{_q(event_id)}", retry=True))
+
+    async def page(
+        self, *, stream: Optional[str] = None, limit: Optional[int] = None, cursor: Optional[str] = None
+    ) -> tuple[list[StoredEvent], Optional[str]]:
+        out = await self._c._request(
+            "GET", "/v1/events", query={"stream": stream, "limit": limit, "cursor": cursor}, retry=True
+        )
+        return [StoredEvent.model_validate(e) for e in out["items"]], out["nextCursor"]
+
+    async def iter(self, *, stream: str = "default", page_size: int = 100) -> AsyncIterator[StoredEvent]:
+        cursor: Optional[str] = None
+        while True:
+            items, cursor = await self.page(stream=stream, limit=page_size, cursor=cursor)
+            for e in items:
+                yield e
+            if not cursor:
+                return
+
+
+class _AMemories:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def get(self, memory_id: MemoryRef) -> Memory:
+        return Memory.model_validate(await self._c._request("GET", f"/v1/memories/{_q(memory_id)}", retry=True))
+
+    async def edit(self, memory_id: MemoryRef, content: str) -> Memory:
+        """Replace what a current memory says. The old text stays as a superseded version. Needs `write`. Not retried."""
+        return Memory.model_validate(
+            await self._c._request("PATCH", f"/v1/memories/{_q(memory_id)}", body={"content": content}, retry=False)
+        )
+
+    async def page(
+        self, *, status: Optional[MemoryStatus] = None, limit: Optional[int] = None, cursor: Optional[str] = None
+    ) -> tuple[list[Memory], Optional[str]]:
+        out = await self._c._request(
+            "GET", "/v1/memories", query={"status": status, "limit": limit, "cursor": cursor}, retry=True
+        )
+        return [Memory.model_validate(m) for m in out["items"]], out["nextCursor"]
+
+    async def iter(self, *, status: Optional[MemoryStatus] = None, page_size: int = 100) -> AsyncIterator[Memory]:
+        cursor: Optional[str] = None
+        while True:
+            items, cursor = await self.page(status=status, limit=page_size, cursor=cursor)
+            for m in items:
+                yield m
+            if not cursor:
+                return
+
+
+class _AJobs:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def get(self, job_id: str) -> JobStatus:
+        return JobStatus.model_validate(await self._c._request("GET", f"/v1/jobs/{_q(job_id)}", retry=True))
+
+    async def wait(self, job_id: str, *, timeout: float = 10.0, interval: float = 0.1) -> JobStatus:
+        deadline = time.monotonic() + timeout
+        delay = interval
+        while True:
+            s = await self.get(job_id)
+            if s.done:
+                return s
+            if time.monotonic() + delay > deadline:
+                raise TimeoutError(f"job {job_id} did not finish in time")
+            await self._c._sleep(delay)
+            delay = min(1.0, delay * 2)
+
+
+class _AWebhooks:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def create(self, url: str, *, events: Optional[Sequence[WebhookEventType]] = None) -> CreatedWebhook:
+        body: dict[str, Any] = {"url": url}
+        if events is not None:
+            body["events"] = list(events)
+        return CreatedWebhook.model_validate(await self._c._request("POST", "/v1/webhooks", body=body, retry=False))
+
+    async def list(self) -> List[Webhook]:
+        return [Webhook.model_validate(w) for w in (await self._c._request("GET", "/v1/webhooks", retry=True))["items"]]
+
+    async def deliveries(self, webhook_id: str, *, limit: Optional[int] = None) -> List[WebhookDelivery]:
+        out = await self._c._request(
+            "GET", f"/v1/webhooks/{_q(webhook_id)}/deliveries", query={"limit": limit}, retry=True
+        )
+        return [WebhookDelivery.model_validate(d) for d in out["items"]]
+
+    async def test(self, webhook_id: str) -> None:
+        await self._c._request("POST", f"/v1/webhooks/{_q(webhook_id)}/test", retry=False)
+
+    async def delete(self, webhook_id: str) -> None:
+        await self._c._request("DELETE", f"/v1/webhooks/{_q(webhook_id)}", retry=False)
+
+
+class _AImports:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def create(
+        self,
+        source: ImportSource,
+        content: str,
+        *,
+        stream: Optional[str] = None,
+        name: Optional[str] = None,
+        include_invalid: bool = False,
+        dry_run: bool = False,
+    ) -> ImportResult:
+        out = await self._c._request(
+            "POST",
+            "/v1/imports",
+            query={"dryRun": "true"} if dry_run else None,
+            body=_import_body(source, content, stream, name, include_invalid),
+            retry=True,
+        )
+        return ImportResult.model_validate(out)
+
+
+class _AExports:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def stream(
+        self,
+        *,
+        stream: Optional[str] = None,
+        include: Optional[Sequence[str]] = None,
+        memory_status: Optional[Sequence[str]] = None,
+    ) -> AsyncIterator[ExportRecord]:
+        p = self._c._cfg.prepare("POST", "/v1/exports", body=_export_body(stream, include, memory_status))
+        res = await self._c._send(p, retry=False, stream=True)
+        ended = False
+        try:
+            async for line in res.aiter_lines():
+                if not line.strip():
+                    continue
+                rec = _export_record(line)
+                ended = ended or rec.record == "end"
+                yield rec
+        finally:
+            await res.aclose()
+        if not ended:
+            raise KaiboError("the export was cut off before its end record")
+
+    async def create(
+        self,
+        *,
+        stream: Optional[str] = None,
+        include: Optional[Sequence[str]] = None,
+        memory_status: Optional[Sequence[str]] = None,
+    ) -> List[ExportRecord]:
+        return [r async for r in self.stream(stream=stream, include=include, memory_status=memory_status)]
+
+
+class _AProjects:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def list(self) -> List[Project]:
+        return [Project.model_validate(p) for p in (await self._c._request("GET", "/v1/projects", retry=True))["items"]]
+
+    async def create(self, name: str) -> Project:
+        return Project.model_validate(await self._c._request("POST", "/v1/projects", body={"name": name}, retry=False))
+
+    async def delete(self, project_id: str) -> None:
+        await self._c._request("DELETE", f"/v1/projects/{_q(project_id)}", retry=False)
+
+    async def rename(self, project_id: str, name: str) -> Project:
+        return Project.model_validate(
+            await self._c._request("PATCH", f"/v1/projects/{_q(project_id)}", body={"name": name}, retry=True)
+        )
+
+
+class _AKeys:
+    def __init__(self, c: AsyncKaibo) -> None:
+        self._c = c
+
+    async def list(self) -> List[ApiKey]:
+        return [ApiKey.model_validate(k) for k in (await self._c._request("GET", "/v1/keys", retry=True))["items"]]
+
+    async def create(
+        self,
+        label: str,
+        scopes: Sequence[KeyScope],
+        *,
+        project_id: Optional[str] = None,
+        expires_at: Optional[str] = None,
+    ) -> CreatedApiKey:
+        body = _key_body(label, scopes, project_id, expires_at)
+        return CreatedApiKey.model_validate(await self._c._request("POST", "/v1/keys", body=body, retry=False))
+
+    async def revoke(self, key_id: str) -> ApiKey:
+        return ApiKey.model_validate(await self._c._request("DELETE", f"/v1/keys/{_q(key_id)}", retry=True))