rosen-sdk 0.2.0__py3-none-any.whl

rosen/__init__.py

@@ -0,0 +1,314 @@
+"""Rosen Python agent SDK (local MVP).
+
+An agent is a function ``handler(ctx: Context) -> str | Result``. The Rosen
+runtime invokes the agent as a subprocess, writing a JSON request to stdin and
+reading a JSON response from stdout. ``run(handler)`` implements that protocol
+so agent authors only write ``handler``.
+"""
+from __future__ import annotations
+
+import dataclasses
+import http.server
+import threading
+import os
+import json
+import sys
+import traceback
+from typing import Callable, Union, TypedDict
+from .model import Engine, Model
+from .tool import Tool, Tools, Transport
+from .memory import Memories, Memory, MemoryBackend
+
+__all__ = [
+    "Context",
+    "Engine",
+    "Result",
+    "Model",
+    "Tool",
+    "Tools",
+    "Transport",
+    "Memory",
+    "Memories",
+    "MemoryBackend",
+    "run",
+    "serve"]
+
+
+@dataclasses.dataclass
+class Context:
+    """Everything an agent receives for a single run."""
+
+    input: str
+    trace_id: str = ""
+    agent: str = ""
+    namespace: str = ""
+    conversation_id: str = ""
+    model: Model = dataclasses.field(default_factory=Model)
+    tools: Tools = dataclasses.field(default_factory=Tools)
+    memory: Memories = dataclasses.field(default_factory=Memories)
+
+    def log(self, message: str) -> None:
+        """Emit a trace line to stderr (stdout is the JSON result channel)."""
+        print(f"[rosen trace {self.trace_id}] {message}", file=sys.stderr)
+
+
+@dataclasses.dataclass
+class Result:
+    """An explicit agent result. Returning a plain ``str`` is also accepted."""
+
+    output: str
+
+
+class ToolSpec(TypedDict, total=False):
+    name: str
+    transport: str
+    endpoint: str
+    permissions: str
+
+
+class ModelSpec(TypedDict, total=False):
+    engine: str
+    model: str
+
+
+class MemorySpec(TypedDict, total=False):
+    name: str
+    backend: str
+    endpoint: str
+    collection: str
+    mode: str
+    embedder: str
+
+
+class Request(TypedDict, total=False):  # total=False = every key optional
+    input: str
+    traceId: str
+    namespace: str
+    conversationId: str
+    model: ModelSpec
+    tools: list[ToolSpec]
+    memory: list[MemorySpec]
+
+
+Handler = Callable[[Context], Union[str, Result]]
+
+
+def _build_context(req: dict, memory: Memories | None = None) -> Context:
+    """
+    Build a Context from a decoded JSON request (shared by run and server).
+    """
+    model_info = req.get("model") or {}
+    tools = Tools([
+        Tool(
+            name=t.get("name", ""),
+            transport=t.get("transport", ""),
+            endpoint=t.get("endpoint", ""),
+            permissions=t.get("permissions", "")
+        )
+        for t in (req.get("tools") or [])
+    ])
+
+    if memory is None:
+        memory = _build_memories(req)
+
+    conversation_id = req.get("conversationId", "")
+
+    return Context(
+        input=req.get("input", ""),
+        trace_id=req.get("traceId", ""),
+        agent=req.get("agent", ""),
+        namespace=req.get("namespace", ""),
+        conversation_id=conversation_id,
+        model=Model(engine=model_info.get("engine", ""),
+                    model=model_info.get("model", ""),
+                    tools=tools,
+                    memory=memory,
+                    conversation_id=conversation_id),
+        tools=tools,
+        memory=memory
+    )
+
+
+def _build_memories(req: dict) -> Memories:
+    return Memories([
+        Memory(
+            name=m.get("name", ""),
+            backend=m.get("backend", ""),
+            endpoint=m.get("endpoint", ""),
+            collection=m.get("collection", ""),
+            mode=m.get("mode", ""),
+            embedder=m.get("embedder")
+        )
+        for m in (req.get("memory") or [])
+    ])
+
+
+def _invoke(
+        handler: Handler,
+        req: dict,
+        base: dict | None = None,
+        base_memory: Memories | None = None) -> dict:
+    """
+    Run the handler for one request, returning an {output}/{error} payload.
+    """
+    request_memory = "memory" in req
+    if base:
+        req = {**base, **req}
+    memory = None if request_memory else base_memory
+    ctx = _build_context(req, memory=memory)
+    try:
+        result = handler(ctx)
+    except Exception as exc:
+        traceback.print_exc(file=sys.stderr)
+        return {"error": str(exc)}
+    output = result.output if isinstance(result, Result) else str(result)
+    return {"output": output}
+
+
+def run(handler: Handler, stdin=None, stdout=None) -> None:
+    """
+    Read one JSON request, invoke ``handler``, write one JSON response.
+
+    The request comes from ``stdin``; if that is empty (a one-shot k8s Job pod has
+    no stdin from the controller) it falls back to the ``ROSEN_REQUEST`` env var.
+    The response is written to ``stdout`` and, whe ``ROSEN_RESULT_PATH`` is set,
+    mirrored there (the Job's termination-log) so k8s surfaces it in the pod's
+    termination.message -- a structured result channel, no log scrapping.
+    """
+    stdin = stdin if stdin is not None else sys.stdin
+    stdout = stdout if stdout is not None else sys.stdout
+
+    raw = stdin.read()
+    if not raw.strip():
+        raw = os.environ.get("ROSEN_REQUEST", "")
+    try:
+        req = json.loads(raw) if raw.strip() else {}
+    except json.JSONDecodeError as exc:
+        result = {"error": f"invalid request json: {exc}"}
+    else:
+        result = _invoke(handler, req)
+    _write(stdout, result)
+    _mirror_result(result)
+
+
+def serve(handler: Handler, addr: str | None = None) -> None:
+    """
+    Run the agent as a long-lived HTTP server (one process, many requests).
+
+    GET     /healthz        -> 200: the readiness signal the controller probes.
+    POST    /invoke         -> 202: accept a run, execute the handler on a background
+                                    thread, and ack immediatly with {runId, worker}.
+                                    The worker is this pod's own address so the controller polls
+                                    the exact pod that holds the result.
+    GET     /result/{runId} -> 200 {status: running|done|error, ...} | 404 unknown.
+
+    Asynchronous by design: the handler never holds the response open, so a slow run
+    can't trip a client-side invoke timeout (the "Failed-but-persisted" split-brain).
+
+    The bind address somes from ``addr`` or ``ROSEN_SERVE_ADDR`` (default 127.0.0.1:8080);
+    ``ROSEN_POD_IP`` (injected by the controller via the downward API) is the address
+    advertised for polling, falling back to the bind host.
+    """
+
+    addr = addr or os.environ.get("ROSEN_SERVE_ADDR", "127.0.0.1:8080")
+    host, _, port = addr.rpartition(":")
+    worker = f"{os.environ.get('ROSEN_POD_IP', host)}:{port}"
+
+    raw_bindings = os.environ.get("ROSEN_AGENT_BINDINGS")
+    base = json.loads(raw_bindings) if raw_bindings else {}
+    base_memory = _build_memories(base)
+
+    results: dict[str, dict] = {}
+    lock = threading.Lock()
+
+    def _run(run_id: str, req: Request) -> None:
+        out = _invoke(handler, req, base, base_memory)
+        with lock:
+            if "error" in out:
+                results[run_id] = {"status": "error", "error": out["error"]}
+            else:
+                results[run_id] = {"status": "done", "output": out["output"]}
+
+    class _Handler(http.server.BaseHTTPRequestHandler):
+        def do_GET(self) -> None:
+            if self.path == "/healthz":
+                self._respond(200, {"status": "ok"})
+            elif self.path.startswith("/result/"):
+                run_id = self.path[len("/result/"):]
+                with lock:
+                    result = results.get(run_id)
+                if result is None:
+                    self._respond(404, {"status": "unknown"})
+                else:
+                    self._respond(200, result)
+            else:
+                self._respond(404, {"error": "not found"})
+
+        def do_POST(self) -> None:
+            if self.path != "/invoke":
+                self._respond(404, {"error": "not found"})
+                return
+            lenght = int(self.headers.get("Content-Length", 0))
+            raw = self.rfile.read(lenght).decode("utf-8")
+            try:
+                req = json.loads(raw) if raw.strip() else {}
+            except json.JSONDecodeError as exc:
+                self._respond(400, {"error": f"invalid request json: {exc}"})
+                return
+
+            run_id = req.get("traceId")
+            if not run_id:
+                self._respond(400, {"error": "missing traceId (run id)"})
+                return
+
+            with lock:
+                fresh = run_id not in results
+                if fresh:
+                    results[run_id] = {"status": "running"}
+            if fresh:
+                threading.Thread(target=_run, args=(run_id, req), daemon=True).start()
+
+            self._respond(202, {"runId": run_id, "worker": worker})
+
+        def _respond(self, code: int, payload: dict) -> None:
+            body = json.dumps(payload).encode("utf-8")
+            try:
+                self.send_response(code)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                self.wfile.write(body)
+            except (BrokenPipeError, ConnectionResetError):
+                print(f"rosen: client disconnected before response (path={self.path})",
+                      file=sys.stderr, flush=True)
+
+        def log_message(self, *args) -> None:
+            pass
+
+    server = http.server.ThreadingHTTPServer((host, int(port)), _Handler)
+    bound = server.server_address
+    print(f"rosen: serving on {bound[0]}:{bound[1]}", file=sys.stderr, flush=True)
+    server.serve_forever()
+
+
+def _write(stdout, payload: dict) -> None:
+    json.dump(payload, stdout)
+    stdout.write("\n")
+    stdout.flush()
+
+
+def _mirror_result(payload: dict) -> None:
+    """
+    Best-effort: mirror the result to ROSEN_RESULT_PATH (the Job's termination-log)
+    so k8s surfaces it in the pod's termination.message. A one-shot Job as no other
+    structured result channel; write failures (e.g. path unwritable locally) are logged
+    and ignored so they never mask the real result on stdout.
+    """
+    path = os.environ.get("ROSEN_RESULT_PATH")
+    if not path:
+        return
+    try:
+        with open(path, "w") as f:
+            f.write(json.dumps(payload))
+    except OSError as exc:
+        print(f"rosen: could not write result to {path}: {exc}", file=sys.stderr, flush=True)

