pinecall-runtime 0.1.6__py3-none-any.whl

pinecall/tenancy/export.py

@@ -0,0 +1,167 @@
+"""An org's world as JSON Lines: calls, memories, settings, words, documents, consent, drift."""
+
+import json
+import time
+from collections.abc import AsyncIterator
+
+from pinecall.domain.names import Env
+from pinecall.postgres.pool import Pool, unbounded
+
+# One page of calls at a time, each with its whole log, so an org of any size streams in bounded
+# memory. Every read is unbounded: an org's list of anything has no size. Postgres writes the
+# JSON: timestamps, jsonb and nulls come out as they are stored.
+CALLS = """
+SELECT jsonb_build_object(
+    'kind', 'call', 'call', head.call, 'agent', head.agent, 'holder', head.holder,
+    'started_at', head.started_at, 'sealed', head.sealed,
+    'facts', to_jsonb(facts) - 'call',
+    'entries', (
+        SELECT coalesce(jsonb_agg(jsonb_build_object(
+            'seq', entry.seq, 'ts', entry.ts, 'type', entry.type, 'data', entry.data
+        ) ORDER BY entry.seq), '[]'::jsonb)
+        FROM call_log entry WHERE entry.log = head.log
+    )
+)::text AS line, coalesce(head.started_at, 0) AS at, head.call
+FROM call_log_head head
+LEFT JOIN call_facts facts ON facts.call = head.call
+WHERE head.org = %(org)s AND head.env = %(env)s AND head.call IS NOT NULL
+  AND (coalesce(head.started_at, 0), head.call) > (%(at)s, %(call)s)
+ORDER BY coalesce(head.started_at, 0), head.call
+LIMIT %(page)s
+"""
+
+# The embedding is the box's to compute again; the words are the org's.
+MEMORIES = """
+SELECT jsonb_build_object(
+    'kind', 'memory', 'contact', contact, 'holder', holder, 'text', text, 'category', category,
+    'valid_from', valid_from, 'invalidated_at', invalidated_at, 'source_call', source_call
+)::text AS line
+FROM contact_memories WHERE org = %(org)s AND env = %(env)s
+ORDER BY contact, valid_from
+"""
+
+SETTINGS = """
+SELECT jsonb_build_object(
+    'kind', 'agent_config', 'agent', agent, 'holder', holder, 'version', version,
+    'config', config, 'author', author, 'note', note, 'set_at', set_at
+)::text AS line
+FROM agent_config WHERE org = %(org)s AND env = %(env)s
+ORDER BY agent, holder, version
+"""
+
+# Every canary set and cleared, as the settings are: nothing is ever updated.
+CANARIES = """
+SELECT jsonb_build_object(
+    'kind', 'canary', 'agent', agent, 'holder', holder, 'version', version, 'share', share,
+    'author', author, 'note', note, 'set_at', set_at
+)::text AS line
+FROM agent_canaries WHERE org = %(org)s AND env = %(env)s
+ORDER BY agent, holder, set_at, id
+"""
+
+# The dataset is the org's in both worlds: every export carries it whole, with the world its
+# call ran in.
+CASES = """
+SELECT jsonb_build_object(
+    'kind', 'eval_case', 'id', id, 'agent', agent, 'name', name, 'golden', golden,
+    'source_call', source_call, 'source_env', source_env, 'held_out', held_out,
+    'author', author, 'created_at', created_at
+)::text AS line
+FROM eval_cases WHERE org = %(org)s
+ORDER BY agent, name
+"""
+
+# Each block of prompt the org's calls were told, once per text, the same in both worlds.
+PROMPTS = """
+SELECT jsonb_build_object(
+    'kind', 'prompt', 'hash', hash, 'text', text, 'first_used_at', first_used_at
+)::text AS line
+FROM prompts WHERE org = %(org)s
+ORDER BY first_used_at, hash
+"""
+
+WORDS = """
+SELECT jsonb_build_object(
+    'kind', 'lexicon', 'agent', agent, 'holder', holder, 'version', version,
+    'said', said, 'heard', heard, 'author', author, 'note', note, 'set_at', set_at
+)::text AS line
+FROM lexicon WHERE org = %(org)s AND env = %(env)s
+ORDER BY agent, holder, version
+"""
+
+DOCUMENTS = """
+SELECT jsonb_build_object(
+    'kind', 'knowledge_file', 'base', base, 'holder', holder, 'path', path, 'text', text,
+    'pushed_at', pushed_at
+)::text AS line
+FROM knowledge_files WHERE org = %(org)s AND env = %(env)s
+ORDER BY base, path
+"""
+
+CONSENTS = """
+SELECT jsonb_build_object(
+    'kind', 'consent', 'number', number, 'consent', kind, 'source', source, 'text', text,
+    'evidence', evidence, 'given_by', given_by, 'call', call, 'given_at', given_at
+)::text AS line
+FROM contact_consents WHERE org = %(org)s AND env = %(env)s
+ORDER BY number, given_at, id
+"""
+
+# What the seal counted of each day, by agent and version: numbers and the judges' names.
+STAGE_DAYS = """
+SELECT jsonb_build_object(
+    'kind', 'stage_day', 'holder', holder, 'agent', agent, 'day', day,
+    'config_version', config_version, 'stage', stage, 'vendor', vendor, 'model', model,
+    'turns', turns, 'buckets', buckets, 'confidence_sum', confidence_sum,
+    'confidence_turns', confidence_turns
+)::text AS line
+FROM stage_days WHERE org = %(org)s AND env = %(env)s
+ORDER BY day, agent, config_version, stage, vendor, model
+"""
+
+JUDGE_DAYS = """
+SELECT jsonb_build_object(
+    'kind', 'judge_day', 'holder', holder, 'agent', agent, 'day', day,
+    'config_version', config_version, 'judge', judge, 'criteria', criteria, 'held', held,
+    'broken', broken
+)::text AS line
+FROM judge_days WHERE org = %(org)s AND env = %(env)s
+ORDER BY day, agent, config_version, judge, criteria
+"""
+
+A_PAGE_OF_CALLS = 100
+
+
+# What follows the calls, in this order.
+AFTER_THE_CALLS = (
+    MEMORIES,
+    SETTINGS,
+    CANARIES,
+    WORDS,
+    DOCUMENTS,
+    CONSENTS,
+    CASES,
+    PROMPTS,
+    STAGE_DAYS,
+    JUDGE_DAYS,
+)
+
+
+async def lines(pool: Pool, org: str, env: Env) -> AsyncIterator[str]:
+    """Every line: the header, calls, memories, settings, words, documents, consent, drift."""
+    yield json.dumps({"kind": "export", "org": org, "env": env, "exported_at": time.time()})
+    at, call = -1.0, ""
+    while True:
+        async with unbounded(pool) as connection:
+            params = {"org": org, "env": env, "at": at, "call": call, "page": A_PAGE_OF_CALLS}
+            rows = await (await connection.execute(CALLS, params)).fetchall()
+        for row in rows:
+            yield str(row["line"])
+        if len(rows) < A_PAGE_OF_CALLS:
+            break
+        at, call = float(rows[-1]["at"]), str(rows[-1]["call"])
+    for query in AFTER_THE_CALLS:
+        async with unbounded(pool) as connection:
+            rows = await (await connection.execute(query, {"org": org, "env": env})).fetchall()
+        for row in rows:
+            yield str(row["line"])

