nitpick-core 0.1.0__py3-none-any.whl

nitpick_core/__init__.py

@@ -0,0 +1,51 @@
+"""nitpick-core — the agent runtime for nitpick.
+
+Loads agent definitions from a minidb registry, dispatches them through
+Prefect, and serves the voice rooms so people can talk to them.
+
+    POST /tasks  {"handle": "assistant"}  ->  a session
+    WS   /ws                              ->  talk to it
+"""
+from __future__ import annotations
+
+from .builder import AgentSpec, build_agent, build_spec
+from .client import CoreClient, CoreError
+from .config import Config, load_config
+from .errors import (
+    AgentNotFound,
+    BuildError,
+    DispatchError,
+    ModelNotFound,
+    NitpickCoreError,
+    RegistryError,
+    SessionNotFound,
+)
+from .registry import Registry, get_registry, resolve_agent
+from .runtime import AgentRuntime, Session, SessionStatus, get_runtime
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AgentNotFound",
+    "AgentRuntime",
+    "AgentSpec",
+    "BuildError",
+    "Config",
+    "CoreClient",
+    "CoreError",
+    "DispatchError",
+    "ModelNotFound",
+    "NitpickCoreError",
+    "Registry",
+    "RegistryError",
+    "Session",
+    "SessionNotFound",
+    "SessionStatus",
+    "__version__",
+    "build_agent",
+    "build_spec",
+    "get_registry",
+    "get_runtime",
+    "load_config",
+    "resolve_agent",
+]

nitpick_core/__main__.py

@@ -0,0 +1,6 @@
+import sys
+
+from .cli import main
+
+if __name__ == "__main__":
+    sys.exit(main())

nitpick_core/api.py