rosen/embedder.py

@@ -0,0 +1,66 @@
+import hashlib
+import os
+
+
+def _stable_hash(token: str) -> int:
+    # Python's hash() is salted per process. Vectors must be
+    # reproducible accross the serve subprocess boundry, so
+    # hash deterministically.
+    return int.from_bytes(hashlib.md5(token.encode()).digest()[:4], "big")
+
+
+class Embedder:
+    dim = 0
+
+    def embed(self, texts: list) -> list:
+        raise NotImplementedError
+
+
+class MockEmbedder(Embedder):
+    """Deterministic, offline embeddings for tests and local dev"""
+
+    def __init__(self, model: str = "", dim: int = 64):
+        self.model = model
+        self.dim = dim
+
+    def embed(self, texts: list) -> list:
+        vectors = []
+        for text in texts:
+            vec = [0.0] * self.dim
+            for token in text.lower().split():
+                vec[_stable_hash(token) % self.dim] += 1
+            vectors.append(vec)
+        return vectors
+
+
+class OpenAIEmbedder(Embedder):
+
+    def __init__(self, model: str):
+        self.model = model
+        self.dim = 0
+
+    def _client(self):
+        base_url = os.environ.get("ROSEN_MODEL_BASE_URL")
+        api_key = os.environ.get("ROSEN_MODEL_API_KEY")
+        if not base_url or not api_key:
+            raise RuntimeError("ROSEN_MODEL_BASE_URL and ROSEN_MODEL_API_KEY must be set")
+        from openai import OpenAI
+        return OpenAI(base_url=base_url, api_key=api_key)
+
+    def embed(self, texts: list) -> list:
+        resp = self._client().embeddings.create(model=self.model, input=texts)
+        vectors = [d.embedding for d in resp.data]
+        if vectors:
+            self.dim = len(vectors[0])
+        return vectors
+
+
+_EMBEDDERS = {
+    "mock": MockEmbedder,
+    "openai": OpenAIEmbedder,
+    "azure-openai": OpenAIEmbedder
+}
+
+
+def build_embedder(engine: str, model: str) -> Embedder:
+    return _EMBEDDERS.get(engine, MockEmbedder)(model)