pinecall/tenancy/hosted_running.py

@@ -0,0 +1,166 @@
+"""A hosted app while it runs: stopped and started, the logs asked for, the time it served."""
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import UTC, date, datetime
+from typing import LiteralString
+
+from pinecall.domain.errors import NotFound
+from pinecall.domain.names import Env
+from pinecall.postgres.pool import Pool
+from pinecall.tenancy.hosting import NOT_HOSTED, HostedApp
+
+# A gap between two beats longer than this is a runner that was away: it counts this much.
+LONGEST_BEAT_S = 30.0
+
+
+# What a person reads of a process's output: its last lines, and no more than this.
+LONGEST_LOGS = 64 * 1024
+
+
+STOP = """
+UPDATE hosted_apps SET stopped_at = now(), stopped_by = %(by)s
+WHERE org = %(org)s AND env = %(env)s AND name = %(name)s RETURNING name
+"""
+
+
+START = """
+UPDATE hosted_apps SET stopped_at = NULL, stopped_by = ''
+WHERE org = %(org)s AND env = %(env)s AND name = %(name)s RETURNING name
+"""
+
+
+ASK_LOGS = """
+UPDATE hosted_apps SET logs_asked_at = now()
+WHERE org = %(org)s AND env = %(env)s AND name = %(name)s
+RETURNING logs, logs_host, logs_at
+"""
+
+
+KEEP_LOGS = """
+UPDATE hosted_apps SET logs = %(lines)s, logs_host = %(host)s, logs_at = now()
+WHERE org = %(org)s AND env = %(env)s AND name = %(name)s
+"""
+
+
+# One statement for a world's beat. Every app that was being counted, or serves now, is locked and
+# read once: the time since its last count goes on today's row, capped, and it is counted from now
+# if it serves and from nothing if it does not. Two runners beating at once wait on the lock and
+# read the other's count, so no span is counted twice.
+COUNT = """
+WITH serving AS (
+    SELECT * FROM unnest(%(orgs)s::text[], %(names)s::text[]) AS app (org, name)
+), previous AS (
+    SELECT app.org, app.name, app.metered_at,
+           EXISTS (SELECT 1 FROM serving WHERE org = app.org AND name = app.name) AS serves
+    FROM hosted_apps app
+    WHERE app.env = %(env)s AND (app.metered_at IS NOT NULL OR EXISTS (
+        SELECT 1 FROM serving WHERE org = app.org AND name = app.name))
+    FOR UPDATE OF app
+), counted AS (
+    INSERT INTO hosted_usage (org, env, name, day, seconds)
+    SELECT org, %(env)s, name, %(day)s,
+           least(greatest(extract(epoch FROM %(now)s::timestamptz - metered_at), 0), %(longest)s)
+    FROM previous WHERE serves AND metered_at IS NOT NULL
+    ON CONFLICT (org, env, name, day)
+    DO UPDATE SET seconds = hosted_usage.seconds + excluded.seconds
+)
+UPDATE hosted_apps app
+SET metered_at = CASE WHEN previous.serves THEN %(now)s::timestamptz END
+FROM previous
+WHERE app.env = %(env)s AND app.org = previous.org AND app.name = previous.name
+"""
+
+
+USAGE = """
+SELECT org, env, name, day, seconds FROM hosted_usage
+WHERE day >= %(since)s AND day < %(until)s AND (%(org)s::text IS NULL OR org = %(org)s)
+  AND (%(env)s::text IS NULL OR env = %(env)s)
+ORDER BY org, env, day, name
+"""
+
+
+@dataclass(frozen=True)
+class Logs:
+    """The last lines of an app's process as the runner read them; empty before the first read."""
+
+    host: str | None
+    lines: str
+    at: datetime | None
+
+
+@dataclass(frozen=True)
+class Served:
+    """The time an app served on one UTC day, in seconds."""
+
+    org: str
+    env: str
+    name: str
+    day: date
+    seconds: float
+
+
+async def stop_app(pool: Pool, app: HostedApp, *, by: str) -> None:
+    """Stop running the app: the runner drains it, and its releases and token stay."""
+    await _changed(pool, STOP, {**app.columns, "by": by}, app)
+
+
+async def start_app(pool: Pool, app: HostedApp) -> None:
+    """Run a stopped app again, its newest release."""
+    await _changed(pool, START, app.columns, app)
+
+
+async def ask_for_logs(pool: Pool, app: HostedApp) -> Logs:
+    """What the runner last sent of the app's output; the runner is told to send it again."""
+    async with pool.connection() as connection:
+        row = await (await connection.execute(ASK_LOGS, app.columns)).fetchone()
+    if row is None:
+        raise NotFound(NOT_HOSTED.format(name=app.name, env=app.env))
+    return Logs(host=row["logs_host"], lines=row["logs"], at=row["logs_at"])
+
+
+async def keep_logs(pool: Pool, app: HostedApp, *, host: str, lines: str) -> None:
+    """Keep what the runner read of the app's output, its last part when it is long."""
+    values = {**app.columns, "host": host, "lines": lines[-LONGEST_LOGS:]}
+    async with pool.connection() as connection:
+        await connection.execute(KEEP_LOGS, values)
+
+
+async def count_serving(
+    pool: Pool, env: Env, serving: Sequence[tuple[str, str]], now: datetime
+) -> None:
+    """Count the time the world's serving apps, each an (org, name), served since the last count."""
+    values = {
+        "env": env,
+        "orgs": [org for org, _ in serving],
+        "names": [name for _, name in serving],
+        "now": now,
+        "day": now.astimezone(UTC).date(),
+        "longest": LONGEST_BEAT_S,
+    }
+    async with pool.connection() as connection, connection.transaction():
+        await connection.execute(COUNT, values)
+
+
+async def served(
+    pool: Pool, since: date, until: date, *, org: str | None = None, env: str | None = None
+) -> list[Served]:
+    """The time apps served per UTC day in [since, until), one org's or every org's."""
+    values = {"since": since, "until": until, "org": org, "env": env}
+    async with pool.connection() as connection:
+        rows = await (await connection.execute(USAGE, values)).fetchall()
+    return [
+        Served(
+            org=row["org"], env=row["env"], name=row["name"], day=row["day"], seconds=row["seconds"]
+        )
+        for row in rows
+    ]
+
+
+async def _changed(
+    pool: Pool, statement: LiteralString, values: dict[str, str], app: HostedApp
+) -> None:
+    async with pool.connection() as connection:
+        row = await (await connection.execute(statement, values)).fetchone()
+    if row is None:
+        raise NotFound(NOT_HOSTED.format(name=app.name, env=app.env))