@@ -0,0 +1,298 @@
+"""The nitpick-core HTTP API.
+
+Two things live on one port:
+
+* the **control API** — dispatch tasks, look at sessions, browse the agent
+  registry
+* the **signaling WebSocket** at ``/ws`` — the nitpick room protocol, so
+  somebody can talk to a running agent
+
+The browser client is served at ``/``.
+
+    POST /tasks           {"handle": "assistant", "room_id": "demo"}
+      -> 202 {"id": "ses-...", "status": "queued", ...}
+
+    GET  /tasks/{id}      the session
+    GET  /sessions        every session
+    GET  /agents          the agent registry
+    GET  /agents/{handle} one agent, resolved
+    GET  /health
+    WS   /ws              the signaling socket
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from typing import Any
+
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import HTMLResponse, JSONResponse
+from pydantic import BaseModel, Field
+
+from nitpick.signaling.server import SignalingServer
+
+from .config import Config, load_config
+from .errors import (
+    AgentNotFound,
+    BuildError,
+    DispatchError,
+    ModelNotFound,
+    RegistryError,
+    SessionNotFound,
+)
+from .flows import run_flow
+from .registry import get_registry
+from .runtime import AgentRuntime, get_runtime
+
+log = logging.getLogger(__name__)
+
+__all__ = ["create_app"]
+
+
+class TaskRequest(BaseModel):
+    """A request to run an agent."""
+
+    handle: str = Field(..., description="The agent's handle in the registry")
+    room_id: str | None = Field(None, description="The room to join. Defaults to the configured one.")
+    overrides: dict[str, Any] = Field(default_factory=dict, description="Overrides for the built agent.")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="Arbitrary data to record on the session.")
+    wait: bool = Field(False, description="Wait for the session to finish before replying.")
+
+
+def _encoder() -> dict[type, Any]:
+    """Encoders for the types the registry stores.
+
+    minidb documents carry an ``ObjectId`` and a ``created_at`` timestamp;
+    FastAPI knows about the timestamp but not about the object id.
+    """
+    encoders: dict[type, Any] = {}
+    try:
+        from minidb import ObjectId
+
+        encoders[ObjectId] = str
+    except ImportError:  # pragma: no cover - minidb is a hard dependency
+        pass
+    return encoders
+
+
+def create_app(
+    config: Config | None = None,
+    *,
+    runtime: AgentRuntime | None = None,
+) -> FastAPI:
+    """Build the application."""
+    settings = config or load_config()
+    agent_runtime = runtime or get_runtime(settings)
+
+    # Configure Prefect before anything else touches it, so its state
+    # lives where this service can find it.
+    from .flows import prefect_settings
+
+    prefect_settings(settings)
+
+    # The signaling server shares the runtime's room manager, so every
+    # session is visible on the same socket.
+    signaling = SignalingServer(agent_runtime.manager)
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    @asynccontextmanager
+    async def lifespan(_app: FastAPI):
+        await agent_runtime.start()
+        try:
+            yield
+        finally:
+            await agent_runtime.stop()
+
+    app = FastAPI(
+        title="nitpick-core", version="0.1.0", lifespan=lifespan
+    )
+    app.state.config = settings
+    app.state.runtime = agent_runtime
+    app.state.signaling = signaling
+
+    # ------------------------------------------------------------------
+    # Control API
+    # ------------------------------------------------------------------
+
+    @app.get("/health")
+    async def health() -> JSONResponse:
+        return JSONResponse(
+            {
+                "status": "ok",
+                "config": settings.to_dict(),
+                "sessions": len(agent_runtime.list_sessions()),
+            }
+        )
+
+    @app.post("/tasks", status_code=202)
+    async def create_task(payload: TaskRequest) -> JSONResponse:
+        """Dispatch a task. The agent runtime picks it up."""
+        try:
+            # The flow is synchronous and takes about a second while Prefect
+            # records the run, so it goes in a thread rather than blocking
+            # the event loop.
+            session = await asyncio.to_thread(
+                run_flow,
+                payload.handle,
+                room_id=payload.room_id,
+                overrides=payload.overrides or None,
+                metadata=payload.metadata or None,
+                runtime=agent_runtime,
+                config=settings,
+            )
+        except DispatchError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+        except (AgentNotFound, ModelNotFound) as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        except BuildError as exc:
+            raise HTTPException(status_code=422, detail=str(exc)) from exc
+        except RegistryError as exc:
+            raise HTTPException(status_code=503, detail=str(exc)) from exc
+
+        return JSONResponse(session, status_code=202)
+
+    @app.get("/tasks/{session_id}")
+    async def get_task(session_id: str) -> JSONResponse:
+        """One session."""
+        try:
+            session = agent_runtime.get_session(session_id)
+        except SessionNotFound as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        return JSONResponse(session.to_dict())
+
+    @app.get("/sessions")
+    async def list_sessions() -> JSONResponse:
+        """Every session."""
+        sessions = [session.to_dict() for session in agent_runtime.list_sessions()]
+        return JSONResponse({"count": len(sessions), "results": sessions})
+
+    @app.get("/sessions/{session_id}")
+    async def get_session(session_id: str) -> JSONResponse:
+        """One session."""
+        try:
+            session = agent_runtime.get_session(session_id)
+        except SessionNotFound as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        return JSONResponse(session.to_dict())
+
+    @app.get("/agents")
+    async def list_agents() -> JSONResponse:
+        """The agent registry."""
+        try:
+            agents = get_registry(settings).list_agents()
+        except RegistryError as exc:
+            raise HTTPException(status_code=503, detail=str(exc)) from exc
+        return JSONResponse(
+            jsonable_encoder(
+                {"count": len(agents), "results": agents}, custom_encoder=_encoder()
+            )
+        )
+
+    @app.get("/agents/{handle}")
+    async def get_agent(handle: str) -> JSONResponse:
+        """One agent, with the models it names resolved."""
+        try:
+            resolved = get_registry(settings).resolve_agent(handle)
+        except AgentNotFound as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        except ModelNotFound as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        except RegistryError as exc:
+            raise HTTPException(status_code=503, detail=str(exc)) from exc
+        return JSONResponse(jsonable_encoder(resolved, custom_encoder=_encoder()))
+
+    # ------------------------------------------------------------------
+    # The signaling socket and the browser client
+    # ------------------------------------------------------------------
+
+    @app.websocket("/ws")
+    async def websocket_endpoint(websocket: WebSocket) -> None:
+        await _serve_signaling(signaling, websocket)
+
+    @app.get("/", response_class=HTMLResponse)
+    async def index() -> HTMLResponse:
+        html = _load_web_client(settings)
+        if html is None:
+            return HTMLResponse(
+                "<h1>nitpick-core</h1><p>The browser client is not installed.</p>"
+            )
+        return HTMLResponse(html)
+
+    return app
+
+
+async def _serve_signaling(signaling: SignalingServer, websocket: WebSocket) -> None:
+    """Run nitpick's WebSocket protocol on this app's socket."""
+    import json
+
+    await websocket.accept()
+    participant = signaling.manager.register_participant(websocket)
+    await websocket.send_json(
+        {
+            "type": "connected",
+            "participant_id": participant.participant_id,
+            "stt": {"provider": "pcm", "transport": "audio_pcm"},
+            "audio": {"output_sample_rate": signaling.audio_output_sample_rate},
+        }
+    )
+    try:
+        while True:
+            raw = await websocket.receive_text()
+            try:
+                payload = json.loads(raw)
+            except json.JSONDecodeError:
+                await websocket.send_json(
+                    {
+                        "type": "error",
+                        "code": "invalid_json",
+                        "message": "Payload is not valid JSON",
+                    }
+                )
+                continue
+            if not isinstance(payload, dict):
+                await websocket.send_json(
+                    {
+                        "type": "error",
+                        "code": "invalid_payload",
+                        "message": "Payload must be a JSON object",
+                    }
+                )
+                continue
+            response = await signaling.handle_message(websocket, payload)
+            if response is not None:
+                await websocket.send_json(response)
+    except WebSocketDisconnect:
+        pass
+    finally:
+        signaling.manager.unregister_participant_by_id(participant.participant_id)
+
+
+def _load_web_client(config: Config) -> str | None:
+    """The browser client.
+
+    nitpick ships one; ``NITPICK_CORE_WEB_ROOT`` overrides it.
+    """
+    from pathlib import Path
+
+    if config.web_root:
+        path = Path(config.web_root)
+        if path.is_dir():
+            path = path / "index.html"
+        if path.exists():
+            return path.read_text(encoding="utf-8")
+        return None
+
+    try:
+        from nitpick.signaling import server as nitpick_server
+    except ImportError:  # pragma: no cover - nitpick is a hard dependency
+        return None
+
+    path = Path(nitpick_server.__file__).resolve().parent.parent / "web" / "index.html"
+    if path.exists():
+        return path.read_text(encoding="utf-8")
+    return None

nitpick_core/builder.py

@@ -0,0 +1,292 @@
+"""Turning registry documents into a runnable nitpick agent.
+
+The registry stores *what* an agent is — its prompt and the models it talks
+through — in a provider-neutral way. This module translates that into the
+provider-specific strings nitpick's ``Agent`` facade expects.
+
+The translation is driven by the model's ``handle``, which names the
+provider (``openai``, ``deepgram``, ``sarvam``, ...):
+
+============  ==========  ==========================================
+Registry      nitpick     How the value is derived
+============  ==========  ==========================================
+``llm``       ``llm``     ``openai:<model>``
+``stt``       ``stt``     ``openai`` / ``deepgram`` / ``sarvam``
+``tts``       ``tts``     ``openai:<voice>``
+============  ==========  ==========================================
+
+A model may override the derived value through its ``config`` — for
+example ``{"model": "gpt-4o-mini"}`` on an LLM, or ``{"voices": ["nova"]}``
+on a TTS model.
+"""
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+from nitpick import Agent
+
+from .errors import BuildError
+
+log = logging.getLogger(__name__)
+
+__all__ = ["AgentSpec", "build_agent", "build_spec", "provider_for"]
+
+
+#: Handles nitpick can talk to directly.
+SUPPORTED_PROVIDERS = {
+    "openai": "openai",
+    "deepgram": "deepgram",
+    "sarvam": "sarvam",
+    "ollama": "openai",  # OpenAI-compatible wire format
+    "vllm": "openai",
+    "together": "openai",
+    "groq": "openai",
+}
+
+#: A fallback for registries where a model's handle was left as its slug
+#: instead of naming the provider. The handle is the right way to say this;
+#: the table only rescues data that was seeded before the handle was
+#: introduced, or written by hand.
+KNOWN_MODEL_PROVIDERS = {
+    "gpt-4o": "openai",
+    "gpt-4o-mini": "openai",
+    "gpt-4.1": "openai",
+    "gpt-4-turbo": "openai",
+    "o1": "openai",
+    "o3": "openai",
+    "alloy": "openai",
+    "echo": "openai",
+    "fable": "openai",
+    "onyx": "openai",
+    "nova": "openai",
+    "whisper": "openai",
+    "deepgram": "deepgram",
+    "deepgram-nova": "deepgram",
+    "nova-3": "deepgram",
+    "saaras": "sarvam",
+    "saaras:v3": "sarvam",
+    "sarvam": "sarvam",
+}
+
+
+@dataclass
+class AgentSpec:
+    """Everything needed to build one agent, resolved from the registry."""
+
+    handle: str
+    name: str
+    prompt: str
+    stt: str
+    llm: str
+    tts: str
+    language: str | None = None
+    #: The base URL of the LLM endpoint, for OpenAI-compatible servers.
+    llm_base_url: str | None = None
+    room_id: str = "demo"
+    config: dict[str, Any] = field(default_factory=dict)
+    #: The raw documents, kept for the session record.
+    documents: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "handle": self.handle,
+            "name": self.name,
+            "prompt": self.prompt,
+            "stt": self.stt,
+            "llm": self.llm,
+            "tts": self.tts,
+            "language": self.language,
+            "llm_base_url": self.llm_base_url,
+            "room_id": self.room_id,
+            "config": self.config,
+        }
+
+
+def provider_for(model: dict[str, Any]) -> str:
+    """The nitpick provider name for a registry model.
+
+    The model's ``handle`` names the provider. A model whose provider
+    nitpick does not know about is only usable if it declares an
+    OpenAI-compatible ``base_url``.
+    """
+    handle = str(model.get("handle") or "").strip().lower()
+    config = model.get("config") or {}
+
+    if handle in SUPPORTED_PROVIDERS:
+        return SUPPORTED_PROVIDERS[handle]
+
+    # The handle was left as the model's own name — fall back to a lookup
+    # by name, so a registry written before the handle existed still works.
+    for key in ("slug", "name"):
+        candidate = str(model.get(key) or "").strip().lower()
+        if candidate in KNOWN_MODEL_PROVIDERS:
+            return KNOWN_MODEL_PROVIDERS[candidate]
+
+    if config.get("base_url"):
+        # Any OpenAI-compatible endpoint works.
+        return "openai"
+
+    raise BuildError(
+        f"the model {model.get('slug')!r} uses the provider {handle!r}, which "
+        "nitpick does not support. Supported providers are: "
+        + ", ".join(sorted(set(SUPPORTED_PROVIDERS.values())))
+        + ". Give the model a 'base_url' in its config if it speaks the "
+        "OpenAI wire format, or set its handle to the provider's name."
+    )
+
+
+def _model_id(model: dict[str, Any], fallback: str) -> str:
+    """The provider-specific model identifier."""
+    config = model.get("config") or {}
+    for key in ("model", "model_id", "id"):
+        value = config.get(key)
+        if isinstance(value, str) and value:
+            return value
+    # The slug is a better identifier than the display name.
+    slug = model.get("slug")
+    if isinstance(slug, str) and slug:
+        return slug
+    return fallback
+
+
+def _voice_for(model: dict[str, Any]) -> str:
+    """The voice a TTS model should speak with."""
+    config = model.get("config") or {}
+    for key in ("voice", "voices"):
+        value = config.get(key)
+        if isinstance(value, str) and value:
+            return value
+        if isinstance(value, list) and value:
+            first = value[0]
+            if isinstance(first, str) and first:
+                return first
+    slug = model.get("slug")
+    return slug if isinstance(slug, str) and slug else "alloy"
+
+
+def _stt_spec(model: dict[str, Any]) -> str:
+    """The ``stt`` argument for nitpick's Agent."""
+    provider = provider_for(model)
+    if provider == "openai":
+        return "openai"
+    if provider == "deepgram":
+        return "deepgram"
+    if provider == "sarvam":
+        return "sarvam"
+    raise BuildError(f"unsupported speech-to-text provider {provider!r}")
+
+
+def _llm_spec(model: dict[str, Any]) -> str:
+    """The ``llm`` argument for nitpick's Agent."""
+    provider = provider_for(model)
+    if provider != "openai":
+        raise BuildError(
+            f"unsupported language-model provider {provider!r} — nitpick speaks "
+            "the OpenAI wire format"
+        )
+    return f"openai:{_model_id(model, 'gpt-4o-mini')}"
+
+
+def _tts_spec(model: dict[str, Any]) -> str:
+    """The ``tts`` argument for nitpick's Agent."""
+    provider = provider_for(model)
+    if provider != "openai":
+        raise BuildError(f"unsupported text-to-speech provider {provider!r}")
+    return f"openai:{_voice_for(model)}"
+
+
+def build_spec(
+    resolved: dict[str, Any],
+    *,
+    room_id: str = "demo",
+    overrides: dict[str, Any] | None = None,
+) -> AgentSpec:
+    """Build the spec for one agent from its resolved documents."""
+    agent_document = resolved["agent"]
+    config = dict(agent_document.get("config") or {})
+    overrides = overrides or {}
+
+    stt_model = resolved.get("stt")
+    llm_model = resolved.get("llm")
+    tts_model = resolved.get("tts")
+
+    if stt_model is None:
+        raise BuildError(
+            f"the agent {agent_document.get('slug')!r} has no speech-to-text model"
+        )
+    if llm_model is None:
+        raise BuildError(
+            f"the agent {agent_document.get('slug')!r} has no language model"
+        )
+    if tts_model is None:
+        raise BuildError(
+            f"the agent {agent_document.get('slug')!r} has no text-to-speech model"
+        )
+
+    stt = overrides.get("stt") or _stt_spec(stt_model)
+    llm = overrides.get("llm") or _llm_spec(llm_model)
+    tts = overrides.get("tts") or _tts_spec(tts_model)
+
+    language = (
+        overrides.get("language")
+        or config.get("language")
+        or (stt_model.get("config") or {}).get("language")
+    )
+
+    # A model may point at an OpenAI-compatible endpoint — a self-hosted
+    # model, or a gateway. That is how a provider nitpick does not know
+    # about can still be used.
+    llm_base_url = (
+        overrides.get("llm_base_url")
+        or config.get("base_url")
+        or (llm_model.get("config") or {}).get("base_url")
+    )
+
+    return AgentSpec(
+        handle=str(agent_document.get("handle") or agent_document.get("slug") or ""),
+        name=str(agent_document.get("name") or agent_document.get("slug") or "Agent"),
+        prompt=str(agent_document.get("prompt") or ""),
+        stt=stt,
+        llm=llm,
+        tts=tts,
+        language=language if isinstance(language, str) else None,
+        llm_base_url=llm_base_url if isinstance(llm_base_url, str) else None,
+        room_id=room_id,
+        config=config,
+        documents=resolved,
+    )
+
+
+def build_agent(
+    spec: AgentSpec,
+    *,
+    manager: Any = None,
+    session_id: str | None = None,
+) -> Agent:
+    """Build the runnable nitpick agent for a spec."""
+    kwargs: dict[str, Any] = {
+        "name": spec.name,
+        "prompt": spec.prompt or None,
+        "stt": spec.stt,
+        "llm": spec.llm,
+        "tts": spec.tts,
+    }
+    if spec.language:
+        kwargs["language"] = spec.language
+    if spec.llm_base_url:
+        kwargs["llm_base_url"] = spec.llm_base_url
+    if manager is not None:
+        kwargs["manager"] = manager
+    if session_id:
+        kwargs["session_id"] = session_id
+
+    # The agent's own config may carry pipeline settings.
+    config = spec.config or {}
+    if "vad" in config:
+        kwargs["vad"] = config["vad"]
+    if "allow_interruptions" in config:
+        kwargs["allow_interruptions"] = bool(config["allow_interruptions"])
+
+    return Agent(**kwargs)