rosen/memory.py

@@ -0,0 +1,277 @@
+import json
+import uuid
+
+from .embedder import build_embedder
+
+
+class MemoryBackend:
+    def remember(self, text, metadata=None, scope=None) -> str:
+        raise NotImplementedError
+
+    def recall(self, query, k=5, scope=None) -> list:
+        raise NotImplementedError
+
+    def history(self, scope, k=None) -> list:
+        raise NotImplementedError
+
+
+class MockMemory(MemoryBackend):
+
+    def __init__(self, name: str, endpoint: str = "", collection: str = "", embedder=None):
+        self.name = name
+        self.endpoint = endpoint
+        self.collection = collection
+        self._items = []
+
+    def remember(self, text, metadata=None, scope=None) -> str:
+        item_id = str(len(self._items))
+        self._items.append(
+            {"id": item_id, "scope": scope, "text": text, "metadata": metadata or {}})
+        return item_id
+
+    def recall(self, query, k=5, scope=None) -> list:
+        needle = (query or "").lower()
+        hits = [
+            {"id": it["id"], "text": it["text"], "metadata": it["metadata"]}
+            for it in reversed(self._items)
+            if it["scope"] == scope and needle in it["text"].lower()
+        ]
+        return hits
+
+    def history(self, scope, k=None) -> list:
+        turns = [
+            {"id": it["id"], "text": it["text"], "metadata": it["metadata"]}
+            for it in self._items
+            if it["scope"] == scope
+        ]
+        return turns[-k:] if k else turns
+
+
+class RedisMemory(MemoryBackend):
+    """Conversational memory over a shared Redis. Each scope (conversationId) is a
+     Redis list of turns, so a turn landing on any replica reads/writes the same
+     list -- conversations survive across replicas and restarts"""
+
+    def __init__(self, name: str, endpoint: str = "", collection: str = "", embedder=None):
+        import redis
+        self.name = name
+        self.endpoint = endpoint or "redis://127.0.0.1:6379/0"
+        self.collection = collection or "rosen"
+        self._client = redis.Redis.from_url(self.endpoint, decode_responses=True)
+
+    def _key(self, scope) -> str:
+        return f"{self.collection}:{self.name}:{scope or 'global'}"
+
+    def remember(self, text, metadata=None, scope=None) -> str:
+        item_id = uuid.uuid4().hex
+        self._client.rpush(
+            self._key(scope),
+            json.dumps({"id": item_id, "text": text, "metadata": metadata or {}}))
+        return item_id
+
+    def history(self, scope, k=None) -> list:
+        start = -k if k else 0
+        return [json.loads(r) for r in self._client.lrange(self._key(scope), start, -1)]
+
+    def recall(self, query, k=5, scope=None) -> list:
+        """simple keyword scan"""
+        needle = (query or "").lower()
+        items = [json.loads(r) for r in self._client.lrange(self._key(scope), 0, -1)]
+        return [it for it in reversed(items) if needle in it["text"].lower()][:k]
+
+
+class FaissMemory(MemoryBackend):
+    """In-process semantic memory"""
+
+    def __init__(self, name: str, endpoint: str = "", collection: str = "", embedder=None):
+        self.name = name
+        self.endpoint = endpoint
+        self.collection = collection
+        self._embedder = embedder
+        self._items = []
+
+    def _embed(self, text: str) -> list:
+        if self._embedder is None:
+            raise RuntimeError("FaissMemory needs and embedder; bind one via Memory.spec.embedder")
+        return self._embedder.embed([text])[0]
+
+    def remember(self, text, metadata=None, scope=None) -> str:
+        item_id = str(len(self._items))
+        self._items.append({
+            "id": item_id, "scope": scope, "text": text,
+            "metadata": metadata or {}, "vector": self._embed(text)})
+
+        return item_id
+
+    def recall(self, query, k=5, scope=None) -> list:
+        import faiss
+        import numpy as np
+        hits = [it for it in self._items if it["scope"] == scope]
+        if not hits:
+            return []
+
+        vectors = np.array([it["vector"] for it in hits], dtype="float32")
+        faiss.normalize_L2(vectors)
+        index = faiss.IndexFlatIP(vectors.shape[1])
+        index.add(vectors)
+
+        q = np.array([self._embed(query)], dtype="float32")
+        faiss.normalize_L2(q)
+
+        _, idxs = index.search(q, min(k, len(hits)))
+        return [
+            {"id": hits[i]["id"], "text": hits[i]["text"], "metadata": hits[i]["metadata"]}
+            for i in idxs[0]
+        ]
+
+    def history(self, scope, k=None) -> list:
+        raise NotImplementedError("semantic memory has no transcript; use conversational store")
+
+
+class PgVectorMemory(MemoryBackend):
+    """Shared/persistent vector memory over Postgres + pgvector"""
+
+    def __init__(self, name: str, endpoint: str = "", collection: str = "", embedder=None):
+        import psycopg
+        from psycopg import sql
+        from pgvector.psycopg import register_vector
+        self.name = name
+        self.endpoint = endpoint or "postgresql://127.0.0.1:5432/rosen"
+        self.collection = collection or "rosen_semantic"
+        self._embedder = embedder
+        self._table = sql.Identifier(self.collection)
+        self._conn = psycopg.connect(self.endpoint, autocommit=True)
+        self._conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+        register_vector(self._conn)
+        self._ready = False
+
+    def _embed(self, text: str) -> list:
+        if self._embedder is None:
+            raise RuntimeError(
+                "PgVectorMemory needs an embedder; bind one via Memory.spec.embedder")
+        return self._embedder.embed([text])[0]
+
+    def _ensure_table(self, dim: int):
+        # Lazily create the table on first use: the dimension comes from the
+        # embedder's vectors, so we never hardcode it.
+        if self._ready:
+            return
+        from psycopg import sql
+        self._conn.execute(sql.SQL(
+            "CREATE TABLE IF NOT EXISTS {table} ("
+            "id TEXT PRIMARY KEY, scope TEXT, text TEXT, "
+            "metadata JSONB, embedding vector({dim}))"
+        ).format(table=self._table, dim=sql.Literal(int(dim))))
+        self._ready = True
+
+    def remember(self, text, metadata=None, scope=None) -> str:
+        from psycopg import sql
+        from psycopg.types.json import Jsonb
+        from pgvector import Vector
+        vector = self._embed(text)
+        self._ensure_table(len(vector))
+        item_id = uuid.uuid4().hex
+        self._conn.execute(
+            sql.SQL("INSERT INTO {table} (id, scope, text, metadata, embedding) "
+                    "VALUES (%s, %s, %s, %s, %s)").format(table=self._table),
+            (item_id, scope, text, Jsonb(metadata or {}), Vector(vector)))
+        return item_id
+
+    def recall(self, query, k=5, scope=None) -> list:
+        from psycopg import sql
+        from pgvector import Vector
+        vector = self._embed(query)
+        self._ensure_table(len(vector))
+        where = sql.SQL("WHERE scope = %s") if scope is not None else sql.SQL("")
+        params = ([scope] if scope is not None else []) + [Vector(vector), k]
+        rows = self._conn.execute(
+            sql.SQL("SELECT id, text, metadata FROM {table} {where} "
+                    "ORDER BY embedding <=> %s LIMIT %s").format(table=self._table, where=where),
+            params).fetchall()
+        return [{"id": r[0], "text": r[1], "metadata": r[2]} for r in rows]
+
+    def history(self, scope, k=None) -> list:
+        raise NotImplementedError("semantic memory has no transcript; use conversational store")
+
+    def close(self):
+        self._conn.close()
+
+
+_BACKENDS = {
+    "mock": MockMemory,
+    "inmemory": MockMemory,
+    "redis": RedisMemory,
+    "faiss": FaissMemory,
+    "pgvector": PgVectorMemory
+}
+
+
+class Memory:
+
+    def __init__(self, name="", backend="", endpoint="", collection="", mode="", embedder=None):
+        self.name = name
+        self.backend = backend
+        self.endpoint = endpoint
+        self.collection = collection
+        self.mode = mode or "semantic"
+        built = build_embedder(embedder.get("engine", ""),
+                               embedder.get("model", "")) if embedder else None
+        adapter_cls = _BACKENDS.get(backend, MockMemory)
+        self._adapter = adapter_cls(name, endpoint, collection, embedder=built)
+
+    def remember(self, text, metadata=None, scope=None) -> str:
+        return self._adapter.remember(text, metadata, scope)
+
+    def recall(self, query, k=5, scope=None) -> list:
+        return self._adapter.recall(query, k, scope)
+
+    def history(self, scope, k=None) -> list:
+        return self._adapter.history(scope, k)
+
+
+class Memories:
+
+    def __init__(self, memories=None):
+        self._by_name = {m.name: m for m in (memories or [])}
+
+    def _by_mode(self, mode: str) -> Memory:
+        matches = [m for m in self._by_name.values() if m.mode == mode]
+        if not matches:
+            raise KeyError(f"no {mode} memory is bound to this agent")
+        if len(matches) > 1:
+            raise KeyError(f"multiple {mode} memory stores bound; address one by name")
+        return matches[0]
+
+    def has_mode(self, mode: str) -> bool:
+        return any(m.mode == mode for m in self._by_name.values())
+
+    def remember(self, text, metadata=None, scope=None) -> str:
+        mode = "conversational" if scope else "semantic"
+        return self._by_mode(mode).remember(text, metadata, scope)
+
+    def recall(self, query, k=5, scope=None) -> list:
+        return self._by_mode("semantic").recall(query, k, scope)
+
+    def history(self, scope, k=None) -> list:
+        return self._by_mode("conversational").history(scope, k)
+
+    def get(self, name: str) -> Memory:
+        try:
+            return self._by_name[name]
+        except KeyError:
+            raise KeyError(f"memory {name!r} not available to this agent")
+
+    def __getitem__(self, name: str) -> Memory:
+        return self.get(name)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._by_name
+
+    def __iter__(self):
+        return iter(self._by_name.values())
+
+    def __len__(self) -> int:
+        return len(self._by_name)
+
+    def name(self):
+        return list(self._by_name.keys())