pinecall-runtime 0.1.6__py3-none-any.whl

pinecall/postgres/pool.py

@@ -0,0 +1,190 @@
+"""The pool an instance opens on its database, and the names a schema and a DSN go by."""
+
+import asyncio
+import contextvars
+import re
+import weakref
+from collections.abc import AsyncGenerator, Coroutine, Generator
+from contextlib import asynccontextmanager, contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+
+import psycopg
+from psycopg.conninfo import conninfo_to_dict
+from psycopg.rows import DictRow, dict_row
+from psycopg_pool import AsyncConnectionPool, PoolTimeout
+
+from pinecall.domain.errors import DeclarationRefused, StoreUnreachable
+
+type Connection = psycopg.AsyncConnection[DictRow]
+
+
+type Pool = AsyncConnectionPool[Connection]
+
+
+DEFAULT_SCHEMA = "public"
+
+
+# A schema name also travels in the pool's libpq options string, where quoting does not reach.
+_A_SCHEMA_NAME = re.compile(r"^[a-z_][a-z0-9_]*$")
+
+
+CONNECT_TIMEOUT_S = 10
+
+
+# What a gateway holds open to Postgres unless PINECALL_DB_POOL says otherwise: its writer's and
+# its doors' together.
+POOL_SIZE = 10
+
+
+# SET LOCAL: they end with the transaction, and the connection goes back to the pool as it came.
+UNBOUNDED = (
+    "SET LOCAL statement_timeout = 0",
+    "SET LOCAL idle_in_transaction_session_timeout = 0",
+)
+
+
+# Row-level security (migration 0096): a connection that names an org sees and writes that org's
+# rows alone; one that names none is the box's, as the operator's doors and the loops are. The
+# gateway names the org of a tenant's key for the whole request; a connection taken then says it.
+SCOPED_ORG: ContextVar[str | None] = ContextVar("scoped_org", default=None)
+
+
+SCOPE = "select set_config('pinecall.org', %s, false)"
+
+
+@dataclass(frozen=True)
+class Timeouts:
+    """How long a pool's statement may run, its transaction sit idle, and a request wait for it."""
+
+    # A door's statement past this is cancelled; what is long on purpose runs `unbounded`.
+    statement_ms: int = 30_000
+    # A transaction its process stopped talking in is ended, and its locks let go.
+    idle_in_transaction_ms: int = 60_000
+    # A request that finds the pool full this long is answered 503, which a worker outlasts.
+    wait_s: float = 2.0
+
+
+TIMEOUTS = Timeouts()
+
+
+# What each connection is known to say, so one that already says the org is asked nothing: a
+# request's doors share a few connections, and the box's loops never name one. A connection is
+# written here only once the server answered; absent means unknown, and unknown is asked.
+_SAYS: weakref.WeakKeyDictionary[Connection, str] = weakref.WeakKeyDictionary()
+
+
+async def open_pool(
+    dsn: str,
+    *,
+    schema: str = DEFAULT_SCHEMA,
+    max_size: int = POOL_SIZE,
+    min_size: int = 1,
+    timeouts: Timeouts = TIMEOUTS,
+) -> Pool:
+    """Open a pool on the schema, rows as dicts; raise StoreUnreachable when nothing answers."""
+    options = (
+        f"-c search_path={','.join(schemas_of(schema))}"
+        f" -c statement_timeout={timeouts.statement_ms}"
+        f" -c idle_in_transaction_session_timeout={timeouts.idle_in_transaction_ms}"
+    )
+    # A DSN's own options (a role to act as, say) are kept, and the pool's follow them.
+    given = conninfo_to_dict(dsn).get("options")
+    pool: Pool = AsyncConnectionPool(
+        dsn,
+        connection_class=psycopg.AsyncConnection[DictRow],
+        # A statement alone is one round trip; what belongs together opens connection.transaction().
+        kwargs={
+            "autocommit": True,
+            "row_factory": dict_row,
+            "options": options if not given else f"{given} {options}",
+        },
+        min_size=min_size,
+        max_size=max_size,
+        timeout=timeouts.wait_s,
+        open=False,
+        check=_scoped,
+    )
+    try:
+        await pool.open(wait=True, timeout=CONNECT_TIMEOUT_S)
+    except (PoolTimeout, psycopg.OperationalError) as refused:
+        await pool.close()
+        raise StoreUnreachable(f"{database_named(dsn)}: {refused}") from refused
+    return pool
+
+
+# Erasure, an export, the nightly retention and a knowledge base's index are long on purpose.
+@asynccontextmanager
+async def unbounded(pool: Pool) -> AsyncGenerator[Connection]:
+    """A connection in a transaction that no statement or idle timeout of the pool ends."""
+    async with pool.connection() as connection, connection.transaction():
+        for statement in UNBOUNDED:
+            await connection.execute(statement)
+        yield connection
+
+
+def scope_to(org: str) -> None:
+    """Name the org every connection taken from here on in this request acts for."""
+    SCOPED_ORG.set(org)
+
+
+# What a tenant's request reads across orgs on purpose (is this address in another org, does
+# another org hold this number) is read box-wide, in a block that says so.
+@contextmanager
+def box_wide() -> Generator[None]:
+    """Take connections as the box's within the block, whatever the request's org."""
+    token = SCOPED_ORG.set(None)
+    try:
+        yield
+    finally:
+        SCOPED_ORG.reset(token)
+
+
+# A task of the process (the log's writer, the relay, a cache's listener) may start inside a
+# tenant's request: it starts box-wide, or it would act for that org for the rest of its life.
+def box_task[T](work: Coroutine[object, object, T]) -> asyncio.Task[T]:
+    """Start the work as a task of the box's, whatever org the request starting it acts for."""
+    context = contextvars.copy_context()
+    context.run(SCOPED_ORG.set, None)
+    return asyncio.create_task(work, context=context)
+
+
+async def connect(dsn: str) -> Connection:
+    """Open one autocommit connection, rows as dicts; the caller closes it."""
+    return await psycopg.AsyncConnection[DictRow].connect(
+        dsn, autocommit=True, row_factory=dict_row, connect_timeout=CONNECT_TIMEOUT_S
+    )
+
+
+# Extension types and operators (halfvec, <=>, bm25's <@>) live in public, so another schema
+# keeps public on its search path.
+def schemas_of(schema: str) -> tuple[str, ...]:
+    """Return the search path for a schema: the schema, then public."""
+    name = check_schema_name(schema)
+    return (name,) if name == DEFAULT_SCHEMA else (name, DEFAULT_SCHEMA)
+
+
+def check_schema_name(schema: str) -> str:
+    """Return the schema name, raising DeclarationRefused unless it is a plain lowercase word."""
+    if not _A_SCHEMA_NAME.match(schema):
+        raise DeclarationRefused(f"a schema name is a lowercase word, not {schema!r}")
+    return schema
+
+
+# Every message that names the database goes through this: a DSN carries the password.
+def database_named(dsn: str) -> str:
+    """Return `host/database` from the DSN, with no credentials."""
+    parts = conninfo_to_dict(dsn)
+    return f"{parts.get('host', 'localhost')}/{parts.get('dbname', '')}"
+
+
+# The memory is dropped before the set and written after it: a set cut between the server's
+# answer and ours (a request cancelled) leaves a connection the server has moved and we have
+# forgotten, which the next checkout asks again, never one we remember wrong.
+async def _scoped(connection: Connection) -> None:
+    wanted = SCOPED_ORG.get() or ""
+    if _SAYS.get(connection) == wanted:
+        return
+    _SAYS.pop(connection, None)
+    await connection.execute(SCOPE, (wanted,))
+    _SAYS[connection] = wanted

pinecall/process/_objects.py

@@ -0,0 +1,97 @@
+"""The box's object store, spoken in S3: any endpoint that takes AWS's version 4 signature."""
+
+import hashlib
+import hmac
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from urllib.parse import quote, urlsplit
+
+import httpx
+
+from pinecall.process.settings import Settings
+
+ALGORITHM = "AWS4-HMAC-SHA256"
+
+# What S3 is told a request without a body hashes to.
+NO_BODY = hashlib.sha256(b"").hexdigest()
+
+
+@dataclass(frozen=True)
+class Signing:
+    """What a request is signed with: the region the store names, and the key's id and secret."""
+
+    region: str
+    access_key_id: str
+    secret_access_key: str
+
+
+@dataclass(frozen=True)
+class ObjectStore:
+    """An S3-compatible endpoint and the key the box reads and writes it with."""
+
+    endpoint: str
+    signing: Signing
+    http: httpx.AsyncClient
+
+    # Path-style: the one address every S3-compatible store answers; a bucket as a host name is
+    # AWS's and a few others' alone.
+    def url_of(self, bucket: str, name: str) -> str:
+        """The object's address."""
+        return f"{self.endpoint.rstrip('/')}/{bucket}/{quote(name, safe='/')}"
+
+    def signed(
+        self, method: str, url: str, headers: Mapping[str, str], payload_sha256: str = NO_BODY
+    ) -> dict[str, str]:
+        """The headers to send, signed now: the ones given, the date, the hash and the proof."""
+        return signature(method, url, headers, self.signing, (payload_sha256, datetime.now(UTC)))
+
+
+def object_store_of(settings: Settings, http: httpx.AsyncClient) -> ObjectStore | None:
+    """The store the settings name whole; None when any of its four is missing."""
+    endpoint, region = settings.s3_endpoint, settings.s3_region
+    key, secret = settings.s3_access_key_id, settings.s3_secret_access_key
+    if endpoint is None or region is None or key is None or secret is None:
+        return None
+    return ObjectStore(endpoint, Signing(region, key, secret), http)
+
+
+# https://docs.aws.amazon.com/AmazonS3/latest/API/sig-v4-header-based-auth.html. The store's
+# requests name no query, so the canonical query is the empty line.
+def signature(
+    method: str,
+    url: str,
+    headers: Mapping[str, str],
+    signing: Signing,
+    payload: tuple[str, datetime],
+) -> dict[str, str]:
+    """The headers with the date, the payload's hash and the Authorization that signs them all."""
+    payload_sha256, now = payload
+    stamp = now.strftime("%Y%m%dT%H%M%SZ")
+    parts = urlsplit(url)
+    signed = {
+        **{name.lower(): value.strip() for name, value in headers.items()},
+        "host": parts.netloc,
+        "x-amz-content-sha256": payload_sha256,
+        "x-amz-date": stamp,
+    }
+    names = sorted(signed)
+    canonical = "\n".join(
+        (
+            method,
+            parts.path or "/",
+            "",
+            "".join(f"{name}:{signed[name]}\n" for name in names),
+            ";".join(names),
+            payload_sha256,
+        )
+    )
+    scope = f"{stamp[:8]}/{signing.region}/s3/aws4_request"
+    to_sign = "\n".join((ALGORITHM, stamp, scope, hashlib.sha256(canonical.encode()).hexdigest()))
+    key = f"AWS4{signing.secret_access_key}".encode()
+    for step in (stamp[:8], signing.region, "s3", "aws4_request"):
+        key = hmac.new(key, step.encode(), hashlib.sha256).digest()
+    proof = hmac.new(key, to_sign.encode(), hashlib.sha256).hexdigest()
+    credential = f"Credential={signing.access_key_id}/{scope}"
+    authorization = f"{ALGORITHM} {credential}, SignedHeaders={';'.join(names)}, Signature={proof}"
+    return {**signed, "authorization": authorization}

pinecall/process/box_settings.py

@@ -0,0 +1,23 @@
+"""The box_settings rows: what the operator configured for this box, one JSON value per name."""
+
+from psycopg.types.json import Jsonb
+
+from pinecall.domain.names import JsonObject
+from pinecall.postgres.pool import Connection
+
+READ = "SELECT value FROM box_settings WHERE name = %(name)s"
+WRITE = """
+INSERT INTO box_settings (name, value) VALUES (%(name)s, %(value)s)
+ON CONFLICT (name) DO UPDATE SET value = excluded.value, set_at = now()
+"""
+
+
+async def read(connection: Connection, name: str) -> JsonObject | None:
+    """The row's value, or None when the operator never set it."""
+    row = await (await connection.execute(READ, {"name": name})).fetchone()
+    return None if row is None else row["value"]
+
+
+async def write(connection: Connection, name: str, value: JsonObject) -> None:
+    """Set the row's value, whole."""
+    await connection.execute(WRITE, {"name": name, "value": Jsonb(value)})

pinecall/process/connections.py

@@ -0,0 +1,124 @@
+"""What one process opens at start and closes at the end: the database, the vault, HTTP, LiveKit."""
+
+from collections.abc import AsyncGenerator, Mapping
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+
+import httpx
+from cryptography.fernet import Fernet, MultiFernet
+from livekit import api
+
+from pinecall.domain.errors import NotAvailable, SettingsRefused
+from pinecall.domain.names import PRODUCTION, SANDBOX, Env
+from pinecall.postgres.pool import Pool, open_pool
+from pinecall.process.settings import Settings
+from pinecall.process.signal import LocalSignal, Signal, opened_signal
+
+VARIABLE = "PINECALL_VAULT_KEY"
+
+
+# The log's writer holds a connection per lane (log/_writer.py LANES), open for the process's
+# life; the doors share the rest of PINECALL_DB_POOL.
+WRITING = 2
+
+
+NO_LIVEKIT = "LIVEKIT_API_KEY and LIVEKIT_API_SECRET are unset: this process cannot reach the SFU"
+
+
+@dataclass(frozen=True)
+class Connections:
+    """The settings a process was given, and the five things it holds open on them."""
+
+    settings: Settings
+    # The doors'.
+    pool: Pool
+    # The log writer's own: never waits for a door, and a door never waits for it.
+    writing: Pool
+    vault: MultiFernet
+    # The carriers', Meta's and the identity providers' HTTP, one pool per process.
+    http: httpx.AsyncClient
+    # Each world's LiveKit server API, one client per LiveKit: both worlds hold the one client
+    # when the sandbox has no LiveKit of its own.
+    servers: Mapping[Env, api.LiveKitAPI]
+    # What the gateways tell each other; a process given no Redis keeps it to itself.
+    signal: Signal = field(default_factory=LocalSignal)
+
+
+UNSET = (
+    f"{VARIABLE} is unset: every vendor key, mailbox and sign-in secret is sealed under it, "
+    "so the gateway does not start without it. `Fernet.generate_key()` makes one"
+)
+
+
+NOT_A_KEY = (
+    f"{VARIABLE} holds something that is not a Fernet key: a comma-separated list of "
+    "32-byte url-safe base64 keys, the newest first"
+)
+
+
+@asynccontextmanager
+async def opened(settings: Settings) -> AsyncGenerator[Connections]:
+    """Open everything, the vault first: a box without its key starts nothing; closed in reverse."""
+    sealed = vault_of(settings.vault_key)
+    servers = servers_of(settings)
+    pool = await open_pool(settings.database_url, max_size=settings.db_pool - WRITING)
+    try:
+        writing = await open_pool(settings.database_url, max_size=WRITING, min_size=WRITING)
+    except BaseException:
+        await pool.close()
+        raise
+    try:
+        async with httpx.AsyncClient() as http, opened_signal(settings) as signal:
+            yield Connections(
+                settings=settings,
+                pool=pool,
+                writing=writing,
+                vault=sealed,
+                http=http,
+                servers=servers,
+                signal=signal,
+            )
+    finally:
+        await closed(servers)
+        await writing.close()
+        await pool.close()
+
+
+# livekit-api opens an HTTP session per client: one per LiveKit and process, closed with it. Both
+# LiveKits take the box's one key pair, so a token the gateway signs opens a room on either.
+def servers_of(settings: Settings) -> dict[Env, api.LiveKitAPI]:
+    """Each world's LiveKit server API, on the box's key pair: one client where they share one."""
+    if not settings.livekit_api_key or not settings.livekit_api_secret:
+        raise NotAvailable(NO_LIVEKIT)
+    key, secret = settings.livekit_api_key, settings.livekit_api_secret
+    production = api.LiveKitAPI(settings.livekit_url_of(PRODUCTION), key, secret)
+    if settings.livekit_url_of(SANDBOX) == settings.livekit_url_of(PRODUCTION):
+        return {PRODUCTION: production, SANDBOX: production}
+    return {
+        PRODUCTION: production,
+        SANDBOX: api.LiveKitAPI(settings.livekit_url_of(SANDBOX), key, secret),
+    }
+
+
+async def closed(servers: Mapping[Env, api.LiveKitAPI]) -> None:
+    """Close each LiveKit's client once, however many worlds hold it."""
+    for server in dict.fromkeys(servers.values()):
+        await server.aclose()
+
+
+def vault_of(keys: str | None) -> MultiFernet:
+    """The vault built from the setting: the first key seals, any key listed opens."""
+    return MultiFernet(keyring_of(keys))
+
+
+def keyring_of(keys: str | None) -> list[Fernet]:
+    """Every key the setting lists, in its order: the first is the one that seals."""
+    if keys is None:
+        raise SettingsRefused(UNSET)
+    listed = [item.strip() for item in keys.split(",") if item.strip()]
+    if not listed:
+        raise SettingsRefused(NOT_A_KEY)
+    try:
+        return [Fernet(listed_one) for listed_one in listed]
+    except ValueError:
+        raise SettingsRefused(NOT_A_KEY) from None

pinecall/process/metrics.py

@@ -0,0 +1,127 @@
+"""What a process counts as it works, in memory, and the Prometheus text it is read as."""
+
+from bisect import bisect_left
+from collections import Counter, deque
+from collections.abc import Iterable, Mapping
+from typing import Literal
+
+# Seconds: an append waits on one Postgres transaction, so the buckets are a transaction's.
+APPEND_BOUNDS_S = (0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0)
+
+# A stage's fallbacks are ordered by the last two minutes of each vendor: judged once five calls
+# were handed it as a stage's first, and over the line once half of those calls saw it fail.
+VENDOR_WINDOW_S = 120.0
+CALLS_TO_JUDGE = 5
+FAILED_SHARE = 0.5
+# A busy window is thousands of calls: past this many, a vendor's oldest are forgotten early.
+MOST_KEPT = 10_000
+
+
+type Rows = Iterable[tuple[Mapping[str, str], float]]
+
+type Kind = Literal["gauge", "counter"]
+
+
+class Histogram:
+    """Observations counted into fixed buckets, with their sum, as Prometheus reads one."""
+
+    def __init__(self, bounds: tuple[float, ...]) -> None:
+        """Every bucket empty."""
+        self.bounds = bounds
+        self.counts = [0] * (len(bounds) + 1)
+        self.total = 0.0
+
+    def observe(self, value: float) -> None:
+        """Count one observation in the first bucket whose bound holds it."""
+        self.counts[bisect_left(self.bounds, value)] += 1
+        self.total += value
+
+
+# Called on the append path: a bisect and three additions, no lock (one event loop).
+class Counters:
+    """The gateway's counts since it started: appends, their time, and the errors workers wrote."""
+
+    def __init__(self) -> None:
+        """Nothing counted yet."""
+        self.append_seconds = Histogram(APPEND_BOUNDS_S)
+        self.appended = 0
+        # By the entry's code and the vendor whose plugin failed, "" when none did.
+        self.errors: Counter[tuple[str, str]] = Counter()
+        # By vendor: when a call was handed it first for a stage, and when a call saw it fail.
+        self.handed: dict[str, deque[float]] = {}
+        self.failures: dict[str, deque[tuple[float, str]]] = {}
+        # By org: how many times its usual day today's spend is, while it is unusual.
+        self.unusual: dict[str, float] = {}
+
+    def appended_in(self, seconds: float, entries: int) -> None:
+        """One append at the door: how long it took and how many entries it wrote."""
+        self.append_seconds.observe(seconds)
+        self.appended += entries
+
+    def failed(self, code: str, vendor: str) -> None:
+        """One error entry a worker wrote."""
+        self.errors[(code, vendor)] += 1
+
+    def failed_on(self, vendor: str, call: str, at: float) -> None:
+        """A call saw the vendor fail: an error naming it, or its stage's switch away from it."""
+        self.failures.setdefault(vendor, deque(maxlen=MOST_KEPT)).append((at, call))
+
+    def handed_out(self, vendors: Iterable[str], at: float) -> None:
+        """A call's stages were resolved with these vendors first, each counted once."""
+        for vendor in set(vendors):
+            self.handed.setdefault(vendor, deque(maxlen=MOST_KEPT)).append(at)
+
+    def spending(self, org: str, multiple: float | None) -> None:
+        """What the last seal found of the org's spend: unusual by this much, or None, usual."""
+        if multiple is None:
+            self.unusual.pop(org, None)
+        else:
+            self.unusual[org] = multiple
+
+    def failing(self, at: float) -> frozenset[str]:
+        """The vendors over their error line in the window up to now."""
+        since = at - VENDOR_WINDOW_S
+        for failures in self.failures.values():
+            while failures and failures[0][0] < since:
+                failures.popleft()
+        over: set[str] = set()
+        for vendor, handed in self.handed.items():
+            while handed and handed[0] < since:
+                handed.popleft()
+            calls = len({call for _, call in self.failures.get(vendor, ())})
+            if len(handed) >= CALLS_TO_JUDGE and calls >= FAILED_SHARE * len(handed):
+                over.add(vendor)
+        return frozenset(over)
+
+
+def family(name: str, what: str, kind: Kind, rows: Rows) -> str:
+    """One gauge or counter family, a line per label set, with its help and type lines."""
+    lines = [f"# HELP {name} {what}", f"# TYPE {name} {kind}"]
+    lines += [f"{name}{_labelled(labels)} {_number(value)}" for labels, value in rows]
+    return "\n".join(lines) + "\n"
+
+
+def histogram(name: str, what: str, observed: Histogram) -> str:
+    """One histogram: its cumulative buckets, the sum and the count."""
+    lines = [f"# HELP {name} {what}", f"# TYPE {name} histogram"]
+    running = 0
+    for bound, count in zip((*observed.bounds, float("inf")), observed.counts, strict=True):
+        running += count
+        lines.append(f'{name}_bucket{{le="{_number(bound)}"}} {running}')
+    lines += [f"{name}_sum {_number(observed.total)}", f"{name}_count {running}"]
+    return "\n".join(lines) + "\n"
+
+
+# The text format's three escapes in a label value: backslash, quote, newline.
+def _labelled(labels: Mapping[str, str]) -> str:
+    if not labels:
+        return ""
+    pairs = (
+        f'{key}="{value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")}"'
+        for key, value in labels.items()
+    )
+    return "{" + ",".join(pairs) + "}"
+
+
+def _number(value: float) -> str:
+    return "+Inf" if value == float("inf") else repr(float(value))