From 5959f93f5ff67880845292ad07e9a83dca840d2f Mon Sep 17 00:00:00 2001 From: GuillaumePELLUET Date: Wed, 5 Aug 2026 14:53:45 +0200 Subject: [PATCH 01/11] Remove OpenRouter dependencies with local llm to opencode --- AGENTS.md | 9 +- README.md | 14 + agent_config.yaml | 34 +- docs/agent.rst | 195 +++++- docs/configuration.rst | 126 ++-- docs/usage/parameters.rst | 31 +- docs/user_commands.rst | 4 +- docs/weights_studio.rst | 42 +- pyproject.toml | 8 +- tests/backend/test_cli_additional_unit.py | 18 +- tests/gRPC/test_grpc_user_actions.py | 2 +- .../test_agent_live_prompt_evaluation.py | 83 +-- .../test_agent_model_and_safety_unit.py | 81 +-- .../services/test_agent_opencode_provider.py | 291 ++++++++ .../services/test_agent_prompt_unit.py | 95 +-- .../services/test_agent_service_unit.py | 68 +- tests/trainer/services/test_opencode_chat.py | 250 +++++++ tests/ui/test_server_agent.py | 238 +++++++ tests/ui/test_server_loop.py | 356 ++++++++++ weightslab/backend/cli.py | 56 +- weightslab/proto/experiment_service.proto | 23 + weightslab/proto/experiment_service_pb2.py | 386 +++++----- .../proto/experiment_service_pb2_grpc.py | 433 +++++++----- weightslab/trainer/services/agent/agent.py | 525 ++++---------- .../trainer/services/agent/opencode_chat.py | 229 ++++++ weightslab/trainer/services/agent_service.py | 49 +- weightslab/trainer/services/data_service.py | 2 +- weightslab/trainer/services/utils/tools.py | 7 +- weightslab/ui/server.py | 662 ++++++++++++++++++ 29 files changed, 3126 insertions(+), 1191 deletions(-) create mode 100644 tests/trainer/services/test_agent_opencode_provider.py create mode 100644 tests/trainer/services/test_opencode_chat.py create mode 100644 tests/ui/test_server_agent.py create mode 100644 tests/ui/test_server_loop.py create mode 100644 weightslab/trainer/services/agent/opencode_chat.py diff --git a/AGENTS.md b/AGENTS.md index b7b63b3e..9629fab8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -222,10 +222,11 @@ gates and separates phases. **Large weights/images fail to transfer.** Raise `GRPC_MAX_MESSAGE_BYTES`. -**The agent bar says it's unconfigured.** The LLM agent needs a provider: a -local **Ollama** server (`provider: ollama`, available immediately) or **cloud -OpenRouter** initialized from the UI via `/init` (then `/model` to switch, -`/reset` to clear). See `weightslab/docs/weights_studio.rst`. +**The agent bar says it's unconfigured.** The LLM agent is backed entirely by +a local **OpenCode** server (`OPENCODE_URL`, default `http://127.0.0.1:4096`) +— WeightsLab starts one for you on first use. Initialize from the UI via +`/init` (then `/model` to switch, `/reset` to clear). See +`weightslab/docs/agent.rst` and `weightslab/docs/weights_studio.rst`. --- diff --git a/README.md b/README.md index 7297a035..776001c5 100644 --- a/README.md +++ b/README.md @@ -321,6 +321,20 @@ Find our documentation [online](https://grayboxtech.github.io/weightslab/latest/ +
+Agent: chat with your training run (OpenCode) + +
+ +WeightsLab ships two distinct agent surfaces — the backend SDK agent for data-manipulation +queries, and a local [OpenCode](https://opencode.ai)-backed agent with a full bash/file +toolset that can restart training, edit your code, and run recurring `/loop` monitoring +jobs. See the [Agent docs](https://grayboxtech.github.io/weightslab/latest/agent.html) for +how the two connect, how to point either one at a local model, and the full `/loop` +reference. + +
+
diff --git a/agent_config.yaml b/agent_config.yaml index ca1ef1d5..ad7d5c38 100644 --- a/agent_config.yaml +++ b/agent_config.yaml @@ -4,25 +4,19 @@ # or directly in the config file. # Config file values will override env. variables if both are set. # If cloned, env. variables can be defined in a .env file at the root of the repository. +# +# OpenCode (opencode.ai) is the only supported agent backend: a local OpenCode +# server backs every LLM call. There is no API key here -- the credential +# lives in OpenCode's own config, entered once via `opencode auth login` or +# the Weights Studio landing page's login modal. agent: - # Select the model provider. - # Local: 'ollama' - # Remote: 'openrouter' - # provider: openrouter # Default to OpenRouter if API key is provided, otherwise fallback to local Ollama. This can be overridden by env variable PREFERRED_PROVIDER. + # URL of the local OpenCode server (can also be set as env variable + # OPENCODE_URL). Defaults to http://127.0.0.1:4096. This is the SAME shared + # root env var the frontend reads, so set it once and both sides point at + # one server. + opencode_url: http://127.0.0.1:4096 - # Local Settings - fallback_to_local: false - ollama_model: llama3.2:3b - - # # Remote Model Selection - # Default is a fast flash-class model. The intent-planning task is simple JSON - # generation, so a small/fast model responds in ~2-4s where a 70B model took - # ~15-30s for no accuracy gain. Switch back to a large model here if you see - # accuracy issues (speed/accuracy tradeoff). - openrouter_model: bytedance-seed/seed-2.0-lite # Open router model name (can also be set as env variable OPENROUTER_MODEL). Fast alternatives: openai/gpt-4o-mini, meta-llama/llama-3.1-8b-instruct. Accurate/slow: ~google/gemini-flash-latest - # openrouter_api_key: # Open router API key (can also be set as env variable OPENROUTER_API_KEY) - openrouter_base_url: https://openrouter.ai/api/v1 # Open router base URL (can also be set as env variable OPENROUTER_BASE_URL) - openrouter_request_timeout: 60.0 # Timeout for OpenRouter API requests in seconds (can also be set as env variable OPENROUTER_REQUEST_TIMEOUT) - openrouter_max_tokens: 2048 # Max completion length. OpenRouter reserves max_tokens*price against the key budget BEFORE generating, so an uncapped value can 402 ("more credits, or fewer max_tokens") on a credit/weekly-limited key. Raise only if responses get truncated (env: OPENROUTER_MAX_TOKENS) - openrouter_provider_sort: throughput # Bias OpenRouter's upstream routing to avoid slow providers: 'throughput' | 'latency' | 'price'. Empty string = let OpenRouter choose (env: OPENROUTER_PROVIDER_SORT) - openrouter_structured_output: false # Ask the model for a schema-validated plan directly (skips free-form JSON + regex repair). More reliable, but only works on models whose OpenRouter route supports structured/JSON-schema output (e.g. Gemini, GPT-4o). env: OPENROUTER_STRUCTURED_OUTPUT=1 + # OpenCode model, "providerID/modelID" (can also be set as env variable + # OPENCODE_MODEL). Empty string lets the OpenCode server use its own + # configured default model. + opencode_model: "" diff --git a/docs/agent.rst b/docs/agent.rst index 9f00ed89..9d8b3b44 100644 --- a/docs/agent.rst +++ b/docs/agent.rst @@ -16,6 +16,52 @@ leaves the process except the prompt text sent to the configured LLM provider. Describe what you want in plain English; the agent translates it into a safe, reviewable plan of dataframe and model operations and executes it. +Two agent surfaces, one OpenCode server +----------------------------------------- + +WeightsLab's agent capability is backed entirely by `OpenCode +`_ — a local ``opencode serve`` process that WeightsLab +starts (or reuses) for you. There is no separate OpenRouter/Ollama +integration to configure: OpenCode itself is the provider layer, and its own +config (``opencode auth login``, or the login modal described below) holds +whatever credentials you use — OpenRouter, Anthropic, a local Ollama model, +anything OpenCode supports. + +That one server backs **two very different agent surfaces**, and knowing +which one you're talking to matters — everything on the rest of this page +describes the first one: + +.. list-table:: + :header-rows: 1 + :widths: 20 40 40 + + * - + - Backend SDK agent + - "Frontend" / OpenCode agent + * - Drives + - The normal query bar and chat-history-panel conversation + (``DataManipulationAgent``, ``weightslab/trainer/services/agent/agent.py``) + - The landing-page chat (pre-experiment) and ``/loop`` (during an experiment) + * - Toolset + - None — every mutating tool (``write``/``edit``/``patch``/``bash``) is + explicitly disabled on every call (``opencode_chat.py``'s + ``_MUTATING_TOOLS``) + - Full toolset — bash, file read/write/edit/patch + * - Memory + - ``self.history``, cleared/summarized by ``/clear`` and ``/compact`` + - An OpenCode session (server-side); cleared/summarized the same way, via + OpenCode's own session delete/summarize endpoints + * - Talks to OpenCode + - In one-shot mode: send a prompt, get text back, no side effects + - Interactively: it can restart training, edit your code, discard/tag + data, run reports + +**During an active experiment, the only way to reach the frontend/OpenCode +agent is** ``/loop`` **from the experiment agent bar.** The landing-page chat +only exists pre-experiment — once you're connected to a running experiment, +that surface is gone, and ``/loop`` (see the "``/loop`` reference" section +near the end of this page) is the sole entry point to the same kind of agent. + What the agent can do --------------------- @@ -171,12 +217,15 @@ dataframe state: .. code-block:: bash - export UTEST_AGENT_PROMPT_EVALUATION=sk-or-... # OpenRouter API key - export UTEST_AGENT_PROMPT_EVALUATION_MODEL=openai/gpt-4o-mini # optional + # Requires a local OpenCode server already running and authenticated + # (opencode has no API-key env var of its own -- see "Initializing the + # agent" below). + export UTEST_AGENT_PROMPT_EVALUATION=1 + export OPENCODE_MODEL=openrouter/anthropic/claude-opus-4.6 # optional pytest weightslab/tests/trainer/services/test_agent_live_prompt_evaluation.py -v Without ``UTEST_AGENT_PROMPT_EVALUATION`` set, the suite logs a note and -skips entirely (it never runs by accident in CI or consumes API credits +skips entirely (it never runs by accident in CI or against a real model unintentionally). A small always-on sanity check for the harness itself (fixture shape, op-runner correctness) still runs regardless. @@ -241,15 +290,42 @@ scenario end-to-end against a real model. Initializing the agent ---------------------- -The agent needs an LLM provider before it can serve requests. Two provider -families are supported: +The agent needs a reachable OpenCode server before it can serve requests -- +OpenCode is the only supported backend. Nothing to install beyond WeightsLab +itself: ``opencode-ai``'s bundled binary ships with the UI's dependencies, and +the UI server (``weightslab/ui/server.py``) starts an ``opencode serve`` child +process on first use, rooted at your experiment directory, tearing it down +when the UI server exits. -- **OpenRouter** — cloud-hosted models (recommended; interactive onboarding in - the UI). -- **Ollama** — local inference, available immediately at backend startup when - configured in ``agent_config.yaml``. +Both agent surfaces (see above) converge on the **same** OpenCode server via +one shared environment variable: -You can initialize it three ways. +.. code-block:: bash + + export OPENCODE_URL=http://127.0.0.1:4096 # or wherever your own `opencode serve` is running + +If ``OPENCODE_URL`` is set and reachable, the UI server adopts it directly +instead of spawning a child; the backend SDK agent reads the same variable +(``agent.py``'s ``_load_config``) — set it once and both sides talk to the one +server, so a model you authenticate once is available everywhere. +``OPENCODE_MODEL`` (or ``agent_config.yaml``'s ``agent.opencode_model``) picks +the default model for the backend SDK agent, as an OpenCode +``providerID/modelID`` string (e.g. ``openrouter/anthropic/claude-opus-4.6``). +Leave it unset to use OpenCode's own configured default. + +Credentials and provider setup live in OpenCode itself, never in WeightsLab: + +.. code-block:: bash + + opencode auth login # OpenRouter, Anthropic, a local Ollama endpoint, anything OpenCode supports + +or, from the browser, the landing page's login modal drives the same flow +without a terminal. For a fully local setup, point OpenCode's own config at +Ollama (or any other local provider it supports) — WeightsLab needs no +changes on its side; it just asks OpenCode for whichever model you've +selected. + +You can initialize the backend SDK agent three ways. Option 1 — Weights Studio UI (recommended) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -265,13 +341,12 @@ Type one of these commands into the chat bar: * - Command - Effect * - ``/init`` - - Opens the OpenRouter onboarding modal. Choose **A — Enter OpenRouter API - key** (paste an ``sk-or-…`` key) or **B — Get API key from OpenRouter** - (OAuth flow), then pick a model. On success the placeholder switches to a + - Connects to the OpenCode server (see ``OPENCODE_URL`` above) and lets + you pick a model. On success the placeholder switches to a ready-to-use example query. * - ``/model`` - - Opens the model browser to switch the active OpenRouter model without - re-entering the API key. + - Opens the model browser to switch the active OpenCode model without + reconnecting. * - ``/reset`` - Clears the current connection and returns the agent to the uninitialized state. @@ -294,9 +369,9 @@ the interactive console exposes an ``agent`` verb: .. code-block:: text agent status # Is the agent available? - agent init --api-key sk-or-... --model openai/gpt-4o-mini [--timeout 20] - agent models # List available OpenRouter models - agent model ~google/gemini-flash-latest # Switch model + agent init [--model openrouter/anthropic/claude-opus-4.6] + agent models # List available OpenCode models + agent model openrouter/openai/gpt-5 # Switch model agent reset # Clear the connection agent query # Run a natural-language request query # Shortcut for `agent query` @@ -314,40 +389,33 @@ UI). For example: Option 3 — Startup configuration file ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -To have a provider ready the moment the backend starts (no ``/init`` needed), -configure ``agent_config.yaml`` and/or environment variables. This is the only -way to enable the local Ollama provider. +To have the agent ready the moment the backend starts (no ``/init`` needed), +configure ``agent_config.yaml`` and/or environment variables. .. code-block:: yaml # agent_config.yaml (repo root, package root, cwd, or $AGENT_CONFIG_PATH) agent: - provider: openrouter # or "ollama" - openrouter_model: ~google/gemini-flash-latest - fallback_to_local: false - # Local Ollama alternative: - ollama_model: llama3.2:3b - ollama_host: localhost - ollama_port: 11435 + opencode_url: http://127.0.0.1:4096 + opencode_model: "" # empty = use OpenCode's own configured default .. code-block:: bash - # Prefer secrets via environment variables over YAML. - export OPENROUTER_API_KEY=your_openrouter_key + # Equivalent environment variables (config file wins if both are set). + export OPENCODE_URL=http://127.0.0.1:4096 + export OPENCODE_MODEL=openrouter/anthropic/claude-opus-4.6 See :doc:`configuration` for the full list of agent environment variables, the ``agent_config.yaml`` lookup order, and every supported YAML key. .. note:: - "Available" means the credentials were actually confirmed to work, not - just that a client object was constructed. A key configured via - ``agent_config.yaml``/environment variables (Option 3) is probed once at - backend startup exactly like the ``/init`` UI flow already does, and if a - live query ever gets rejected with 401, the connection is immediately - marked unavailable rather than continuing to report "ready" until the - next restart. If health checks and real requests ever disagree, that's a - bug — the two are kept in sync by design. + "Available" only means a client object was constructed against + ``OPENCODE_URL`` -- OpenCode's own constructor never eagerly connects, so + there is nothing to probe at backend startup the way a cloud API key + needed a connectivity check. Actual unreachability (server not running, or + later restarted) surfaces on the first real query instead, which is + reported through the normal "Internal Agent Error"/reconnect path. Using the agent effectively ---------------------------- @@ -646,6 +714,57 @@ expression it builds for you. comparisons are ``False``) — the query never errors, it just excludes those rows. +``/loop`` reference +---------------------- + +``/loop``, typed into the **experiment agent bar**, is the other agent +surface described at the top of this page: it starts a recurring check-in +against a dedicated OpenCode session — the same kind of session the +landing-page chat uses, with the same full toolset. It never touches the +backend SDK agent directly. + +.. code-block:: text + + /loop 30m Watch the training loss and loss_shape trends; if the run stalls or diverges, pause it and tell me why + /loop list + /loop stop + +- **Syntax**: ``/loop m|h `` to start (minimum interval: 60s), + ``/loop list`` to see running jobs, ``/loop stop `` to cancel one. +- **What it can do**: the loop's OpenCode session is told about the local + ``weightslab`` CLI, reachable over bash against the live training process: + + - ``weightslab pause`` / ``weightslab resume`` — freeze/resume weight updates + - ``weightslab discard `` — discard a sample by id + - ``weightslab agent query ""`` — hands the request to the + **backend SDK agent's** own intent pipeline, e.g. ``weightslab agent + query "discard samples where loss > 5 and tag them hard_examples"``. This + is how the loop reaches back into the database/history: it can't ask the + backend agent directly, but it can drive it through the CLI. + - ``weightslab status`` — a snapshot of hyperparameters/model/training state + + These four are what the loop's system prompt explicitly calls out, but bash + access means any other ``weightslab`` CLI verb is reachable too — e.g. + ``weightslab report`` to generate a narrative report for the loop to read + and act on. It may also read/edit training code directly and attempt to + restart a crashed process via bash — this is best-effort (no supervisor or + PID handoff): it looks for the process, stops it if still running, and + re-launches from whatever it can determine (shell history, a run script, + logs). There is no dedicated restart command. +- **Concurrency cap**: at most 3 loops at once, shared across both chat + surfaces (they hit the same registry). A 4th ``/loop start`` is rejected + with an error rather than silently stopping an older job — stop one first + with ``/loop stop ``. +- **Managing running jobs**: a panel pinned at the top of the chat-history + window lists every running job with a live countdown to its next check-in, + and lets you edit a job's prompt/interval in place or stop it — no need to + remember ``/loop stop `` if the panel is in view. ``/loop list``/``/loop + stop`` also work from the landing-page chat pre-experiment, hitting the + same registry. +- **Persistence**: a loop is tied to the running ``weightslab start`` process, + not the browser tab — it survives a page reload or closed tab, but not a + full restart of the UI server. + Workflow pattern ---------------- diff --git a/docs/configuration.rst b/docs/configuration.rst index cd006a70..94bd97bc 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -683,10 +683,15 @@ Evaluation Mode ``0`` disables the absolute override and uses the dynamic formula only. -AI / LLM API Keys -~~~~~~~~~~~~~~~~~ +AI / LLM Configuration +~~~~~~~~~~~~~~~~~~~~~~ -These keys are required only when using the agentic data-query features. +OpenCode (`opencode.ai `_) is the agent's only supported +backend -- there is no API key to set here. The credential lives in +OpenCode's own config (``opencode auth login``, or the Weights Studio landing +page's login modal), which can point at OpenRouter, Anthropic, a local Ollama +endpoint, or anything else OpenCode supports; WeightsLab itself only needs to +know which server to talk to. .. list-table:: :header-rows: 1 @@ -695,19 +700,23 @@ These keys are required only when using the agentic data-query features. * - Variable - Default - Description - * - ``OPENROUTER_API_KEY`` + * - ``OPENCODE_URL`` + - ``http://127.0.0.1:4096`` + - URL of the local OpenCode server. Shared with the frontend (Weights + Studio's landing-page chat and ``/loop``) via the same variable, so + set it once and both sides talk to the one server. + * - ``OPENCODE_MODEL`` - *(empty)* - - OpenRouter API key ? required for cloud agent setup in Weights Studio. + - Default model for the backend SDK agent, as an OpenCode + ``providerID/modelID`` string (e.g. + ``openrouter/anthropic/claude-opus-4.6``). Empty uses OpenCode's own + configured default. Agent Configuration ~~~~~~~~~~~~~~~~~~~ These variables control how the data-query agent finds its YAML configuration. -The agent supports two provider families: - -- ``ollama`` for local inference -- ``openrouter`` for cloud-hosted models .. list-table:: :header-rows: 1 @@ -744,7 +753,7 @@ Agent Provider Setup ~~~~~~~~~~~~~~~~~~~~ The runtime agent is configured from ``agent_config.yaml`` plus optional -environment variables such as ``OPENROUTER_API_KEY``. +environment variables such as ``OPENCODE_URL``. Supported YAML keys ^^^^^^^^^^^^^^^^^^^ @@ -756,73 +765,33 @@ Supported YAML keys * - Key - Example - Description - * - ``agent.provider`` - - ``ollama`` - - Active provider. Common values: ``ollama`` or ``openrouter``. - * - ``agent.ollama_model`` - - ``llama3.2:3b`` - - Local Ollama model name. - * - ``agent.ollama_host`` - - ``localhost`` - - Ollama host. - * - ``agent.ollama_port`` - - ``11435`` - - Ollama HTTP port used by WeightsLab. - * - ``agent.openrouter_model`` - - ``~google/gemini-flash-latest`` - - Default OpenRouter model. - * - ``agent.openrouter_base_url`` - - ``https://openrouter.ai/api/v1`` - - OpenRouter-compatible base URL. - * - ``agent.openrouter_request_timeout`` - - ``15.0`` - - Request timeout in seconds for OpenRouter calls. - * - ``agent.openrouter_api_key`` - - *(secret)* - - Optional API key in YAML. Prefer environment variables or UI init when possible. - * - ``agent.fallback_to_local`` - - ``false`` - - If enabled, WeightsLab also tries the local Ollama provider as fallback. - -Local Ollama example -^^^^^^^^^^^^^^^^^^^^ - -Use this mode when you want the agent available immediately at backend startup. - -.. code-block:: yaml - - agent: - provider: ollama - ollama_model: llama3.2:3b - ollama_host: localhost - ollama_port: 11435 - fallback_to_local: false - -Operational steps: + * - ``agent.opencode_url`` + - ``http://127.0.0.1:4096`` + - URL of the local OpenCode server. + * - ``agent.opencode_model`` + - ``openrouter/anthropic/claude-opus-4.6`` + - Default model, as an OpenCode ``providerID/modelID`` string. Empty + uses OpenCode's own configured default. -1. Install Ollama. -2. Pull a model, for example ``ollama pull llama3.2:3b``. -3. Start the Ollama server. -4. Start WeightsLab. -5. Open Weights Studio and query the agent directly. - -Cloud OpenRouter example -^^^^^^^^^^^^^^^^^^^^^^^^ - -Use this mode when you want hosted models and interactive setup from Weights Studio. +Example +^^^^^^^ .. code-block:: yaml agent: - provider: openrouter - openrouter_model: ~google/gemini-flash-latest - fallback_to_local: false - -Recommended secret handling: + opencode_url: http://127.0.0.1:4096 + opencode_model: "" # empty = use OpenCode's own configured default -.. code-block:: bash +Setup steps: - export OPENROUTER_API_KEY=your_openrouter_key +1. Have a local OpenCode server running (WeightsLab starts one for you on + first use; see :doc:`agent`), or point ``OPENCODE_URL`` at your own. +2. Authenticate it once: ``opencode auth login`` (or the login modal from + Weights Studio's landing page) -- OpenRouter, Anthropic, a local Ollama + endpoint, anything OpenCode supports. +3. Start WeightsLab. +4. Open Weights Studio and query the agent directly, or type ``/init`` first + to pick a specific model. Weights Studio commands ^^^^^^^^^^^^^^^^^^^^^^^ @@ -830,11 +799,10 @@ Weights Studio commands When using Weights Studio, the agent bar supports these runtime commands: 1. ``/init`` - Opens the OpenRouter onboarding flow. - Users can enter an API key manually or use the OAuth flow, then select a model. + Connects to the OpenCode server and lets you pick a model. 2. ``/model`` - Opens the model browser and switches the active OpenRouter model without - requiring a full reinitialization. + Opens the model browser and switches the active OpenCode model without + reconnecting. 3. ``/reset`` Clears the current runtime connection state and returns the agent to the uninitialized status. @@ -842,13 +810,13 @@ When using Weights Studio, the agent bar supports these runtime commands: Notes ^^^^^ -- The default OpenRouter model is ``~google/gemini-flash-latest``. -- The model browser fetches the available models from OpenRouter using the - configured API key. +- The model browser fetches the available models from the OpenCode server's + own provider catalog. - Connection and model-change actions are recorded in the agent history as log-style entries. -- ``/reset`` clears the current runtime agent state. If your startup config is - local-only and you want that provider back immediately, restart the backend. +- ``/reset`` clears the current runtime agent state. If your startup config + points at a server that's still running, ``/init`` reconnects immediately; + otherwise restart the backend once the server is back. Testing diff --git a/docs/usage/parameters.rst b/docs/usage/parameters.rst index c7fb22f0..254fa0f9 100644 --- a/docs/usage/parameters.rst +++ b/docs/usage/parameters.rst @@ -463,6 +463,9 @@ Audit logging LLM / agent integration (optional) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +The agent is backed entirely by a local OpenCode server (see :doc:`../agent`) +— there is no API key here; the credential lives in OpenCode's own config. + .. list-table:: :header-rows: 1 :widths: 35 15 50 @@ -470,27 +473,15 @@ LLM / agent integration (optional) * - Variable - Default - Description - * - ``OPENROUTER_API_KEY`` - - *(unset)* - - API key for OpenRouter. Required only when using WeightsLab's - LLM-assisted analysis features. - * - ``OPENROUTER_MODEL`` - - *(unset)* - - Model identifier forwarded to OpenRouter (e.g. - ``"openai/gpt-4o"``). - * - ``OPENROUTER_REQUEST_TIMEOUT`` + * - ``OPENCODE_URL`` + - ``http://127.0.0.1:4096`` + - URL of the local OpenCode server. Shared with the frontend, so set it + once and both sides talk to the one server. + * - ``OPENCODE_MODEL`` - *(unset)* - - Per-request timeout in seconds for OpenRouter calls. - * - ``OPENROUTER_MAX_TOKENS`` - - ``2048`` - - Maximum completion length requested from OpenRouter. OpenRouter - pre-authorizes ``max_tokens × completion_price`` against the key's - remaining budget *before* generating, so leaving this uncapped makes - the model request its full output window and can fail with a ``402`` - ("requires more credits, or fewer max_tokens") on a credit- or - weekly-limited key — even though the model is otherwise usable. The - default is ample for intent planning; raise it only if you see - truncated responses. + - Default model, as an OpenCode ``providerID/modelID`` string (e.g. + ``"openrouter/anthropic/claude-opus-4.6"``). Unset uses OpenCode's own + configured default. Telemetry ~~~~~~~~~~ diff --git a/docs/user_commands.rst b/docs/user_commands.rst index 8eebbdbe..9ef44d1f 100644 --- a/docs/user_commands.rst +++ b/docs/user_commands.rst @@ -487,9 +487,9 @@ sub-verb reference, examples, and setup: see :doc:`agent`. .. code-block:: text agent status - agent init --api-key sk-or-... --model openai/gpt-4o-mini --timeout 20 + agent init --model openrouter/anthropic/claude-opus-4.6 agent models - agent model google/gemini-flash-latest + agent model openrouter/openai/gpt-5 ask tag train samples with loss > 1.2 as goldset Experiment report diff --git a/docs/weights_studio.rst b/docs/weights_studio.rst index b90e5fd1..3f66bbab 100644 --- a/docs/weights_studio.rst +++ b/docs/weights_studio.rst @@ -195,45 +195,41 @@ Agent Usage in Weights Studio ------------------------------ Weights Studio includes an agent bar and an expandable agent history window. -The agent can run with either: +The agent is backed entirely by a local OpenCode server (`opencode.ai +`_) — see :doc:`agent` for the full setup story and the +distinction between this chat-bar agent and the separate ``/loop``/landing-page +OpenCode agent. -- a local Ollama provider configured on the backend -- a cloud OpenRouter provider configured at startup or initialized from the UI - -Local Ollama workflow -~~~~~~~~~~~~~~~~~~~~~ +OpenCode workflow +~~~~~~~~~~~~~~~~~~ -If the backend is configured with ``provider: ollama`` and the Ollama server is -running, the agent is available immediately after backend startup. +WeightsLab starts (or reuses) a local ``opencode serve`` process for you, so +there's normally nothing to configure before the agent is available. If the +backend isn't connected to it yet, Weights Studio shows the agent as +unconfigured and the input placeholder instructs the user to type ``/init``. -Typical local setup: +Typical setup: -1. Start Ollama. +1. Authenticate OpenCode once, if you haven't already: ``opencode auth + login`` (or the landing page's login modal) — OpenRouter, Anthropic, a + local Ollama endpoint, anything OpenCode supports. 2. Start WeightsLab (``wl.serve(serving_grpc=True)``). 3. Start Weights Studio (``weightslab start``). -4. Ask questions in the agent bar. - -Cloud OpenRouter workflow -~~~~~~~~~~~~~~~~~~~~~~~~~ - -If the backend is not initialized with a cloud key yet, Weights Studio shows -the agent as unconfigured and the input placeholder instructs the user to type -``/init``. +4. Ask questions in the agent bar, or type ``/init`` first to pick a specific + model. ``/init`` flow: 1. Type ``/init`` in the agent input. -2. Choose manual API key entry or the OpenRouter OAuth flow. +2. Weights Studio connects to the OpenCode server. 3. Select a model from the available model list. 4. Confirm to initialize the runtime connection. -The default cloud model is ``~google/gemini-flash-latest``. - Available agent commands ~~~~~~~~~~~~~~~~~~~~~~~~ -- ``/init`` — initialize OpenRouter from the UI -- ``/model`` — open the model chooser to switch the active OpenRouter model +- ``/init`` — connect to the OpenCode server from the UI +- ``/model`` — open the model chooser to switch the active OpenCode model - ``/reset`` — clear the current agent runtime connection and status History behavior diff --git a/pyproject.toml b/pyproject.toml index 170624f9..3a741a79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,10 +81,12 @@ dependencies = [ # Environment variable loading (used by agent service) "python-dotenv>=1,<2", - # Agent service runtime deps (imported by default module graph) + # Agent service runtime deps (imported by default module graph). Only + # langchain-core is needed -- it provides the Runnable/ChatPromptTemplate + # abstraction OpenCodeChat.as_runnable() plugs into. No OpenCode-specific + # SDK package exists; its client is hand-rolled in opencode_chat.py via + # stdlib urllib/threading. "langchain-core>=0.3,<2", - "langchain-ollama>=0.2,<2", - "langchain-openai>=0.2,<2", # Jupyter "ipykernel>=6.29,<7", diff --git a/tests/backend/test_cli_additional_unit.py b/tests/backend/test_cli_additional_unit.py index 0cbe947c..aa60576e 100644 --- a/tests/backend/test_cli_additional_unit.py +++ b/tests/backend/test_cli_additional_unit.py @@ -71,29 +71,27 @@ def test_add_tag_uses_sample_id_helper_for_multiple_samples(self): self.assertTrue(result["ok"]) tag_mock.assert_called_once_with(sample_ids=["sample_001", "sample_002", "sample_003"], tag="goldset", mode="add") - def test_agent_init_accepts_api_key_model_and_timeout(self): + def test_agent_init_accepts_a_model(self): agent = MagicMock() - agent.openrouter_request_timeout = 15.0 - agent.openrouter_model = "initial-model" - agent.initialize_with_cloud_key.return_value = (True, "Agent initialized successfully. Ready to help you.") + agent.opencode_model = "initial-model" + agent.initialize_with_cloud_key.return_value = (True, "Agent initialized successfully via OpenCode. Ready to help you.") cli_backend.set_cli_agent(agent) - result = _handle_command("agent init --api-key test-key --model openai/gpt-4o-mini --timeout 22") + result = _handle_command("agent init --model openrouter/openai/gpt-5") self.assertTrue(result["ok"]) - agent.initialize_with_cloud_key.assert_called_once_with("test-key", "openrouter", "openai/gpt-4o-mini") - self.assertEqual(agent.openrouter_request_timeout, 22.0) + agent.initialize_with_cloud_key.assert_called_once_with("", "opencode", "openrouter/openai/gpt-5") def test_agent_model_command_switches_model(self): agent = MagicMock() - agent.openrouter_model = "google/gemini-2.5-flash" + agent.opencode_model = "openrouter/anthropic/claude-opus-4.6" agent.change_model.return_value = (True, "Model switched") cli_backend.set_cli_agent(agent) - result = _handle_command("agent model google/gemini-2.5-flash") + result = _handle_command("agent model openrouter/anthropic/claude-opus-4.6") self.assertTrue(result["ok"]) - agent.change_model.assert_called_once_with("google/gemini-2.5-flash") + agent.change_model.assert_called_once_with("openrouter/anthropic/claude-opus-4.6") def test_agent_query_uses_data_service_when_available(self): mock_response = MagicMock( diff --git a/tests/gRPC/test_grpc_user_actions.py b/tests/gRPC/test_grpc_user_actions.py index 530d8e7d..6f2fcafe 100644 --- a/tests/gRPC/test_grpc_user_actions.py +++ b/tests/gRPC/test_grpc_user_actions.py @@ -290,7 +290,7 @@ def _make_real_data_service(self): # vectorized path and doesn't need it, but GetDataSamples does). ds._data_executor = ThreadPoolExecutor(max_workers=2) ds._agent = MagicMock() - ds._agent.is_ollama_available.return_value = True + ds._agent.is_available.return_value = True ds.audit_logger = MagicMock() return ds, df_manager diff --git a/tests/trainer/services/test_agent_live_prompt_evaluation.py b/tests/trainer/services/test_agent_live_prompt_evaluation.py index ceaedeaa..1139289d 100644 --- a/tests/trainer/services/test_agent_live_prompt_evaluation.py +++ b/tests/trainer/services/test_agent_live_prompt_evaluation.py @@ -2,30 +2,25 @@ Live-LLM evaluation of the Data Manipulation Agent against a batch of realistic user prompts. -This suite is OPT-IN: every test calls a REAL LLM through OpenRouter -(consuming API credits and real wall-clock time), so it is skipped entirely -unless an OpenRouter API key can be resolved. The key/model are resolved, in -priority order, from: - - 1. The dedicated ``UTEST_AGENT_PROMPT_EVALUATION`` / - ``UTEST_AGENT_PROMPT_EVALUATION_MODEL`` env vars (explicit opt-in). - 2. The standard ``OPENROUTER_API_KEY`` / ``OPENROUTER_MODEL`` env vars — - including any loaded from a repo ``.env`` file — so the same credentials - the running agent uses also drive this suite with no extra setup. - -So any of these work from the command line: - - # reuse your existing OpenRouter config (.env or exported env var) +This suite is OPT-IN: every test calls a REAL LLM through a local OpenCode +server (opencode.ai) (consuming real wall-clock time), so it is skipped +entirely unless explicitly turned on via the ``UTEST_AGENT_PROMPT_EVALUATION`` +env var (any non-empty value). Unlike the old OpenRouter-backed version of +this suite, there is no API key to resolve here -- OpenCode's credential +lives in its own config (``opencode auth login``), not in an env var -- so +opting in just means "run me", and ``OPENCODE_URL``/``OPENCODE_MODEL`` (the +SAME vars the agent itself reads) pick which server/model to run against. + +So this works from the command line, with a local OpenCode server already +running and authenticated: + + set UTEST_AGENT_PROMPT_EVALUATION=1 pytest weightslab/tests/trainer/services/test_agent_live_prompt_evaluation.py -v - # or pass explicitly for this run only (PowerShell) - $env:OPENROUTER_API_KEY="sk-or-..."; $env:OPENROUTER_MODEL="google/gemini-flash-latest"; pytest ... -v + # or target a non-default server/model for this run only (PowerShell) + $env:UTEST_AGENT_PROMPT_EVALUATION="1"; $env:OPENCODE_URL="http://127.0.0.1:4096"; $env:OPENCODE_MODEL="openrouter/anthropic/claude-opus-4.6"; pytest ... -v - # or the dedicated opt-in vars (cmd.exe) - set UTEST_AGENT_PROMPT_EVALUATION=sk-or-... - pytest ... -v - -The model defaults to the agent's own default OpenRouter model when unset. +The model defaults to the OpenCode server's own configured default when unset. Each test reproduces a specific, previously-reported bug/scenario and verifies the agent's plan, once executed against a realistic synthetic @@ -58,13 +53,16 @@ logger = logging.getLogger(__name__) -def _resolve_live_credentials() -> "tuple[str, str | None]": - """Resolve the OpenRouter (key, model) for the live suite. +def _resolve_live_opencode_config() -> "tuple[bool, str]": + """Resolve whether to run the live suite, and which OpenCode model to + request. - Mirrors the agent's own config loading: pull in any repo ``.env`` first, - then prefer the dedicated UTEST_* opt-in vars, falling back to the standard - OPENROUTER_* vars so the same credentials the agent runs on also drive this - suite without duplicating them. + OpenCode has no API-key concept (its credential lives in the OpenCode + server's own config, entered once via ``opencode auth login``), so opting + in is just a plain on/off switch -- ``UTEST_AGENT_PROMPT_EVALUATION`` set + to any non-empty value. ``OPENCODE_URL``/``OPENCODE_MODEL`` (the SAME env + vars the agent itself reads) pick which server/model to run against; any + repo ``.env`` is loaded first so they can live there too. """ if load_dotenv is not None: # weightslab/tests/trainer/services/ -> parents[4] = repo root, @@ -75,26 +73,19 @@ def _resolve_live_credentials() -> "tuple[str, str | None]": load_dotenv(dotenv_path=candidate, override=False) load_dotenv(override=False) - key = ( - os.environ.get("UTEST_AGENT_PROMPT_EVALUATION", "").strip() - or os.environ.get("OPENROUTER_API_KEY", "").strip() - ) - model = ( - os.environ.get("UTEST_AGENT_PROMPT_EVALUATION_MODEL", "").strip() - or os.environ.get("OPENROUTER_MODEL", "").strip() - or None - ) - return key, model + run_live = bool(os.environ.get("UTEST_AGENT_PROMPT_EVALUATION", "").strip()) + model = os.environ.get("OPENCODE_MODEL", "").strip() or None + return run_live, model -API_KEY, MODEL = _resolve_live_credentials() +RUN_LIVE, MODEL = _resolve_live_opencode_config() -if not API_KEY: +if not RUN_LIVE: logger.info( - "[test_agent_live_prompt_evaluation] No OpenRouter key found " - "(checked UTEST_AGENT_PROMPT_EVALUATION and OPENROUTER_API_KEY, incl. .env) -- " - "skipping live-LLM agent prompt evaluation tests. Set one of those (and optionally " - "OPENROUTER_MODEL) to run this suite against a real model." + "[test_agent_live_prompt_evaluation] UTEST_AGENT_PROMPT_EVALUATION not set -- " + "skipping live-LLM agent prompt evaluation tests. Set it to any value, with a " + "local OpenCode server running and authenticated (see OPENCODE_URL/OPENCODE_MODEL), " + "to run this suite against a real model." ) @@ -146,7 +137,7 @@ def _make_live_agent(df: pd.DataFrame, exp_ctx=None) -> DataManipulationAgent: # for the data-only tests (no model registered). ctx = SimpleNamespace(_all_datasets_df=df, _ctx=exp_ctx) agent = DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key(API_KEY, "openrouter", MODEL or agent.openrouter_model) + ok, message = agent.initialize_with_cloud_key("", "opencode", MODEL) if not ok: raise RuntimeError(f"Failed to initialize live agent for testing: {message}") return agent @@ -246,7 +237,7 @@ def _run_ops(df: pd.DataFrame, ops: list, model_service=None) -> "tuple[pd.DataF return df, messages -@unittest.skipUnless(API_KEY, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") +@unittest.skipUnless(RUN_LIVE, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") class TestAgentLivePromptEvaluation(unittest.TestCase): """Runs a battery of realistic user prompts against a REAL LLM and verifies the resulting dataframe/message state. Each test is a @@ -403,7 +394,7 @@ def test_reset_view_is_recognized(self): self.assertTrue(any(op.get("params", {}).get("__agent_reset__") for op in ops), ops) -@unittest.skipUnless(API_KEY, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") +@unittest.skipUnless(RUN_LIVE, "UTEST_AGENT_PROMPT_EVALUATION not set; skipping live-LLM agent evaluation") class TestAgentRstDocumentedPrompts(unittest.TestCase): """ One test per example prompt listed in docs/agent.rst's "Example prompts diff --git a/tests/trainer/services/test_agent_model_and_safety_unit.py b/tests/trainer/services/test_agent_model_and_safety_unit.py index 6f475cda..93e4d50d 100644 --- a/tests/trainer/services/test_agent_model_and_safety_unit.py +++ b/tests/trainer/services/test_agent_model_and_safety_unit.py @@ -16,13 +16,9 @@ def _install_agent_dependency_stubs(): stubs = { - "langchain_ollama": types.ModuleType("langchain_ollama"), - "langchain_openai": types.ModuleType("langchain_openai"), "langchain_core": types.ModuleType("langchain_core"), "langchain_core.prompts": types.ModuleType("langchain_core.prompts"), } - stubs["langchain_ollama"].ChatOllama = object - stubs["langchain_openai"].ChatOpenAI = object stubs["langchain_core.prompts"].ChatPromptTemplate = object return stubs @@ -37,9 +33,7 @@ def _make_agent(df=None): # `_ctx=None` means `_setup_model_schema` bails out early (no live model), # matching how a standalone agent behaves before any model is registered. ctx = SimpleNamespace(_all_datasets_df=df, _ctx=None) - - with mock.patch.object(agent_mod, "ChatOpenAI", None), mock.patch.object(agent_mod, "ChatOllama", None): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) return agent_mod, agent @@ -134,8 +128,6 @@ def fake_try_query_provider(provider, instruction, system_prompt): return None agent._try_query_provider = fake_try_query_provider - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False agent.query("keep only validation or test samples") @@ -168,8 +160,6 @@ def fake_try_query_provider(provider, instruction, system_prompt): return next(it) agent._try_query_provider = fake_try_query_provider - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False return agent, calls def test_history_starts_empty(self): @@ -231,64 +221,17 @@ def test_history_unaffected_by_a_failed_query(self): self.assertEqual(agent.history, []) -class TestStartupProviderVerification(unittest.TestCase): +class TestQueryAuthFailureHandling(unittest.TestCase): """Reported bug: CheckAgentHealth said the agent was available, but a - real query then failed with "401 Unauthorized". Root cause: - is_available() only checks that a provider CLIENT OBJECT exists, not - that its credentials were ever confirmed to work -- true for any key - loaded from agent_config.yaml/env vars, which (unlike the /init UI flow) - never runs a connectivity check. _verify_startup_providers() now probes - once at construction time, and a 401 detected during a real query - invalidates the cached connection so is_available() reflects reality - immediately instead of staying stale until the process restarts.""" - - def test_startup_verification_disables_a_bad_openrouter_key(self): - agent_mod, agent = _make_agent() - - class _FailingChatModel: - def __init__(self, *a, **kw): pass - def invoke(self, prompt): - raise RuntimeError("401 Unauthorized") - - agent.openrouter_api_key = "bad-key" - with mock.patch.object(agent_mod, "ChatOpenAI", _FailingChatModel): - agent._setup_providers() - agent._verify_startup_providers() - - self.assertIsNone(agent.chain_openrouter) - self.assertFalse(agent.is_available()) - - def test_startup_verification_keeps_a_good_openrouter_key(self): - agent_mod, agent = _make_agent() - - class _OkChatModel: - def __init__(self, *a, **kw): pass - def invoke(self, prompt): - return SimpleNamespace(content="OK") - - agent.openrouter_api_key = "good-key" - with mock.patch.object(agent_mod, "ChatOpenAI", _OkChatModel): - agent._setup_providers() - agent._verify_startup_providers() - - self.assertIsNotNone(agent.chain_openrouter) - self.assertTrue(agent.is_available()) - - def test_startup_verification_skips_probe_when_no_chain_was_built(self): - # No openrouter_api_key configured at all -> chain_openrouter stays - # None -> nothing to probe, must not raise. - _, agent = _make_agent() - agent.chain_openrouter = None - - agent._verify_startup_providers() # should be a no-op, not raise - - self.assertIsNone(agent.chain_openrouter) + real query then failed with "401 Unauthorized". A 401 detected during a + real query invalidates the cached OpenCode connection so is_available() + reflects reality immediately instead of staying stale until the process + restarts; a non-auth (e.g. transient/timeout) failure must NOT do that, + since the connection might still be perfectly valid.""" def test_401_during_query_invalidates_the_cached_connection(self): _, agent = _make_agent() - agent.chain_openrouter = object() # simulates an already-"available" cached client - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False + agent.chain_opencode = object() # simulates an already-"available" cached client def fake_try_query_provider(provider, instruction, system_prompt): agent._last_query_error = RuntimeError("401 Unauthorized") @@ -299,7 +242,7 @@ def fake_try_query_provider(provider, instruction, system_prompt): self.assertTrue(agent.is_available()) # stale "available" before the query result = agent.query("do something") - self.assertIsNone(agent.chain_openrouter) + self.assertIsNone(agent.chain_opencode) self.assertFalse(agent.is_available()) self.assertIn("Agent not connected", result[0]["params"]["reason"]) @@ -307,9 +250,7 @@ def test_non_auth_failure_does_not_invalidate_the_connection(self): # A transient/non-auth failure (e.g. timeout) must NOT disable a # connection that might still be perfectly valid. _, agent = _make_agent() - agent.chain_openrouter = object() - agent.preferred_provider = "openrouter" - agent.fallback_to_local = False + agent.chain_opencode = object() def fake_try_query_provider(provider, instruction, system_prompt): agent._last_query_error = RuntimeError("Connection timed out") @@ -319,7 +260,7 @@ def fake_try_query_provider(provider, instruction, system_prompt): agent.query("do something") - self.assertIsNotNone(agent.chain_openrouter) + self.assertIsNotNone(agent.chain_opencode) self.assertTrue(agent.is_available()) diff --git a/tests/trainer/services/test_agent_opencode_provider.py b/tests/trainer/services/test_agent_opencode_provider.py new file mode 100644 index 00000000..d1fd4bff --- /dev/null +++ b/tests/trainer/services/test_agent_opencode_provider.py @@ -0,0 +1,291 @@ +"""Tests for DataManipulationAgent's OpenCode provider: config loading, +_setup_providers wiring self.chain_opencode, initialize_with_cloud_key/ +change_model/get_available_models/reset_connection, and clear_history/ +compact_history. OpenCode is the only supported agent backend. + +Follows the exact same "_install_agent_dependency_stubs + _make_agent" pattern +already used in test_agent_model_and_safety_unit.py / test_agent_prompt_unit.py +(duplicated per-file by this repo's own convention, not imported across test +files). OpenCodeChat itself is mocked out here -- its own HTTP/SSE behavior is +covered by test_opencode_chat.py against a real fake server. +""" + +import importlib +import sys +import types +import unittest +from types import SimpleNamespace +from unittest import mock +from unittest.mock import MagicMock + +import pandas as pd + +# `_make_agent()` (like test_agent_model_and_safety_unit.py's identical helper) +# wraps each import in `mock.patch.dict(sys.modules, stubs, clear=False)`, which +# restores sys.modules to its EXACT pre-`with` snapshot on exit -- including +# evicting every module (torch, numpy, and their transitive dependency tree) +# that wasn't already resident when the block started. Those C extensions +# cannot be safely re-initialized after eviction, and the failure is +# order-dependent: it only shows up on the SECOND-and-later `_make_agent()` +# call in a run where nothing had already pulled in agent.py's full transitive +# import tree first. The sibling test file (test_agent_model_and_safety_unit.py) +# avoids this via its own top-level `from weightslab.trainer.services.data_service +# import ...`, which happens to import that whole tree before any stubbing runs. +# Importing the same module here for the same reason, not because this file +# needs DataService itself. +from weightslab.trainer.services.data_service import DataService # noqa: F401 + + +def _install_agent_dependency_stubs(): + stubs = { + "langchain_core": types.ModuleType("langchain_core"), + "langchain_core.prompts": types.ModuleType("langchain_core.prompts"), + } + stubs["langchain_core.prompts"].ChatPromptTemplate = object + return stubs + + +def _make_agent(df=None): + with mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): + agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") + + if df is None: + df = pd.DataFrame({"loss": [0.1, 0.9], "discarded": [False, False]}) + + ctx = SimpleNamespace(_all_datasets_df=df, _ctx=None) + agent = agent_mod.DataManipulationAgent(ctx) + + return agent_mod, agent + + +class TestOpenCodeConfigLoading(unittest.TestCase): + def test_opencode_url_and_model_default(self): + with mock.patch.dict("os.environ", {}, clear=False): + for key in ("OPENCODE_URL", "OPENCODE_MODEL"): + import os + os.environ.pop(key, None) + _, agent = _make_agent() + self.assertEqual(agent.opencode_url, "http://127.0.0.1:4096") + self.assertEqual(agent.opencode_model, "") + + def test_opencode_url_and_model_read_from_env(self): + with mock.patch.dict("os.environ", { + "OPENCODE_URL": "http://127.0.0.1:5555", + "OPENCODE_MODEL": "openrouter/anthropic/claude-opus-4.6", + }, clear=False): + _, agent = _make_agent() + self.assertEqual(agent.opencode_url, "http://127.0.0.1:5555") + self.assertEqual(agent.opencode_model, "openrouter/anthropic/claude-opus-4.6") + + +class TestSetupProvidersOpenCode(unittest.TestCase): + def test_opencode_chain_is_built(self): + agent_mod, agent = _make_agent() + fake_runnable = MagicMock() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = fake_runnable + initialized = agent._setup_providers() + + mock_cls.assert_called_once_with(agent.opencode_url, agent.opencode_model) + self.assertTrue(initialized) + self.assertIs(agent.chain_opencode, fake_runnable) + + def test_no_api_key_gate(self): + """OpenCode has no API-key concept at all -- the credential lives in + OpenCode's own config.""" + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = MagicMock() + initialized = agent._setup_providers() + self.assertTrue(initialized) # succeeded with no key configured anywhere + + def test_setup_error_is_caught_and_reported_as_not_initialized(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat", side_effect=RuntimeError("boom")): + initialized = agent._setup_providers() + self.assertFalse(initialized) + self.assertIsNone(agent.chain_opencode) + + +class TestInitializeWithCloudKeyOpenCode(unittest.TestCase): + def test_accepts_opencode_and_ignores_empty_api_key(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = MagicMock() + success, message = agent.initialize_with_cloud_key("", "opencode", "openrouter/openai/gpt-5") + + self.assertTrue(success) + self.assertIn("OpenCode", message) + self.assertEqual(agent.preferred_provider, "opencode") + self.assertEqual(agent.opencode_model, "openrouter/openai/gpt-5") + + def test_reports_failure_when_opencode_unreachable(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat", side_effect=RuntimeError("connection refused")): + success, message = agent.initialize_with_cloud_key("", "opencode", None) + self.assertFalse(success) + self.assertIn("OpenCode", message) + + def test_rejects_any_provider_other_than_opencode(self): + _, agent = _make_agent() + success, message = agent.initialize_with_cloud_key("key", "anthropic-direct", None) + self.assertFalse(success) + self.assertIn("Only OpenCode", message) + + def test_rejects_openrouter_now_that_it_is_removed(self): + _, agent = _make_agent() + success, message = agent.initialize_with_cloud_key("sk-or-test", "openrouter", "openai/gpt-5") + self.assertFalse(success) + self.assertIn("Only OpenCode", message) + + +class TestChangeModelOpenCode(unittest.TestCase): + def test_switches_opencode_model(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat") as mock_cls: + mock_cls.return_value.as_runnable.return_value = MagicMock() + success, message = agent.change_model("openrouter/openai/gpt-5-mini") + self.assertTrue(success) + self.assertEqual(agent.opencode_model, "openrouter/openai/gpt-5-mini") + + def test_reports_failure_when_opencode_unreachable(self): + agent_mod, agent = _make_agent() + with mock.patch.object(agent_mod, "OpenCodeChat", side_effect=RuntimeError("down")): + success, message = agent.change_model("openrouter/openai/gpt-5-mini") + self.assertFalse(success) + + def test_empty_model_is_rejected(self): + _, agent = _make_agent() + success, message = agent.change_model(" ") + self.assertFalse(success) + self.assertIn("Model cannot be empty", message) + + +class TestGetAvailableModelsOpenCode(unittest.TestCase): + def test_flattens_providers_into_provider_slash_model_strings(self): + _, agent = _make_agent() + agent.opencode_url = "http://127.0.0.1:4096" + + fake_payload = { + "providers": [ + {"id": "openrouter", "models": {"anthropic/claude-opus-4.6": {}, "openai/gpt-5": {}}}, + {"id": "ollama", "models": {"llama3.2:3b": {}}}, + ], + } + + class _FakeResp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + import json + return json.dumps(fake_payload).encode() + + with mock.patch("urllib.request.urlopen", return_value=_FakeResp()): + ok, models, message = agent.get_available_models() + + self.assertTrue(ok) + self.assertEqual( + models, + sorted([ + "openrouter/anthropic/claude-opus-4.6", + "openrouter/openai/gpt-5", + "ollama/llama3.2:3b", + ]), + ) + + def test_reports_failure_when_opencode_server_unreachable(self): + _, agent = _make_agent() + with mock.patch("urllib.request.urlopen", side_effect=OSError("refused")): + ok, models, message = agent.get_available_models() + self.assertFalse(ok) + self.assertEqual(models, []) + self.assertIn("OpenCode", message) + + +class TestResetConnection(unittest.TestCase): + def test_reset_clears_opencode_chain_and_model(self): + agent_mod, agent = _make_agent() + agent.chain_opencode = MagicMock() + agent.opencode_model = "openrouter/openai/gpt-5" + + success, message = agent.reset_connection() + + self.assertTrue(success) + self.assertIsNone(agent.chain_opencode) + self.assertEqual(agent.preferred_provider, "opencode") + + +class _FakePipedRunnable: + """Stands in for `(ChatPromptTemplate | chain)` -- skips actual prompt + formatting (irrelevant to what compact_history does with the result) and + just forwards straight to the underlying chain's `.invoke`, matching real + LangChain's RunnableSequence semantics for this purpose.""" + + def __init__(self, chain): + self._chain = chain + + def invoke(self, variables): + return self._chain.invoke(variables) + + +class _FakeChatPromptTemplate: + @classmethod + def from_messages(cls, messages): + return cls() + + def __or__(self, chain): + return _FakePipedRunnable(chain) + + +class TestClearAndCompactHistory(unittest.TestCase): + def test_clear_history_empties_and_reports_count(self): + _, agent = _make_agent() + agent.history = ["User: a", "Action: 1 ops executed", "User: b", "Action: 2 ops executed"] + + success, message = agent.clear_history() + + self.assertTrue(success) + self.assertEqual(agent.history, []) + self.assertIn("4", message) + + def test_compact_history_on_empty_history_is_a_no_op_success(self): + # Short-circuits before touching ChatPromptTemplate at all -- no patch needed. + _, agent = _make_agent() + agent.history = [] + success, message = agent.compact_history() + self.assertTrue(success) + self.assertEqual(agent.history, []) + + def test_compact_history_replaces_history_with_one_summary(self): + agent_mod, agent = _make_agent() + agent.history = ["User: discard bad samples", "Action: 3 ops executed"] + + fake_reply = SimpleNamespace(content="Discarded 3 low-quality samples per user request.") + agent.chain_opencode = MagicMock(invoke=MagicMock(return_value=fake_reply)) + + with mock.patch.object(agent_mod, "ChatPromptTemplate", _FakeChatPromptTemplate): + success, message = agent.compact_history() + + self.assertTrue(success) + self.assertEqual(len(agent.history), 1) + self.assertIn("Discarded 3 low-quality samples", agent.history[0]) + + def test_compact_history_fails_cleanly_when_no_provider_available(self): + agent_mod, agent = _make_agent() + agent.history = ["User: x"] + agent.chain_opencode = None + + with mock.patch.object(agent_mod, "ChatPromptTemplate", _FakeChatPromptTemplate): + success, message = agent.compact_history() + + self.assertFalse(success) + # History is left untouched on failure -- nothing was actually compacted. + self.assertEqual(agent.history, ["User: x"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/trainer/services/test_agent_prompt_unit.py b/tests/trainer/services/test_agent_prompt_unit.py index 116d7cdb..093d992a 100644 --- a/tests/trainer/services/test_agent_prompt_unit.py +++ b/tests/trainer/services/test_agent_prompt_unit.py @@ -10,13 +10,9 @@ def _install_agent_dependency_stubs(): stubs = { - "langchain_ollama": types.ModuleType("langchain_ollama"), - "langchain_openai": types.ModuleType("langchain_openai"), "langchain_core": types.ModuleType("langchain_core"), "langchain_core.prompts": types.ModuleType("langchain_core.prompts"), } - stubs["langchain_ollama"].ChatOllama = object - stubs["langchain_openai"].ChatOpenAI = object stubs["langchain_core.prompts"].ChatPromptTemplate = object return stubs @@ -37,18 +33,6 @@ def _rewrite_origin_literals(self, code): return code -class _FakeChatModel: - def __init__(self, *args, **kwargs): - self.args = args - self.kwargs = kwargs - - def with_structured_output(self, schema): - return self - - def invoke(self, prompt): - return SimpleNamespace(content="OK") - - class TestAgentPromptUnit(unittest.TestCase): def test_intent_prompt_contains_expected_placeholders(self): self.assertIn("{row_count}", INTENT_PROMPT) @@ -134,44 +118,7 @@ def test_agent_models_and_handlers(self): ) self.assertEqual(action["function"], "action.save") - def test_initialize_with_cloud_key_checks_chat_connectivity(self): - with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): - agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") - - ctx = SimpleNamespace( - _all_datasets_df=agent_mod.pd.DataFrame({"metric": [1.0, 2.0]}), - ) - - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key("test-key", "openrouter", "google/gemini-2.5-flash") - - self.assertTrue(ok) - self.assertIn("initialized successfully", message) - self.assertIsNotNone(agent.chain_openrouter) - self.assertEqual(agent.openrouter_model, "google/gemini-2.5-flash") - - def test_initialize_with_cloud_key_fails_when_probe_fails(self): - with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): - agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") - - class _FailingChatModel(_FakeChatModel): - def invoke(self, prompt): - raise RuntimeError("401 Unauthorized") - - ctx = SimpleNamespace( - _all_datasets_df=agent_mod.pd.DataFrame({"metric": [1.0, 2.0]}), - ) - - with mock.patch.object(agent_mod, "ChatOpenAI", _FailingChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key("bad-key", "openrouter", "~google/gemini-flash-latest") - - self.assertFalse(ok) - self.assertIn("connectivity check failed", message) - self.assertIsNone(agent.chain_openrouter) - - def test_initialize_with_cloud_key_rejects_non_openrouter_provider(self): + def test_initialize_with_cloud_key_rejects_non_opencode_provider(self): with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): agent_mod = importlib.import_module("weightslab.trainer.services.agent.agent") @@ -179,12 +126,11 @@ def test_initialize_with_cloud_key_rejects_non_openrouter_provider(self): _all_datasets_df=agent_mod.pd.DataFrame({"metric": [1.0, 2.0]}), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) - ok, message = agent.initialize_with_cloud_key("test-key", "grok", "grok-3-mini") + agent = agent_mod.DataManipulationAgent(ctx) + ok, message = agent.initialize_with_cloud_key("test-key", "grok", "grok-3-mini") self.assertFalse(ok) - self.assertIn("Only OpenRouter", message) + self.assertIn("Only OpenCode", message) def test_build_python_mask_keeps_string_literals(self): with unittest.mock.patch.dict(sys.modules, _install_agent_dependency_stubs(), clear=False): @@ -202,8 +148,7 @@ def test_build_python_mask_keeps_string_literals(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) mask = agent._build_python_mask( [agent_mod.Condition(column="origin", op="==", value="train")] @@ -234,8 +179,7 @@ def test_compact_schema_for_prompt_separates_index_levels_from_columns(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) schema_text = agent._compact_schema_for_prompt() @@ -321,8 +265,7 @@ def test_filter_by_origin_and_loss(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) cond1 = agent_mod.Condition(column="origin", op="==", value="train") cond2 = agent_mod.Condition(column="loss", op="<", value=0.3) @@ -368,8 +311,7 @@ def test_tag_high_loss_samples(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -401,8 +343,7 @@ def test_tag_from_quantile_computation(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -496,8 +437,7 @@ def test_analysis_train_loss_stddev(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="analysis", @@ -531,8 +471,7 @@ def test_tag_outliers_by_stddev(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -563,8 +502,7 @@ def test_tag_outliers_by_iqr(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -597,8 +535,7 @@ def test_filter_and_tag_combination(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent_mod.DataManipulationAgent(ctx) + agent_mod.DataManipulationAgent(ctx) # First filter: keep only train filt_cond = agent_mod.Condition(column="origin", op="==", value="train") @@ -635,8 +572,7 @@ def test_untag_operation(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", @@ -667,8 +603,7 @@ def test_rename_tag_operation(self): ), ) - with mock.patch.object(agent_mod, "ChatOpenAI", _FakeChatModel), mock.patch.object(agent_mod, "ChatOllama", _FakeChatModel): - agent = agent_mod.DataManipulationAgent(ctx) + agent = agent_mod.DataManipulationAgent(ctx) step = agent_mod.AtomicIntent( kind="transform", diff --git a/tests/trainer/services/test_agent_service_unit.py b/tests/trainer/services/test_agent_service_unit.py index fb26bd6c..36b8f5aa 100644 --- a/tests/trainer/services/test_agent_service_unit.py +++ b/tests/trainer/services/test_agent_service_unit.py @@ -24,9 +24,11 @@ def test_check_agent_health_reports_ready_when_available(self): self.assertTrue(response.available) self.assertIn('Ready to help you.', response.message) - def test_initialize_agent_delegates_to_agent_with_openrouter(self): + def test_initialize_agent_rejects_openrouter_now_that_it_is_removed(self): + """PROVIDER_OPENROUTER (0) is kept in the .proto only for wire + compatibility with older frontends -- requesting it must be rejected + cleanly rather than reaching the (now opencode-only) agent.""" agent = MagicMock() - agent.initialize_with_cloud_key.return_value = (True, 'ok') service, _ = self._make_service(agent=agent) response = service.InitializeAgent( @@ -38,13 +40,9 @@ def test_initialize_agent_delegates_to_agent_with_openrouter(self): None, ) - agent.initialize_with_cloud_key.assert_called_once_with( - 'sk-or-test', - 'openrouter', - '~google/gemini-flash-latest', - ) - self.assertTrue(response.success) - self.assertEqual(response.message, 'ok') + agent.initialize_with_cloud_key.assert_not_called() + self.assertFalse(response.success) + self.assertIn('Only OpenCode', response.message) def test_initialize_agent_rejects_unsupported_provider(self): agent = MagicMock() @@ -61,7 +59,7 @@ def test_initialize_agent_rejects_unsupported_provider(self): agent.initialize_with_cloud_key.assert_not_called() self.assertFalse(response.success) - self.assertIn('Only OpenRouter', response.message) + self.assertIn('Only OpenCode', response.message) def test_change_get_and_reset_agent_delegate_to_agent(self): agent = MagicMock() @@ -99,6 +97,56 @@ def test_methods_fail_cleanly_when_agent_backend_missing(self): self.assertEqual(list(list_response.models), []) self.assertIn('not running', init_response.message) + def test_initialize_agent_accepts_opencode_provider(self): + agent = MagicMock() + agent.initialize_with_cloud_key.return_value = (True, 'Agent initialized successfully via OpenCode. Ready to help you.') + service, _ = self._make_service(agent=agent) + + response = service.InitializeAgent( + pb2.InitializeAgentRequest( + api_key='', # ignored for opencode -- credential lives in OpenCode's own config + provider=pb2.PROVIDER_OPENCODE, + model='openrouter/anthropic/claude-opus-4.6', + ), + None, + ) + + agent.initialize_with_cloud_key.assert_called_once_with( + '', 'opencode', 'openrouter/anthropic/claude-opus-4.6', + ) + self.assertTrue(response.success) + + def test_clear_agent_history_delegates_to_agent(self): + agent = MagicMock() + agent.clear_history.return_value = (True, 'Cleared 4 history entries.') + service, _ = self._make_service(agent=agent) + + response = service.ClearAgentHistory(pb2.Empty(), None) + + agent.clear_history.assert_called_once_with() + self.assertTrue(response.success) + self.assertEqual(response.message, 'Cleared 4 history entries.') + + def test_compact_agent_history_delegates_to_agent(self): + agent = MagicMock() + agent.compact_history.return_value = (True, 'Compacted 4 entries into one summary.') + service, _ = self._make_service(agent=agent) + + response = service.CompactAgentHistory(pb2.Empty(), None) + + agent.compact_history.assert_called_once_with() + self.assertTrue(response.success) + + def test_clear_and_compact_history_fail_cleanly_when_agent_backend_missing(self): + service, _ = self._make_service(agent=None) + + clear_response = service.ClearAgentHistory(pb2.Empty(), None) + compact_response = service.CompactAgentHistory(pb2.Empty(), None) + + self.assertFalse(clear_response.success) + self.assertFalse(compact_response.success) + self.assertIn('not running', clear_response.message) + if __name__ == '__main__': unittest.main() diff --git a/tests/trainer/services/test_opencode_chat.py b/tests/trainer/services/test_opencode_chat.py new file mode 100644 index 00000000..52840b2c --- /dev/null +++ b/tests/trainer/services/test_opencode_chat.py @@ -0,0 +1,250 @@ +"""Tests for OpenCodeChat (weightslab/trainer/services/agent/opencode_chat.py). + +Exercises the module against a REAL minimal HTTP server implementing the +subset of OpenCode's protocol this class needs (POST /session, POST +/session/{id}/message, GET /event as text/event-stream) rather than mocking +urllib -- the class's correctness hinges on stream-first ordering and SSE +event parsing, which a mocked urlopen would not honestly exercise. Mirrors +the fake-server technique already used in tests/ui/test_server_agent.py for +the same reason. +""" + +import json +import threading +import time +import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +from weightslab.trainer.services.agent.opencode_chat import OpenCodeChat, OpenCodeError + + +class _FakeOpenCodeHandler(BaseHTTPRequestHandler): + """Implements just enough of OpenCode's HTTP+SSE surface to drive + OpenCodeChat through a full create -> send -> stream -> idle cycle. + Configured per-server-instance via class attributes the test sets before + starting it (see _FakeOpenCodeServer below).""" + + def log_message(self, *a): + pass + + def _send_json(self, obj): + body = json.dumps(obj).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_POST(self): + length = int(self.headers.get("Content-Length", "0") or 0) + raw = self.rfile.read(length) if length else b"" + try: + body = json.loads(raw.decode("utf-8")) if raw else {} + except ValueError: + body = {} + + if self.path == "/session": + self.server.recorded_session_titles.append(body.get("title")) + self._send_json({"id": self.server.session_id}) + return + + if self.path == f"/session/{self.server.session_id}/message": + self.server.recorded_messages.append(body) + # A real server holds this open for the whole turn; this fake + # returns immediately -- OpenCodeChat must not rely on this + # response for content, only on the SSE stream (see its own + # docstring). Sleep briefly so the message genuinely arrives + # after the stream has had a chance to open, exercising the + # stream-first ordering rather than accidentally passing by luck. + time.sleep(0.05) + self._send_json({}) + return + + self.send_response(404) + self.end_headers() + + def do_GET(self): + if self.path != "/event": + self.send_response(404) + self.end_headers() + return + + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + + def emit(event: dict) -> None: + payload = f"data: {json.dumps(event)}\n\n".encode("utf-8") + self.wfile.write(payload) + self.wfile.flush() + + session_id = self.server.session_id + # Wait for the message POST to actually land before emitting anything + # -- proves OpenCodeChat is genuinely reading a live stream, not just + # replaying a canned response. + deadline = time.monotonic() + 5 + while not self.server.recorded_messages and time.monotonic() < deadline: + time.sleep(0.01) + + emit({"type": "message.updated", "properties": {"info": {"id": "msg_1", "role": "assistant", "sessionID": session_id}}}) + for delta in self.server.reply_deltas: + emit({"type": "message.part.updated", "properties": {"part": { + "id": f"prt_{delta[:4]}", "type": "text", "text": delta, "messageID": "msg_1", "sessionID": session_id, + }}}) + if self.server.emit_error: + emit({"type": "session.error", "properties": {"sessionID": session_id}}) + else: + emit({"type": "session.idle", "properties": {"sessionID": session_id}}) + # Keep the connection open a moment so OpenCodeChat's break-on-idle has + # definitely already fired before we tear down. + time.sleep(0.05) + + +class _FakeOpenCodeServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, *args, reply_deltas, emit_error=False, **kwargs): + super().__init__(*args, **kwargs) + self.session_id = "ses_test1" + self.reply_deltas = reply_deltas + self.emit_error = emit_error + self.recorded_messages = [] + self.recorded_session_titles = [] + + +class _ServerTestCase(unittest.TestCase): + def _start_server(self, reply_deltas, emit_error=False): + self.httpd = _FakeOpenCodeServer( + ("127.0.0.1", 0), _FakeOpenCodeHandler, reply_deltas=reply_deltas, emit_error=emit_error, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + + def tearDown(self): + if hasattr(self, "httpd"): + self.httpd.shutdown() + self.thread.join(timeout=5) + + +class TestOpenCodeChatCall(_ServerTestCase): + def test_collects_streamed_text_and_returns_an_ai_message(self): + self._start_server(reply_deltas=["Here is ", "the answer."]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model="openrouter/openai/gpt-5", timeout=10) + + result = chat.as_runnable().invoke("do the thing") + + self.assertEqual(result.content, "Here is the answer.") + + def test_sends_the_model_ref_split_on_first_slash_only(self): + self._start_server(reply_deltas=["ok"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model="openrouter/anthropic/claude-opus-4.6", timeout=10) + + chat._call("hello") + + self.assertEqual(len(self.httpd.recorded_messages), 1) + self.assertEqual( + self.httpd.recorded_messages[0]["model"], + {"providerID": "openrouter", "modelID": "anthropic/claude-opus-4.6"}, + ) + + def test_disables_every_mutating_tool_on_the_outgoing_message(self): + """This wrapper backs the SDK agent's text/JSON call sites, which parse + the reply themselves -- unlike the Weights Studio landing chat, it must + never let OpenCode write files as a side effect.""" + self._start_server(reply_deltas=["ok"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("hello") + + tools = self.httpd.recorded_messages[0]["tools"] + self.assertEqual(tools, {"write": False, "edit": False, "patch": False, "bash": False}) + + def test_omits_model_field_when_none_configured(self): + self._start_server(reply_deltas=["ok"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("hello") + + self.assertNotIn("model", self.httpd.recorded_messages[0]) + + def test_creates_a_fresh_session_per_call(self): + self._start_server(reply_deltas=["a"]) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("first") + chat._call("second") + + # A fresh session per call (self.history on DataManipulationAgent + # already carries cross-call context) -- both calls hit /session, not + # a reused id. + self.assertEqual(len(self.httpd.recorded_session_titles), 2) + + def test_degrades_to_empty_string_on_session_error_rather_than_raising(self): + self._start_server(reply_deltas=["partial"], emit_error=True) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + # session.error still ends the read loop cleanly; whatever text arrived + # before the error is still returned rather than raising. + result = chat._call("hello") + self.assertEqual(result.content, "partial") + + def test_raises_opencode_error_when_the_server_is_unreachable(self): + chat = OpenCodeChat("http://127.0.0.1:1", model=None, timeout=2) # nothing listens on port 1 + with self.assertRaises(OpenCodeError): + chat._call("hello") + + +class TestModelRefParsing(unittest.TestCase): + def test_splits_on_first_slash_only(self): + chat = OpenCodeChat("http://x", model="openrouter/anthropic/claude-opus-4.6") + self.assertEqual(chat._model_ref(), {"providerID": "openrouter", "modelID": "anthropic/claude-opus-4.6"}) + + def test_none_when_no_model_configured(self): + chat = OpenCodeChat("http://x", model=None) + self.assertIsNone(chat._model_ref()) + + def test_none_when_model_has_no_slash(self): + chat = OpenCodeChat("http://x", model="justamodel") + self.assertIsNone(chat._model_ref()) + + +class TestHandleEvent(unittest.TestCase): + """Unit-level checks on the event state machine, independent of the + network -- complements the end-to-end server tests above.""" + + def test_ignores_parts_from_a_message_not_yet_known_to_be_assistant(self): + text_parts = {} + assistant_ids = set() + payload = json.dumps({ + "type": "message.part.updated", + "properties": {"part": {"id": "p1", "type": "text", "text": "x", "messageID": "msg_unknown", "sessionID": "s1"}}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", assistant_ids, text_parts) + self.assertIsNone(outcome) + self.assertEqual(text_parts, {}) + + def test_ignores_events_for_a_different_session(self): + text_parts = {} + assistant_ids = {"msg_1"} + payload = json.dumps({ + "type": "message.part.updated", + "properties": {"part": {"id": "p1", "type": "text", "text": "x", "messageID": "msg_1", "sessionID": "OTHER"}}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", assistant_ids, text_parts) + self.assertIsNone(outcome) + self.assertEqual(text_parts, {}) + + def test_malformed_payload_is_ignored_not_raised(self): + outcome = OpenCodeChat._handle_event("not json", "s1", set(), {}) + self.assertIsNone(outcome) + + def test_session_idle_signals_completion(self): + outcome = OpenCodeChat._handle_event( + json.dumps({"type": "session.idle", "properties": {"sessionID": "s1"}}), "s1", set(), {}, + ) + self.assertEqual(outcome, "idle") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ui/test_server_agent.py b/tests/ui/test_server_agent.py new file mode 100644 index 00000000..1d40e35a --- /dev/null +++ b/tests/ui/test_server_agent.py @@ -0,0 +1,238 @@ +"""Tests for weightslab/ui/server.py's OpenCode-agent supervisor: + +- POST /agent-server/start -- spawns (or reuses) a local OpenCode server rooted + at the experiment directory, so the browser never has to run `opencode serve` + by hand. Mirrors /local-notebook's "the browser can't spawn a process, so it + asks us to" shape. +- GET /agent-server/status -- none / running / killed, for the composer's status + line to poll. + +CI has no real `opencode` binary, so every test replaces +`ui_server._resolve_opencode_argv` with a tiny stand-in HTTP server (started via +`python -c`) that serves /global/health the same way the real OpenCode server +does. That is the only thing `_OpencodeSession` depends on to decide the child +started successfully, so it exercises the real spawn/health-poll/status code +path without depending on Node or the OpenCode package being installed. +""" + +import json +import os +import sys +import tempfile +import threading +import time +import unittest +import urllib.error +import urllib.request +from unittest.mock import patch + +from weightslab.ui import server as ui_server + +# A minimal stand-in for `opencode serve`: binds the --port it was given and +# answers /global/health like the real server does. Reads --port out of +# sys.argv positionally rather than assuming argv[0] is anything in particular, +# since "python -c serve --hostname H --port N" hands the child an argv +# whose exact shape depends on the platform's python launcher. +_FAKE_OPENCODE_SRC = r""" +import sys, json +from http.server import BaseHTTPRequestHandler, HTTPServer + +def _port(): + for i, a in enumerate(sys.argv): + if a == "--port": + return int(sys.argv[i + 1]) + return 4096 + +class H(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + def do_GET(self): + if self.path == "/global/health": + body = json.dumps({"version": "0.0.0-fake"}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + else: + self.send_response(404) + self.end_headers() + +HTTPServer(("127.0.0.1", _port()), H).serve_forever() +""" + +_FAKE_ARGV = [sys.executable, "-c", _FAKE_OPENCODE_SRC] + +# A stand-in that never becomes healthy -- models a child that starts but never +# binds (a bad flag, a crash loop) so ensure()'s timeout path is exercised. +_HANGING_ARGV = [sys.executable, "-c", "import time; time.sleep(60)"] + + +class TestOpencodeSessionUnit(unittest.TestCase): + """Exercises _OpencodeSession directly -- no HTTP layer, no real OpenCode.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.session = ui_server._OpencodeSession() + # Keep polling fast so a genuine failure test doesn't sit for 45s. + self._timeout_patch = patch.object(ui_server, "_OPENCODE_START_TIMEOUT", 3.0) + self._timeout_patch.start() + + def tearDown(self): + self.session.shutdown() + self._timeout_patch.stop() + + def test_starts_and_reports_running(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + self.assertFalse(result["reused"]) + self.assertEqual(result["workspace"], self.tmp) + self.assertTrue(result["url"].startswith("http://127.0.0.1:")) + + status = self.session.status() + self.assertEqual(status["state"], "running") + self.assertEqual(status["workspace"], self.tmp) + + def test_second_call_reuses_the_same_process(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + first = self.session.ensure(self.tmp, "http://localhost:5173") + second = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertFalse(first["reused"]) + self.assertTrue(second["reused"]) + self.assertEqual(first["url"], second["url"]) + + def test_missing_binary_and_missing_npx_reports_a_clear_error(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=None): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertFalse(result["ok"]) + self.assertIn("opencode", result["error"].lower()) + + status = self.session.status() + self.assertEqual(status["state"], "none") + + def test_child_that_never_becomes_healthy_times_out_and_is_killed(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_HANGING_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertFalse(result["ok"]) + self.assertIn("did not come up", result["error"]) + + # The timed-out child must actually be killed, not leaked. + status = self.session.status() + self.assertEqual(status["state"], "killed") + + def test_shutdown_stops_a_running_process(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"]) + process = self.session._process + self.session.shutdown() + self.assertIsNotNone(process.poll()) # exited + + +class TestCorsOriginVariants(unittest.TestCase): + """The localhost <-> 127.0.0.1 expansion is the #1 way this feature goes + silently wrong -- a mismatch here makes every request look like the agent + server is simply not there.""" + + def test_expands_localhost_to_127_0_0_1(self): + variants = ui_server._cors_origin_variants("http://localhost:5173") + self.assertIn("http://localhost:5173", variants) + self.assertIn("http://127.0.0.1:5173", variants) + + def test_expands_127_0_0_1_to_localhost(self): + variants = ui_server._cors_origin_variants("http://127.0.0.1:8080") + self.assertIn("http://127.0.0.1:8080", variants) + self.assertIn("http://localhost:8080", variants) + + def test_leaves_a_non_loopback_origin_alone(self): + variants = ui_server._cors_origin_variants("https://weightslab.example.com") + self.assertEqual(variants, ["https://weightslab.example.com"]) + + def test_none_origin_yields_no_variants(self): + self.assertEqual(ui_server._cors_origin_variants(None), []) + + +class _ServerTestCase(unittest.TestCase): + """Spins up a real serve_ui() on 127.0.0.1: per test, rooted at a + fresh temp dir passed as experiment_dir -- same shape as + test_server_experiment_reports.py.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.httpd = ui_server.serve_ui( + ui_host="127.0.0.1", ui_port=0, + backend_host="localhost", backend_port=50051, + open_browser=False, block=False, + experiment_dir=self.tmp, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + time.sleep(0.1) + + # The module-level singleton is shared across the whole test process; + # swap in a fresh one per test so a leftover child from another test + # can't make this one see {reused: true} unexpectedly. + self._orig_session = ui_server._opencode_session + ui_server._opencode_session = ui_server._OpencodeSession() + self._timeout_patch = patch.object(ui_server, "_OPENCODE_START_TIMEOUT", 3.0) + self._timeout_patch.start() + + def tearDown(self): + self._timeout_patch.stop() + ui_server._opencode_session.shutdown() + ui_server._opencode_session = self._orig_session + self.httpd.shutdown() + self.thread.join(timeout=5) + + def _get(self, path): + return urllib.request.urlopen(f"http://127.0.0.1:{self.port}{path}", timeout=5) + + def _post(self, path, origin=None): + req = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", method="POST", data=b"", + ) + if origin: + req.add_header("Origin", origin) + return urllib.request.urlopen(req, timeout=10) + + +class TestAgentServerEndpoint(_ServerTestCase): + + def test_status_is_none_before_anything_starts(self): + with self._get("/agent-server/status") as r: + data = json.loads(r.read().decode()) + self.assertEqual(data["state"], "none") + + def test_start_spawns_rooted_at_the_experiment_dir(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + with self._post("/agent-server/start", origin="http://localhost:5173") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + self.assertEqual(data["workspace"], self.tmp) + + with self._get("/agent-server/status") as r: + status = json.loads(r.read().decode()) + self.assertEqual(status["state"], "running") + + def test_missing_opencode_and_npx_returns_a_clear_error_not_a_500_crash(self): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=None): + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post("/agent-server/start", origin="http://localhost:5173") + self.assertEqual(ctx.exception.code, 500) + data = json.loads(ctx.exception.read().decode()) + self.assertFalse(data["ok"]) + self.assertIn("opencode", data["error"].lower()) + + def test_falls_back_to_reconstructing_origin_from_host_header(self): + # No Origin header at all (e.g. a same-origin fetch some browsers omit + # it for) -- must not crash, and must still start the server. + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + with self._post("/agent-server/start") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ui/test_server_loop.py b/tests/ui/test_server_loop.py new file mode 100644 index 00000000..0ca30391 --- /dev/null +++ b/tests/ui/test_server_loop.py @@ -0,0 +1,356 @@ +"""Tests for weightslab/ui/server.py's /loop feature -- recurring +OpenCode-backed monitoring jobs (`_LoopRegistry` + the +/agent-server/loop/{start,list,stop} endpoints). + +_LoopRegistry.start() delegates session/message plumbing to the module-level +_opencode_json_request/_opencode_send_and_collect helpers (already covered at +the wire-protocol level by tests/trainer/services/test_opencode_chat.py's +fake-SSE-server tests for the sibling Python client). Here those two helpers +are mocked so the tests exercise _LoopRegistry's OWN logic -- validation, +session reuse across ticks, error bookkeeping, stop/list, and the HTTP +wiring -- the same split test_server_agent.py uses (direct _OpencodeSession +unit tests, then a separate endpoint-level test class). +""" + +import json +import tempfile +import threading +import time +import unittest +import urllib.request +from unittest.mock import patch + +from weightslab.ui import server as ui_server + + +class TestLoopRegistryUnit(unittest.TestCase): + """Exercises _LoopRegistry directly, with the opencode session assumed + already up (_opencode_session.ensure mocked) and the session-create/ + send-and-collect wire calls mocked -- those are covered elsewhere (see + module docstring).""" + + def setUp(self): + self.registry = ui_server._LoopRegistry() + self._ensure_patch = patch.object( + ui_server._opencode_session, "ensure", + return_value={"ok": True, "url": "http://127.0.0.1:1", "workspace": "/tmp", "reused": False}, + ) + self._ensure_patch.start() + + def tearDown(self): + self.registry.shutdown() + self._ensure_patch.stop() + + def _wait_for_first_tick(self, job_id, timeout=2.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + jobs = {j["id"]: j for j in self.registry.list()} + job = jobs.get(job_id) + if job and (job["lastResult"] is not None or job["lastError"] is not None): + return job + time.sleep(0.01) + self.fail("loop job never completed its first tick") + + def test_rejects_an_empty_prompt(self): + result = self.registry.start(" ", 120, "/tmp", None) + self.assertFalse(result["ok"]) + self.assertIn("prompt", result["error"]) + self.assertEqual(self.registry.list(), []) + + def test_rejects_an_interval_below_the_minimum(self): + result = self.registry.start("watch training", 10, "/tmp", None) + self.assertFalse(result["ok"]) + self.assertIn("Minimum", result["error"]) + + def test_surfaces_an_ensure_failure_without_starting_a_job(self): + self._ensure_patch.stop() + try: + with patch.object(ui_server._opencode_session, "ensure", + return_value={"ok": False, "error": "no opencode binary"}): + result = self.registry.start("watch training", 120, "/tmp", None) + finally: + self._ensure_patch.start() + + self.assertFalse(result["ok"]) + self.assertEqual(result["error"], "no opencode binary") + self.assertEqual(self.registry.list(), []) + + def test_first_tick_creates_a_session_seeded_with_the_system_preamble(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}) as create_mock, \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="all good") as send_mock: + result = self.registry.start("watch the loss", 120, "/tmp", "http://localhost:5173") + self.assertTrue(result["ok"], result) + job = self._wait_for_first_tick(result["id"]) + + self.assertEqual(job["lastResult"], "all good") + self.assertIsNone(job["lastError"]) + create_mock.assert_called_once() + self.assertEqual(create_mock.call_args.args[1], "/session") + sent_text = send_mock.call_args.args[2] + self.assertIn("watch the loss", sent_text) + self.assertIn("weightslab pause", sent_text) # system preamble documents the CLI verbs + + def test_second_tick_reuses_the_session_and_sends_the_bare_prompt(self): + with patch.object(ui_server, "_LOOP_MIN_INTERVAL_SECONDS", 0.01), \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}) as create_mock, \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok") as send_mock: + result = self.registry.start("watch the loss", 0.02, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + deadline = time.monotonic() + 2 + while send_mock.call_count < 2 and time.monotonic() < deadline: + time.sleep(0.01) + self.registry.stop(result["id"]) + + self.assertGreaterEqual(send_mock.call_count, 2) + create_mock.assert_called_once() # session created once, reused on tick 2 + second_call_text = send_mock.call_args_list[1].args[2] + self.assertEqual(second_call_text, "watch the loss") # no preamble wrapper on repeat ticks + + def test_records_last_error_and_does_not_set_last_result_on_a_failed_tick(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", side_effect=RuntimeError("boom")): + result = self.registry.start("watch training", 120, "/tmp", None) + job = self._wait_for_first_tick(result["id"]) + + self.assertEqual(job["lastError"], "boom") + self.assertIsNone(job["lastResult"]) + + def test_stop_removes_the_job_and_prevents_a_further_tick(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok") as send_mock: + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + stop_result = self.registry.stop(result["id"]) + self.assertTrue(stop_result["ok"]) + self.assertEqual(self.registry.list(), []) + + calls_at_stop = send_mock.call_count + time.sleep(0.1) + self.assertEqual(send_mock.call_count, calls_at_stop) # no tick after stop + + def test_stopping_an_unknown_job_id_reports_not_found(self): + result = self.registry.stop("does-not-exist") + self.assertFalse(result["ok"]) + + def test_list_reflects_multiple_concurrent_jobs(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + first = self.registry.start("watch a", 120, "/tmp", None) + second = self.registry.start("watch b", 120, "/tmp", None) + self._wait_for_first_tick(first["id"]) + self._wait_for_first_tick(second["id"]) + + ids = {j["id"] for j in self.registry.list()} + self.assertEqual(ids, {first["id"], second["id"]}) + + def test_shutdown_stops_every_job(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + self.registry.shutdown() + self.assertEqual(self.registry.list(), []) + + def test_rejects_a_fourth_concurrent_job(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + for i in range(ui_server._LOOP_MAX_CONCURRENT): + result = self.registry.start(f"watch {i}", 120, "/tmp", None) + self.assertTrue(result["ok"], result) + + fourth = self.registry.start("one too many", 120, "/tmp", None) + + self.assertFalse(fourth["ok"]) + self.assertIn("already running", fourth["error"]) + self.assertEqual(len(self.registry.list()), ui_server._LOOP_MAX_CONCURRENT) + + def test_update_changes_the_prompt_without_touching_the_interval(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + update_result = self.registry.update(result["id"], prompt="watch training more closely") + + self.assertTrue(update_result["ok"], update_result) + job = {j["id"]: j for j in self.registry.list()}[result["id"]] + self.assertEqual(job["prompt"], "watch training more closely") + self.assertEqual(job["intervalSeconds"], 120) + + def test_update_changes_the_interval_and_reschedules_immediately(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok") as send_mock: + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + with patch.object(ui_server, "_LOOP_MIN_INTERVAL_SECONDS", 0.01): + update_result = self.registry.update(result["id"], interval_seconds=0.02) + self.assertTrue(update_result["ok"], update_result) + self.assertEqual(update_result["intervalSeconds"], 0.02) + + # If the reschedule took effect immediately, a second tick lands + # almost at once rather than after the original 120s interval. + deadline = time.monotonic() + 2 + while send_mock.call_count < 2 and time.monotonic() < deadline: + time.sleep(0.01) + self.registry.stop(result["id"]) + + self.assertGreaterEqual(send_mock.call_count, 2) + + def test_update_on_an_unknown_job_reports_not_found(self): + result = self.registry.update("does-not-exist", prompt="anything") + self.assertFalse(result["ok"]) + self.assertIn("No loop job", result["error"]) + + def test_update_rejects_an_empty_prompt(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + update_result = self.registry.update(result["id"], prompt=" ") + + self.assertFalse(update_result["ok"]) + self.assertIn("prompt", update_result["error"]) + + def test_update_rejects_an_interval_below_the_minimum(self): + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + update_result = self.registry.update(result["id"], interval_seconds=10) + + self.assertFalse(update_result["ok"]) + self.assertIn("Minimum", update_result["error"]) + + +class _LoopServerTestCase(unittest.TestCase): + """Spins up a real serve_ui() on 127.0.0.1:, same shape as + test_server_agent.py's _ServerTestCase -- reused rather than imported + since that one is a module-private helper of its own file.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.httpd = ui_server.serve_ui( + ui_host="127.0.0.1", ui_port=0, + backend_host="localhost", backend_port=50051, + open_browser=False, block=False, + experiment_dir=self.tmp, + ) + self.port = self.httpd.server_address[1] + self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + self.thread.start() + time.sleep(0.1) + + self._orig_registry = ui_server._loop_registry + ui_server._loop_registry = ui_server._LoopRegistry() + + def tearDown(self): + ui_server._loop_registry.shutdown() + ui_server._loop_registry = self._orig_registry + self.httpd.shutdown() + self.thread.join(timeout=5) + + def _get(self, path): + return urllib.request.urlopen(f"http://127.0.0.1:{self.port}{path}", timeout=5) + + def _post_json(self, path, body): + data = json.dumps(body).encode("utf-8") + req = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", method="POST", data=data, + headers={"Content-Type": "application/json"}, + ) + return urllib.request.urlopen(req, timeout=10) + + def _post(self, path): + req = urllib.request.Request(f"http://127.0.0.1:{self.port}{path}", method="POST", data=b"") + return urllib.request.urlopen(req, timeout=10) + + +class TestLoopEndpoints(_LoopServerTestCase): + + def test_list_is_empty_before_anything_starts(self): + with self._get("/agent-server/loop/list") as r: + data = json.loads(r.read().decode()) + self.assertEqual(data["loops"], []) + + def test_start_delegates_to_the_registry_and_the_new_job_shows_up_in_list(self): + with patch.object(ui_server._loop_registry, "start", + return_value={"ok": True, "id": "1", "intervalSeconds": 1800.0}) as start_mock: + with self._post_json("/agent-server/loop/start", {"prompt": "watch the loss", "intervalMinutes": 30}) as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + self.assertEqual(data["id"], "1") + start_mock.assert_called_once() + args = start_mock.call_args.args + self.assertEqual(args[0], "watch the loss") + self.assertEqual(args[1], 30 * 60.0) + self.assertEqual(args[2], self.tmp) # rooted at the experiment dir + + def test_start_with_an_empty_prompt_returns_a_400(self): + import urllib.error + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post_json("/agent-server/loop/start", {"prompt": "", "intervalMinutes": 30}) + self.assertEqual(ctx.exception.code, 400) + data = json.loads(ctx.exception.read().decode()) + self.assertFalse(data["ok"]) + + def test_stop_delegates_to_the_registry_with_the_id_parsed_out_of_the_path(self): + with patch.object(ui_server._loop_registry, "stop", return_value={"ok": True}) as stop_mock: + with self._post("/agent-server/loop/42/stop") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"]) + stop_mock.assert_called_once_with("42") + + def test_stopping_an_unknown_job_returns_a_404(self): + import urllib.error + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post("/agent-server/loop/does-not-exist/stop") + self.assertEqual(ctx.exception.code, 404) + + def test_update_delegates_to_the_registry_with_the_id_parsed_out_of_the_path(self): + with patch.object(ui_server._loop_registry, "update", + return_value={"ok": True, "id": "42", "prompt": "new prompt", "intervalSeconds": 300.0}) as update_mock: + with self._post_json("/agent-server/loop/42/update", {"prompt": "new prompt", "intervalMinutes": 5}) as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + update_mock.assert_called_once_with("42", prompt="new prompt", interval_seconds=5 * 60.0) + + def test_updating_an_unknown_job_returns_a_400(self): + import urllib.error + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post_json("/agent-server/loop/does-not-exist/update", {"prompt": "x"}) + self.assertEqual(ctx.exception.code, 400) + data = json.loads(ctx.exception.read().decode()) + self.assertFalse(data["ok"]) + + def test_list_reflects_a_real_start_end_to_end_through_the_registry(self): + with patch.object(ui_server._opencode_session, "ensure", + return_value={"ok": True, "url": "http://127.0.0.1:1", "workspace": self.tmp, "reused": False}), \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="all quiet"): + with self._post_json("/agent-server/loop/start", {"prompt": "watch training", "intervalMinutes": 30}) as r: + started = json.loads(r.read().decode()) + self.assertTrue(started["ok"], started) + + deadline = time.monotonic() + 2 + job = None + while time.monotonic() < deadline: + with self._get("/agent-server/loop/list") as r: + loops = json.loads(r.read().decode())["loops"] + job = next((j for j in loops if j["id"] == started["id"]), None) + if job and job["lastResult"]: + break + time.sleep(0.02) + + self.assertIsNotNone(job) + self.assertEqual(job["lastResult"], "all quiet") + self.assertEqual(job["prompt"], "watch training") + + +if __name__ == "__main__": + unittest.main() diff --git a/weightslab/backend/cli.py b/weightslab/backend/cli.py index 6e64fa3c..c994a48c 100644 --- a/weightslab/backend/cli.py +++ b/weightslab/backend/cli.py @@ -210,9 +210,9 @@ def _handle_command(cmd: str) -> Any: }, 'agent_examples': { 'check agent': 'agent status', - 'initialize openrouter': 'agent init --api-key sk-or-... --model openai/gpt-4o-mini --timeout 20', + 'initialize opencode': 'agent init --model openrouter/anthropic/claude-opus-4.6', 'list available models': 'agent models', - 'switch model': 'agent model ~google/gemini-flash-latest', + 'switch model': 'agent model openrouter/openai/gpt-5', 'query the agent': 'agent query discard all samples with loss > 5 and tag them as hard_examples', 'query shortcut': 'ask tag train samples with loss > 1.2 as goldset', } @@ -695,12 +695,12 @@ def _handle_command(cmd: str) -> Any: return { 'ok': True, 'available': available, - 'message': 'Agent available.' if available else 'Agent not configured. Use: agent init --api-key KEY [--model MODEL] [--timeout SEC]', + 'message': 'Agent available.' if available else 'Agent not configured. Use: agent init [--model MODEL]', 'commands': { 'agent status': 'Check whether the agent is available', - 'agent init': 'Initialize OpenRouter with API key, model, and optional timeout', - 'agent model': 'Switch the active OpenRouter model', - 'agent models': 'List available OpenRouter models', + 'agent init': 'Initialize the OpenCode agent backend, optionally with a model', + 'agent model': 'Switch the active OpenCode model', + 'agent models': 'List available OpenCode models', 'agent reset': 'Clear the active agent connection', 'agent query': 'Execute a natural-language query through the agent', }, @@ -720,8 +720,8 @@ def _handle_command(cmd: str) -> Any: 'ok': True, 'available': available, 'preferred_provider': getattr(agent, 'preferred_provider', None), - 'openrouter_model': getattr(agent, 'openrouter_model', None), - 'openrouter_timeout': getattr(agent, 'openrouter_request_timeout', None), + 'opencode_url': getattr(agent, 'opencode_url', None), + 'opencode_model': getattr(agent, 'opencode_model', None), 'message': 'Agent available. Ready to help you.' if available else 'Agent not configured. Use agent init.', } @@ -729,51 +729,23 @@ def _handle_command(cmd: str) -> Any: if agent is None: return {'ok': False, 'error': 'agent_unavailable', 'message': 'Agent backend is not attached to the CLI.'} - api_key = None - provider = 'openrouter' model = None - timeout = None i = 0 while i < len(agent_parts): token = agent_parts[i] - if token == '--api-key' and i + 1 < len(agent_parts): - api_key = agent_parts[i + 1] - i += 2 - elif token == '--provider' and i + 1 < len(agent_parts): - provider = agent_parts[i + 1] - i += 2 - elif token == '--model' and i + 1 < len(agent_parts): + if token == '--model' and i + 1 < len(agent_parts): model = agent_parts[i + 1] i += 2 - elif token == '--timeout' and i + 1 < len(agent_parts): - timeout = agent_parts[i + 1] - i += 2 else: - return {'ok': False, 'error': 'usage: agent init --api-key KEY [--provider openrouter] [--model MODEL] [--timeout SEC]'} - - api_key = api_key or os.environ.get('OPENROUTER_API_KEY') - if not api_key: - return {'ok': False, 'error': 'usage: agent init --api-key KEY [--provider openrouter] [--model MODEL] [--timeout SEC]'} - - if timeout is not None: - try: - timeout_value = float(timeout) - except ValueError: - return {'ok': False, 'error': 'Invalid --timeout value'} - os.environ['OPENROUTER_REQUEST_TIMEOUT'] = str(timeout_value) - try: - agent.openrouter_request_timeout = timeout_value - except Exception: - pass + return {'ok': False, 'error': 'usage: agent init [--model MODEL]'} - success, message = agent.initialize_with_cloud_key(api_key, provider, model) + success, message = agent.initialize_with_cloud_key("", "opencode", model) return { 'ok': success, 'message': message, - 'provider': provider, - 'model': getattr(agent, 'openrouter_model', model), - 'timeout': getattr(agent, 'openrouter_request_timeout', None), + 'provider': 'opencode', + 'model': getattr(agent, 'opencode_model', model), } if subverb in ('model', 'set-model'): @@ -783,7 +755,7 @@ def _handle_command(cmd: str) -> Any: return {'ok': False, 'error': 'usage: agent model '} model = ' '.join(agent_parts).strip() success, message = agent.change_model(model) - return {'ok': success, 'message': message, 'model': getattr(agent, 'openrouter_model', model)} + return {'ok': success, 'message': message, 'model': getattr(agent, 'opencode_model', model)} if subverb in ('models', 'list-models'): if agent is None: diff --git a/weightslab/proto/experiment_service.proto b/weightslab/proto/experiment_service.proto index e926ab55..6c686348 100644 --- a/weightslab/proto/experiment_service.proto +++ b/weightslab/proto/experiment_service.proto @@ -40,6 +40,13 @@ service ExperimentService { rpc ChangeAgentModel (ChangeAgentModelRequest) returns (ChangeAgentModelResponse); rpc GetAgentModels (GetAgentModelsRequest) returns (GetAgentModelsResponse); rpc ResetAgent (Empty) returns (ResetAgentResponse); + // Wipe the agent's conversation history (self.history) without touching the + // provider connection -- distinct from ResetAgent, which drops the connection. + rpc ClearAgentHistory (Empty) returns (ClearAgentHistoryResponse); + // Summarize the agent's conversation history via the active model, replacing + // it with the summary. Distinct from OpenCode's own session compaction (which + // this does not touch) -- this is the SDK agent's own self.history. + rpc CompactAgentHistory (Empty) returns (CompactAgentHistoryResponse); // Notebook (shared in-process Python kernel for the studio UI). Cell output is // server-streamed in chunks (stdout / stderr / result / image), mirroring @@ -576,6 +583,12 @@ message AgentHealthResponse { // --- Agent Onboarding --- enum AgentProviderType { PROVIDER_OPENROUTER = 0; + // A local OpenCode server (opencode.ai) backs the LLM calls instead of a + // direct OpenRouter connection. No api_key is used for this provider -- the + // credential lives in OpenCode's own config, entered once via `opencode auth + // login` or the Weights Studio landing page's login modal. `model` carries + // OpenCode's "providerID/modelID" string (e.g. "openrouter/anthropic/claude-opus-4.6"). + PROVIDER_OPENCODE = 1; } message InitializeAgentRequest { @@ -611,6 +624,16 @@ message ResetAgentResponse { string message = 2; } +message ClearAgentHistoryResponse { + bool success = 1; + string message = 2; +} + +message CompactAgentHistoryResponse { + bool success = 1; + string message = 2; +} + // --- Checkpoint Restore --- message RestoreCheckpointRequest { string experiment_hash = 1; // Hash code of checkpoint to restore diff --git a/weightslab/proto/experiment_service_pb2.py b/weightslab/proto/experiment_service_pb2.py index d299a943..9adf459e 100644 --- a/weightslab/proto/experiment_service_pb2.py +++ b/weightslab/proto/experiment_service_pb2.py @@ -1,8 +1,8 @@ # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE -# source: weightslab/proto/experiment_service.proto -# Protobuf Python Version: 6.31.1 +# source: experiment_service.proto +# Protobuf Python Version: 5.28.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool @@ -11,11 +11,11 @@ from google.protobuf.internal import builder as _builder _runtime_version.ValidateProtobufRuntimeVersion( _runtime_version.Domain.PUBLIC, - 6, - 31, + 5, + 28, 1, '', - 'weightslab/proto/experiment_service.proto' + 'experiment_service.proto' ) # @@protoc_insertion_point(imports) @@ -24,11 +24,11 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n)weightslab/proto/experiment_service.proto\"\x89\x01\n\x1aGetLatestLoggerDataRequest\x12\x1c\n\x14request_full_history\x18\x01 \x01(\x08\x12\x12\n\nmax_points\x18\x02 \x01(\x05\x12\x17\n\x0f\x62reak_by_slices\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\"\x81\x02\n\x0fLoggerDataPoint\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x11\n\tmodel_age\x18\x02 \x01(\x05\x12\x14\n\x0cmetric_value\x18\x03 \x01(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x04 \x01(\t\x12\x11\n\ttimestamp\x18\x05 \x01(\x03\x12\x11\n\tsample_id\x18\x06 \x01(\t\x12\x1c\n\x14is_evaluation_marker\x18\x07 \x01(\x08\x12\x12\n\nsplit_name\x18\x08 \x01(\t\x12\x17\n\x0f\x65valuation_tags\x18\t \x03(\t\x12\x12\n\npoint_note\x18\n \x01(\t\x12\x12\n\naudit_mode\x18\x0b \x01(\x08\"[\n\x1bGetLatestLoggerDataResponse\x12 \n\x06points\x18\x01 \x03(\x0b\x32\x10.LoggerDataPoint\x12\x1a\n\x12weightslab_version\x18\x02 \x01(\t\"\x07\n\x05\x45mpty\"/\n\x08NeuronId\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tneuron_id\x18\x02 \x01(\x05\"\x91\x02\n\x0fWeightOperation\x12*\n\x07op_type\x18\x01 \x01(\x0e\x32\x14.WeightOperationTypeH\x00\x88\x01\x01\x12\x15\n\x08layer_id\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1d\n\nneuron_ids\x18\x03 \x03(\x0b\x32\t.NeuronId\x12\x16\n\x0eneurons_to_add\x18\t \x01(\x05\x12 \n\x18zerofy_from_incoming_ids\x18\x0b \x03(\x05\x12\x1c\n\x14zerofy_to_neuron_ids\x18\x0c \x03(\x05\x12+\n\x11zerofy_predicates\x18\r \x03(\x0e\x32\x10.ZerofyPredicateB\n\n\x08_op_typeB\x0b\n\t_layer_id\"_\n\x17WeightsOperationRequest\x12/\n\x10weight_operation\x18\x01 \x01(\x0b\x32\x10.WeightOperationH\x00\x88\x01\x01\x42\x13\n\x11_weight_operation\"<\n\x18WeightsOperationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xc1\x05\n\x0fHyperParameters\x12\x1c\n\x0f\x65xperiment_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12!\n\x14training_steps_to_do\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x17\n\nbatch_size\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12 \n\x13\x66ull_eval_frequency\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12 \n\x13\x63heckpont_frequency\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x18\n\x0bis_training\x18\x07 \x01(\x08H\x06\x88\x01\x01\x12\x15\n\x08nb_steps\x18\x08 \x01(\x05H\x07\x88\x01\x01\x12\x19\n\x0c\x61uditor_mode\x18\t \x01(\x08H\x08\x88\x01\x01\x12\x1d\n\x10train_batch_size\x18\n \x01(\x05H\t\x88\x01\x01\x12\x1b\n\x0eval_batch_size\x18\x0b \x01(\x05H\n\x88\x01\x01\x12\x1c\n\x0ftest_batch_size\x18\x0c \x01(\x05H\x0b\x88\x01\x01\x12\x1c\n\x0f\x65valuation_mode\x18\r \x01(\x08H\x0c\x88\x01\x01\x12\x1e\n\x11\x65valuation_config\x18\x0e \x01(\tH\r\x88\x01\x01\x42\x12\n\x10_experiment_nameB\x17\n\x15_training_steps_to_doB\x10\n\x0e_learning_rateB\r\n\x0b_batch_sizeB\x16\n\x14_full_eval_frequencyB\x16\n\x14_checkpont_frequencyB\x0e\n\x0c_is_trainingB\x0b\n\t_nb_stepsB\x0f\n\r_auditor_modeB\x13\n\x11_train_batch_sizeB\x11\n\x0f_val_batch_sizeB\x12\n\x10_test_batch_sizeB\x12\n\x10_evaluation_modeB\x14\n\x12_evaluation_config\",\n\rMetricsStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"~\n\rAnnotatStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12.\n\x08metadata\x18\x02 \x03(\x0b\x32\x1c.AnnotatStatus.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x90\x02\n\x10TrainingStatusEx\x12\x16\n\ttimestamp\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0f\x65xperiment_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x16\n\tmodel_age\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12+\n\x0emetrics_status\x18\x04 \x01(\x0b\x32\x0e.MetricsStatusH\x03\x88\x01\x01\x12+\n\x0e\x61nnotat_status\x18\x05 \x01(\x0b\x32\x0e.AnnotatStatusH\x04\x88\x01\x01\x42\x0c\n\n_timestampB\x12\n\x10_experiment_nameB\x0c\n\n_model_ageB\x11\n\x0f_metrics_statusB\x11\n\x0f_annotat_status\"]\n\x15HyperParameterCommand\x12/\n\x10hyper_parameters\x18\x01 \x01(\x0b\x32\x10.HyperParametersH\x00\x88\x01\x01\x42\x13\n\x11_hyper_parameters\">\n\x14\x44\x65nySamplesOperation\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\"0\n\x17LoadCheckpointOperation\x12\x15\n\rcheckpoint_id\x18\x01 \x01(\x05\"b\n\x11PlotNoteOperation\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x0c\n\x04note\x18\x04 \x01(\t\"L\n\x17SaveCheckpointOperation\x12\x19\n\x11save_architecture\x18\x01 \x01(\x08\x12\x16\n\x0esave_optimizer\x18\x02 \x01(\x08\"\x1a\n\x18RestartInstanceOperation\"\x8d\x08\n\x0eTrainerCommand\x12\x1c\n\x14get_hyper_parameters\x18\x04 \x01(\x08\x12\x1e\n\x16get_interactive_layers\x18\x05 \x01(\x08\x12\x1d\n\x10get_data_records\x18\x06 \x01(\tH\x00\x88\x01\x01\x12%\n\x18get_single_layer_info_id\x18\x08 \x01(\x05H\x01\x88\x01\x01\x12;\n\x16hyper_parameter_change\x18\x01 \x01(\x0b\x32\x16.HyperParameterCommandH\x02\x88\x01\x01\x12:\n\x16\x64\x65ny_samples_operation\x18\x07 \x01(\x0b\x32\x15.DenySamplesOperationH\x03\x88\x01\x01\x12?\n\x1b\x64\x65ny_eval_samples_operation\x18\n \x01(\x0b\x32\x15.DenySamplesOperationH\x04\x88\x01\x01\x12@\n\x19load_checkpoint_operation\x18\t \x01(\x0b\x32\x18.LoadCheckpointOperationH\x05\x88\x01\x01\x12\x42\n\x1eremove_from_denylist_operation\x18\x0b \x01(\x0b\x32\x15.DenySamplesOperationH\x06\x88\x01\x01\x12G\n#remove_eval_from_denylist_operation\x18\x0c \x01(\x0b\x32\x15.DenySamplesOperationH\x07\x88\x01\x01\x12\x34\n\x13plot_note_operation\x18\r \x01(\x0b\x32\x12.PlotNoteOperationH\x08\x88\x01\x01\x12@\n\x19save_checkpoint_operation\x18\x0e \x01(\x0b\x32\x18.SaveCheckpointOperationH\t\x88\x01\x01\x12\x39\n\x11restart_operation\x18\x0f \x01(\x0b\x32\x19.RestartInstanceOperationH\n\x88\x01\x01\x42\x13\n\x11_get_data_recordsB\x1b\n\x19_get_single_layer_info_idB\x19\n\x17_hyper_parameter_changeB\x19\n\x17_deny_samples_operationB\x1e\n\x1c_deny_eval_samples_operationB\x1c\n\x1a_load_checkpoint_operationB!\n\x1f_remove_from_denylist_operationB&\n$_remove_eval_from_denylist_operationB\x16\n\x14_plot_note_operationB\x1c\n\x1a_save_checkpoint_operationB\x14\n\x12_restart_operation\"\x9d\x01\n\x12HyperParameterDesc\x12\r\n\x05label\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04type\x18\x03 \x01(\t\x12\x1c\n\x0fnumerical_value\x18\x04 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cstring_value\x18\x05 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_numerical_valueB\x0f\n\r_string_value\"\xf2\x02\n\x10NeuronStatistics\x12!\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronIdH\x00\x88\x01\x01\x12\x17\n\nneuron_age\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1f\n\x12train_trigger_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x1e\n\x11\x65val_trigger_rate\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x07 \x01(\x02H\x04\x88\x01\x01\x12\x36\n\x0bincoming_lr\x18\x08 \x03(\x0b\x32!.NeuronStatistics.IncomingLrEntry\x1a\x31\n\x0fIncomingLrEntry\x12\x0b\n\x03key\x18\x01 \x01(\x05\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\x42\x0c\n\n_neuron_idB\r\n\x0b_neuron_ageB\x15\n\x13_train_trigger_rateB\x14\n\x12_eval_trigger_rateB\x10\n\x0e_learning_rate\"\xf0\x02\n\x13LayerRepresentation\x12\x15\n\x08layer_id\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rneurons_count\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12#\n\x16incoming_neurons_count\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x13\n\x06stride\x18\x07 \x01(\x05H\x06\x88\x01\x01\x12-\n\x12neurons_statistics\x18\n \x03(\x0b\x32\x11.NeuronStatisticsB\x0b\n\t_layer_idB\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x10\n\x0e_neurons_countB\x19\n\x17_incoming_neurons_countB\x0e\n\x0c_kernel_sizeB\t\n\x07_stride\"H\n\x11\x41\x63tivationRequest\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tsample_id\x18\x02 \x01(\t\x12\x0e\n\x06origin\x18\x03 \x01(\t\"H\n\rActivationMap\x12\x11\n\tneuron_id\x18\x01 \x01(\x05\x12\x0e\n\x06values\x18\x02 \x03(\x02\x12\t\n\x01H\x18\x03 \x01(\x05\x12\t\n\x01W\x18\x04 \x01(\x05\"d\n\x12\x41\x63tivationResponse\x12\x12\n\nlayer_type\x18\x01 \x01(\t\x12\x15\n\rneurons_count\x18\x02 \x01(\x05\x12#\n\x0b\x61\x63tivations\x18\x03 \x03(\x0b\x32\x0e.ActivationMap\"\x93\x01\n\tTaskField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x15\n\x0b\x66loat_value\x18\x02 \x01(\x02H\x00\x12\x13\n\tint_value\x18\x03 \x01(\x05H\x00\x12\x16\n\x0cstring_value\x18\x04 \x01(\tH\x00\x12\x15\n\x0b\x62ytes_value\x18\x05 \x01(\x0cH\x00\x12\x14\n\nbool_value\x18\x06 \x01(\x08H\x00\x42\x07\n\x05value\"\x87\x03\n\x0eRecordMetadata\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x14\n\x0csample_label\x18\x02 \x03(\x05\x12\x19\n\x11sample_prediction\x18\x03 \x03(\x05\x12=\n\x10sample_last_loss\x18\x04 \x03(\x0b\x32#.RecordMetadata.SampleLastLossEntry\x12\x19\n\x11sample_encounters\x18\x05 \x01(\x05\x12\x18\n\x10sample_discarded\x18\x06 \x01(\x08\x12 \n\x0c\x65xtra_fields\x18\x07 \x03(\x0b\x32\n.TaskField\x12\x16\n\x0eprediction_raw\x18\t \x01(\x0c\x12\x11\n\ttask_type\x18\n \x01(\t\x12\x19\n\x11sample_label_text\x18\x0b \x03(\t\x12\x1e\n\x16sample_prediction_text\x18\x0c \x03(\t\x1a\x35\n\x13SampleLastLossEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x93\x01\n\x10SampleStatistics\x12\x13\n\x06origin\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0csample_count\x18\x07 \x01(\x05H\x01\x88\x01\x01\x12\x11\n\ttask_type\x18\t \x01(\t\x12 \n\x07records\x18\x08 \x03(\x0b\x32\x0f.RecordMetadataB\t\n\x07_originB\x0f\n\r_sample_count\"\xe6\x01\n\x0f\x43ommandResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x33\n\x16hyper_parameters_descs\x18\x03 \x03(\x0b\x32\x13.HyperParameterDesc\x12\x33\n\x15layer_representations\x18\x04 \x03(\x0b\x32\x14.LayerRepresentation\x12\x31\n\x11sample_statistics\x18\x05 \x01(\x0b\x32\x11.SampleStatisticsH\x00\x88\x01\x01\x42\x14\n\x12_sample_statistics\"U\n\rSampleRequest\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_origin\"\xad\x02\n\x15SampleRequestResponse\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05label\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x64\x61ta\x18\x04 \x01(\x0cH\x03\x88\x01\x01\x12\x1a\n\rerror_message\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x15\n\x08raw_data\x18\x06 \x01(\x0cH\x05\x88\x01\x01\x12\x11\n\x04mask\x18\x07 \x01(\x0cH\x06\x88\x01\x01\x12\x17\n\nprediction\x18\x08 \x01(\x0cH\x07\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_originB\x08\n\x06_labelB\x07\n\x05_dataB\x10\n\x0e_error_messageB\x0b\n\t_raw_dataB\x07\n\x05_maskB\r\n\x0b_prediction\"\x92\x01\n\x12\x42\x61tchSampleRequest\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x19\n\x0cresize_width\x18\x03 \x01(\x05H\x00\x88\x01\x01\x12\x1a\n\rresize_height\x18\x04 \x01(\x05H\x01\x88\x01\x01\x42\x0f\n\r_resize_widthB\x10\n\x0e_resize_height\">\n\x13\x42\x61tchSampleResponse\x12\'\n\x07samples\x18\x01 \x03(\x0b\x32\x16.SampleRequestResponse\".\n\x0eWeightsRequest\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\"\x9d\x02\n\x0fWeightsResponse\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x10\n\x08incoming\x18\x04 \x01(\x05\x12\x10\n\x08outgoing\x18\x05 \x01(\x05\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x0f\n\x07weights\x18\x07 \x03(\x02\x12\x0f\n\x07success\x18\x0b \x01(\x08\x12\x1a\n\rerror_message\x18\x0c \x01(\tH\x03\x88\x01\x01\x42\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x0e\n\x0c_kernel_sizeB\x10\n\x0e_error_message\"R\n\x10\x44\x61taQueryRequest\x12\r\n\x05query\x18\x01 \x01(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\x12\x1b\n\x13is_natural_language\x18\x03 \x01(\x08\"5\n\x11\x43\x61tegoricalTagDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\ncategories\x18\x02 \x03(\t\"\xa9\x02\n\x11\x44\x61taQueryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1d\n\x15number_of_all_samples\x18\x03 \x01(\x05\x12%\n\x1dnumber_of_samples_in_the_loop\x18\x04 \x01(\x05\x12#\n\x1bnumber_of_discarded_samples\x18\x05 \x01(\x05\x12\x13\n\x0bunique_tags\x18\x06 \x03(\t\x12+\n\x11\x61gent_intent_type\x18\x07 \x01(\x0e\x32\x10.AgentIntentType\x12\x17\n\x0f\x61nalysis_result\x18\x08 \x01(\t\x12,\n\x10\x63\x61tegorical_tags\x18\t \x03(\x0b\x32\x12.CategoricalTagDef\"\xc2\x01\n\x12\x44\x61taSamplesRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12 \n\x18include_transformed_data\x18\x03 \x01(\x08\x12\x18\n\x10include_raw_data\x18\x04 \x01(\x08\x12\x19\n\x11stats_to_retrieve\x18\x05 \x03(\t\x12\x14\n\x0cresize_width\x18\x06 \x01(\x05\x12\x15\n\rresize_height\x18\x07 \x01(\x05\"m\n\x08\x44\x61taStat\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\r\n\x05shape\x18\x03 \x03(\x05\x12\r\n\x05value\x18\x04 \x03(\x02\x12\x14\n\x0cvalue_string\x18\x05 \x01(\t\x12\x11\n\tthumbnail\x18\x06 \x01(\x0c\">\n\nDataRecord\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x1d\n\ndata_stats\x18\x02 \x03(\x0b\x32\t.DataStat\"Z\n\x13\x44\x61taSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12!\n\x0c\x64\x61ta_records\x18\x03 \x03(\x0b\x32\x0b.DataRecord\"C\n\x0fHistogramSubBar\x12\x0e\n\x06origin\x18\x01 \x01(\t\x12\x11\n\tdiscarded\x18\x02 \x01(\x08\x12\r\n\x05\x63ount\x18\x03 \x01(\x03\"h\n\x0cHistogramBin\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03\x61vg\x18\x03 \x01(\x01\x12\r\n\x05\x63ount\x18\x04 \x01(\x03\x12\"\n\x08sub_bars\x18\x05 \x03(\x0b\x32\x10.HistogramSubBar\"[\n\x17\x43\x61tegoricalHistogramBar\x12\r\n\x05label\x18\x01 \x01(\t\x12\r\n\x05\x63ount\x18\x02 \x01(\x03\x12\"\n\x08sub_bars\x18\x03 \x03(\x0b\x32\x10.HistogramSubBar\"4\n\x10HistogramRequest\x12\x0e\n\x06\x63olumn\x18\x01 \x01(\t\x12\x10\n\x08max_bins\x18\x02 \x01(\x05\"\xb2\x01\n\x11HistogramResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\ntotal_rows\x18\x03 \x01(\x03\x12\x1b\n\x04\x62ins\x18\x04 \x03(\x0b\x32\r.HistogramBin\x12\x16\n\x0eis_categorical\x18\x05 \x01(\x08\x12\x32\n\x10\x63\x61tegorical_bars\x18\x06 \x03(\x0b\x32\x18.CategoricalHistogramBar\"W\n\x12GetMetaDataRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12\x17\n\x0fmodal_sample_id\x18\x03 \x01(\t\"\x99\x01\n\x13GetMetaDataResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1a\n\x12\x61ll_metadata_names\x18\x03 \x03(\t\x12!\n\x0cgrid_records\x18\x04 \x03(\x0b\x32\x0b.DataRecord\x12!\n\x0cmodal_record\x18\x05 \x01(\x0b\x32\x0b.DataRecord\"Y\n\x1aGetSignalTrajectoryRequest\x12\x13\n\x0bsignal_name\x18\x01 \x01(\t\x12\x12\n\nsample_ids\x18\x02 \x03(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"4\n\x10SignalTrajectory\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x03(\x02\"}\n\x1bGetSignalTrajectoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x13\n\x0bsignal_name\x18\x03 \x01(\t\x12\'\n\x0ctrajectories\x18\x04 \x03(\x0b\x32\x11.SignalTrajectory\"J\n\x11PointCloudRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"\xbf\x01\n\x0fPointCloudChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nnum_points\x18\x03 \x01(\x05\x12\x14\n\x0cnum_features\x18\x04 \x01(\x05\x12\x10\n\x08pc_range\x18\x05 \x03(\x02\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\x07 \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x08 \x01(\x05\x12\x15\n\rfeature_names\x18\t \x03(\t\"\xdc\x01\n\x10\x44\x61taEditsRequest\x12\x11\n\tstat_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66loat_value\x18\x02 \x01(\x02\x12\x14\n\x0cstring_value\x18\x03 \x01(\t\x12\x12\n\nbool_value\x18\x04 \x01(\x08\x12\x1d\n\x04type\x18\x05 \x01(\x0e\x32\x0f.SampleEditType\x12\x13\n\x0bsamples_ids\x18\x06 \x03(\t\x12\x16\n\x0esample_origins\x18\x07 \x03(\t\x12\x16\n\x0eis_categorical\x18\x08 \x01(\x08\x12\x12\n\ncategories\x18\t \x03(\t\"5\n\x11\x44\x61taEditsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\":\n\x12\x44\x61taSplitsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x13\n\x0bsplit_names\x18\x02 \x03(\t\"9\n\x13\x41gentHealthResponse\x12\x11\n\tavailable\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"^\n\x16InitializeAgentRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12$\n\x08provider\x18\x02 \x01(\x0e\x32\x12.AgentProviderType\x12\r\n\x05model\x18\x03 \x01(\t\";\n\x17InitializeAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"(\n\x17\x43hangeAgentModelRequest\x12\r\n\x05model\x18\x01 \x01(\t\"<\n\x18\x43hangeAgentModelResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x17\n\x15GetAgentModelsRequest\"J\n\x16GetAgentModelsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0e\n\x06models\x18\x02 \x03(\t\x12\x0f\n\x07message\x18\x03 \x01(\t\"6\n\x12ResetAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"3\n\x18RestoreCheckpointRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\"=\n\x19RestoreCheckpointResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"R\n\x18TriggerEvaluationRequest\x12\x12\n\nsplit_name\x18\x01 \x01(\t\x12\x0c\n\x04tags\x18\x02 \x03(\t\x12\x14\n\x0cuse_full_set\x18\x03 \x01(\x08\"=\n\x19TriggerEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x1c\n\x1aGetEvaluationStatusRequest\"\x81\x01\n\x1bGetEvaluationStatusResponse\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x0f\n\x07\x63urrent\x18\x02 \x01(\x05\x12\r\n\x05total\x18\x03 \x01(\x05\x12\x0f\n\x07message\x18\x04 \x01(\t\x12\r\n\x05\x65rror\x18\x05 \x01(\t\x12\x12\n\nsplit_name\x18\x06 \x01(\t\")\n\x17\x43\x61ncelEvaluationRequest\x12\x0e\n\x06reason\x18\x01 \x01(\t\"<\n\x18\x43\x61ncelEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"7\n\x16RunNotebookCellRequest\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07\x63\x65ll_id\x18\x02 \x01(\t\"2\n\x10NotebookCellDone\x12\x12\n\nexec_count\x18\x01 \x01(\x05\x12\n\n\x02ok\x18\x02 \x01(\x08\"\x1e\n\x1cInterruptNotebookCellRequest\":\n\x1dInterruptNotebookCellResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x65rror\x18\x02 \x01(\t\"\xbd\x01\n\x11NotebookCellChunk\x12\x0f\n\x07\x63\x65ll_id\x18\x01 \x01(\t\x12\x10\n\x06stdout\x18\x02 \x01(\tH\x00\x12\x10\n\x06stderr\x18\x03 \x01(\tH\x00\x12\x15\n\x0bresult_text\x18\x04 \x01(\tH\x00\x12\x13\n\timage_png\x18\x05 \x01(\x0cH\x00\x12\x19\n\x0f\x65rror_traceback\x18\x06 \x01(\tH\x00\x12!\n\x04\x64one\x18\x07 \x01(\x0b\x32\x11.NotebookCellDoneH\x00\x42\t\n\x07payload\"S\n\x10NotebookResponse\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0f\n\x07\x65xisted\x18\x02 \x01(\x08\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"7\n\x13SaveNotebookRequest\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\"M\n\x14SaveNotebookResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"C\n\x1bGenerateNotebookCodeRequest\x12\x0e\n\x06prompt\x18\x01 \x01(\t\x12\x14\n\x0c\x63ontext_code\x18\x02 \x01(\t\"\\\n\x1cGenerateNotebookCodeResponse\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x13\n\x0b\x65xplanation\x18\x02 \x01(\t\x12\n\n\x02ok\x18\x03 \x01(\x08\x12\r\n\x05\x65rror\x18\x04 \x01(\t*d\n\x13WeightOperationType\x12\n\n\x06ZEROFY\x10\x00\x12\x10\n\x0cREINITIALIZE\x10\x01\x12\n\n\x06\x46REEZE\x10\x02\x12\x12\n\x0eREMOVE_NEURONS\x10\t\x12\x0f\n\x0b\x41\x44\x44_NEURONS\x10\n*o\n\x0fZerofyPredicate\x12\x19\n\x15ZEROFY_PREDICATE_NONE\x10\x00\x12 \n\x1cZEROFY_PREDICATE_WITH_FROZEN\x10\x01\x12\x1f\n\x1bZEROFY_PREDICATE_WITH_OLDER\x10\x02*M\n\x0f\x41gentIntentType\x12\x12\n\x0eINTENT_UNKNOWN\x10\x00\x12\x11\n\rINTENT_FILTER\x10\x01\x12\x13\n\x0fINTENT_ANALYSIS\x10\x02*I\n\x0eSampleEditType\x12\x11\n\rEDIT_OVERRIDE\x10\x00\x12\x13\n\x0f\x45\x44IT_ACCUMULATE\x10\x01\x12\x0f\n\x0b\x45\x44IT_REMOVE\x10\x02*,\n\x11\x41gentProviderType\x12\x17\n\x13PROVIDER_OPENROUTER\x10\x00\x32\x9d\x0e\n\x11\x45xperimentService\x12P\n\x13GetLatestLoggerData\x12\x1b.GetLatestLoggerDataRequest\x1a\x1c.GetLatestLoggerDataResponse\x12\x36\n\x11\x45xperimentCommand\x12\x0f.TrainerCommand\x1a\x10.CommandResponse\x12H\n\x11ManipulateWeights\x12\x18.WeightsOperationRequest\x1a\x19.WeightsOperationResponse\x12/\n\nGetWeights\x12\x0f.WeightsRequest\x1a\x10.WeightsResponse\x12\x39\n\x0eGetActivations\x12\x12.ActivationRequest\x1a\x13.ActivationResponse\x12\x37\n\nGetSamples\x12\x13.BatchSampleRequest\x1a\x14.BatchSampleResponse\x12\x37\n\x0e\x41pplyDataQuery\x12\x11.DataQueryRequest\x1a\x12.DataQueryResponse\x12;\n\x0eGetDataSamples\x12\x13.DataSamplesRequest\x1a\x14.DataSamplesResponse\x12\x35\n\x0cGetHistogram\x12\x11.HistogramRequest\x1a\x12.HistogramResponse\x12\x38\n\x0bGetMetaData\x12\x13.GetMetaDataRequest\x1a\x14.GetMetaDataResponse\x12P\n\x13GetSignalTrajectory\x12\x1b.GetSignalTrajectoryRequest\x1a\x1c.GetSignalTrajectoryResponse\x12\x37\n\rGetPointCloud\x12\x12.PointCloudRequest\x1a\x10.PointCloudChunk0\x01\x12\x37\n\x0e\x45\x64itDataSample\x12\x11.DataEditsRequest\x1a\x12.DataEditsResponse\x12,\n\rGetDataSplits\x12\x06.Empty\x1a\x13.DataSplitsResponse\x12\x30\n\x10\x43heckAgentHealth\x12\x06.Empty\x1a\x14.AgentHealthResponse\x12\x44\n\x0fInitializeAgent\x12\x17.InitializeAgentRequest\x1a\x18.InitializeAgentResponse\x12G\n\x10\x43hangeAgentModel\x12\x18.ChangeAgentModelRequest\x1a\x19.ChangeAgentModelResponse\x12\x41\n\x0eGetAgentModels\x12\x16.GetAgentModelsRequest\x1a\x17.GetAgentModelsResponse\x12)\n\nResetAgent\x12\x06.Empty\x1a\x13.ResetAgentResponse\x12@\n\x0fRunNotebookCell\x12\x17.RunNotebookCellRequest\x1a\x12.NotebookCellChunk0\x01\x12V\n\x15InterruptNotebookCell\x12\x1d.InterruptNotebookCellRequest\x1a\x1e.InterruptNotebookCellResponse\x12(\n\x0bGetNotebook\x12\x06.Empty\x1a\x11.NotebookResponse\x12;\n\x0cSaveNotebook\x12\x14.SaveNotebookRequest\x1a\x15.SaveNotebookResponse\x12S\n\x14GenerateNotebookCode\x12\x1c.GenerateNotebookCodeRequest\x1a\x1d.GenerateNotebookCodeResponse\x12J\n\x11RestoreCheckpoint\x12\x19.RestoreCheckpointRequest\x1a\x1a.RestoreCheckpointResponse\x12J\n\x11TriggerEvaluation\x12\x19.TriggerEvaluationRequest\x1a\x1a.TriggerEvaluationResponse\x12P\n\x13GetEvaluationStatus\x12\x1b.GetEvaluationStatusRequest\x1a\x1c.GetEvaluationStatusResponse\x12G\n\x10\x43\x61ncelEvaluation\x12\x18.CancelEvaluationRequest\x1a\x19.CancelEvaluationResponseb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x18\x65xperiment_service.proto\"\x89\x01\n\x1aGetLatestLoggerDataRequest\x12\x1c\n\x14request_full_history\x18\x01 \x01(\x08\x12\x12\n\nmax_points\x18\x02 \x01(\x05\x12\x17\n\x0f\x62reak_by_slices\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\"\x81\x02\n\x0fLoggerDataPoint\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x11\n\tmodel_age\x18\x02 \x01(\x05\x12\x14\n\x0cmetric_value\x18\x03 \x01(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x04 \x01(\t\x12\x11\n\ttimestamp\x18\x05 \x01(\x03\x12\x11\n\tsample_id\x18\x06 \x01(\t\x12\x1c\n\x14is_evaluation_marker\x18\x07 \x01(\x08\x12\x12\n\nsplit_name\x18\x08 \x01(\t\x12\x17\n\x0f\x65valuation_tags\x18\t \x03(\t\x12\x12\n\npoint_note\x18\n \x01(\t\x12\x12\n\naudit_mode\x18\x0b \x01(\x08\"[\n\x1bGetLatestLoggerDataResponse\x12 \n\x06points\x18\x01 \x03(\x0b\x32\x10.LoggerDataPoint\x12\x1a\n\x12weightslab_version\x18\x02 \x01(\t\"\x07\n\x05\x45mpty\"/\n\x08NeuronId\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tneuron_id\x18\x02 \x01(\x05\"\x91\x02\n\x0fWeightOperation\x12*\n\x07op_type\x18\x01 \x01(\x0e\x32\x14.WeightOperationTypeH\x00\x88\x01\x01\x12\x15\n\x08layer_id\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1d\n\nneuron_ids\x18\x03 \x03(\x0b\x32\t.NeuronId\x12\x16\n\x0eneurons_to_add\x18\t \x01(\x05\x12 \n\x18zerofy_from_incoming_ids\x18\x0b \x03(\x05\x12\x1c\n\x14zerofy_to_neuron_ids\x18\x0c \x03(\x05\x12+\n\x11zerofy_predicates\x18\r \x03(\x0e\x32\x10.ZerofyPredicateB\n\n\x08_op_typeB\x0b\n\t_layer_id\"_\n\x17WeightsOperationRequest\x12/\n\x10weight_operation\x18\x01 \x01(\x0b\x32\x10.WeightOperationH\x00\x88\x01\x01\x42\x13\n\x11_weight_operation\"<\n\x18WeightsOperationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xc1\x05\n\x0fHyperParameters\x12\x1c\n\x0f\x65xperiment_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12!\n\x14training_steps_to_do\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x17\n\nbatch_size\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12 \n\x13\x66ull_eval_frequency\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12 \n\x13\x63heckpont_frequency\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x18\n\x0bis_training\x18\x07 \x01(\x08H\x06\x88\x01\x01\x12\x15\n\x08nb_steps\x18\x08 \x01(\x05H\x07\x88\x01\x01\x12\x19\n\x0c\x61uditor_mode\x18\t \x01(\x08H\x08\x88\x01\x01\x12\x1d\n\x10train_batch_size\x18\n \x01(\x05H\t\x88\x01\x01\x12\x1b\n\x0eval_batch_size\x18\x0b \x01(\x05H\n\x88\x01\x01\x12\x1c\n\x0ftest_batch_size\x18\x0c \x01(\x05H\x0b\x88\x01\x01\x12\x1c\n\x0f\x65valuation_mode\x18\r \x01(\x08H\x0c\x88\x01\x01\x12\x1e\n\x11\x65valuation_config\x18\x0e \x01(\tH\r\x88\x01\x01\x42\x12\n\x10_experiment_nameB\x17\n\x15_training_steps_to_doB\x10\n\x0e_learning_rateB\r\n\x0b_batch_sizeB\x16\n\x14_full_eval_frequencyB\x16\n\x14_checkpont_frequencyB\x0e\n\x0c_is_trainingB\x0b\n\t_nb_stepsB\x0f\n\r_auditor_modeB\x13\n\x11_train_batch_sizeB\x11\n\x0f_val_batch_sizeB\x12\n\x10_test_batch_sizeB\x12\n\x10_evaluation_modeB\x14\n\x12_evaluation_config\",\n\rMetricsStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"~\n\rAnnotatStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12.\n\x08metadata\x18\x02 \x03(\x0b\x32\x1c.AnnotatStatus.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x90\x02\n\x10TrainingStatusEx\x12\x16\n\ttimestamp\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0f\x65xperiment_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x16\n\tmodel_age\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12+\n\x0emetrics_status\x18\x04 \x01(\x0b\x32\x0e.MetricsStatusH\x03\x88\x01\x01\x12+\n\x0e\x61nnotat_status\x18\x05 \x01(\x0b\x32\x0e.AnnotatStatusH\x04\x88\x01\x01\x42\x0c\n\n_timestampB\x12\n\x10_experiment_nameB\x0c\n\n_model_ageB\x11\n\x0f_metrics_statusB\x11\n\x0f_annotat_status\"]\n\x15HyperParameterCommand\x12/\n\x10hyper_parameters\x18\x01 \x01(\x0b\x32\x10.HyperParametersH\x00\x88\x01\x01\x42\x13\n\x11_hyper_parameters\">\n\x14\x44\x65nySamplesOperation\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\"0\n\x17LoadCheckpointOperation\x12\x15\n\rcheckpoint_id\x18\x01 \x01(\x05\"b\n\x11PlotNoteOperation\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x0c\n\x04note\x18\x04 \x01(\t\"L\n\x17SaveCheckpointOperation\x12\x19\n\x11save_architecture\x18\x01 \x01(\x08\x12\x16\n\x0esave_optimizer\x18\x02 \x01(\x08\"\x1a\n\x18RestartInstanceOperation\"\x8d\x08\n\x0eTrainerCommand\x12\x1c\n\x14get_hyper_parameters\x18\x04 \x01(\x08\x12\x1e\n\x16get_interactive_layers\x18\x05 \x01(\x08\x12\x1d\n\x10get_data_records\x18\x06 \x01(\tH\x00\x88\x01\x01\x12%\n\x18get_single_layer_info_id\x18\x08 \x01(\x05H\x01\x88\x01\x01\x12;\n\x16hyper_parameter_change\x18\x01 \x01(\x0b\x32\x16.HyperParameterCommandH\x02\x88\x01\x01\x12:\n\x16\x64\x65ny_samples_operation\x18\x07 \x01(\x0b\x32\x15.DenySamplesOperationH\x03\x88\x01\x01\x12?\n\x1b\x64\x65ny_eval_samples_operation\x18\n \x01(\x0b\x32\x15.DenySamplesOperationH\x04\x88\x01\x01\x12@\n\x19load_checkpoint_operation\x18\t \x01(\x0b\x32\x18.LoadCheckpointOperationH\x05\x88\x01\x01\x12\x42\n\x1eremove_from_denylist_operation\x18\x0b \x01(\x0b\x32\x15.DenySamplesOperationH\x06\x88\x01\x01\x12G\n#remove_eval_from_denylist_operation\x18\x0c \x01(\x0b\x32\x15.DenySamplesOperationH\x07\x88\x01\x01\x12\x34\n\x13plot_note_operation\x18\r \x01(\x0b\x32\x12.PlotNoteOperationH\x08\x88\x01\x01\x12@\n\x19save_checkpoint_operation\x18\x0e \x01(\x0b\x32\x18.SaveCheckpointOperationH\t\x88\x01\x01\x12\x39\n\x11restart_operation\x18\x0f \x01(\x0b\x32\x19.RestartInstanceOperationH\n\x88\x01\x01\x42\x13\n\x11_get_data_recordsB\x1b\n\x19_get_single_layer_info_idB\x19\n\x17_hyper_parameter_changeB\x19\n\x17_deny_samples_operationB\x1e\n\x1c_deny_eval_samples_operationB\x1c\n\x1a_load_checkpoint_operationB!\n\x1f_remove_from_denylist_operationB&\n$_remove_eval_from_denylist_operationB\x16\n\x14_plot_note_operationB\x1c\n\x1a_save_checkpoint_operationB\x14\n\x12_restart_operation\"\x9d\x01\n\x12HyperParameterDesc\x12\r\n\x05label\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04type\x18\x03 \x01(\t\x12\x1c\n\x0fnumerical_value\x18\x04 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cstring_value\x18\x05 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_numerical_valueB\x0f\n\r_string_value\"\xf2\x02\n\x10NeuronStatistics\x12!\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronIdH\x00\x88\x01\x01\x12\x17\n\nneuron_age\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1f\n\x12train_trigger_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x1e\n\x11\x65val_trigger_rate\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x07 \x01(\x02H\x04\x88\x01\x01\x12\x36\n\x0bincoming_lr\x18\x08 \x03(\x0b\x32!.NeuronStatistics.IncomingLrEntry\x1a\x31\n\x0fIncomingLrEntry\x12\x0b\n\x03key\x18\x01 \x01(\x05\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\x42\x0c\n\n_neuron_idB\r\n\x0b_neuron_ageB\x15\n\x13_train_trigger_rateB\x14\n\x12_eval_trigger_rateB\x10\n\x0e_learning_rate\"\xf0\x02\n\x13LayerRepresentation\x12\x15\n\x08layer_id\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rneurons_count\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12#\n\x16incoming_neurons_count\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x13\n\x06stride\x18\x07 \x01(\x05H\x06\x88\x01\x01\x12-\n\x12neurons_statistics\x18\n \x03(\x0b\x32\x11.NeuronStatisticsB\x0b\n\t_layer_idB\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x10\n\x0e_neurons_countB\x19\n\x17_incoming_neurons_countB\x0e\n\x0c_kernel_sizeB\t\n\x07_stride\"H\n\x11\x41\x63tivationRequest\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tsample_id\x18\x02 \x01(\t\x12\x0e\n\x06origin\x18\x03 \x01(\t\"H\n\rActivationMap\x12\x11\n\tneuron_id\x18\x01 \x01(\x05\x12\x0e\n\x06values\x18\x02 \x03(\x02\x12\t\n\x01H\x18\x03 \x01(\x05\x12\t\n\x01W\x18\x04 \x01(\x05\"d\n\x12\x41\x63tivationResponse\x12\x12\n\nlayer_type\x18\x01 \x01(\t\x12\x15\n\rneurons_count\x18\x02 \x01(\x05\x12#\n\x0b\x61\x63tivations\x18\x03 \x03(\x0b\x32\x0e.ActivationMap\"\x93\x01\n\tTaskField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x15\n\x0b\x66loat_value\x18\x02 \x01(\x02H\x00\x12\x13\n\tint_value\x18\x03 \x01(\x05H\x00\x12\x16\n\x0cstring_value\x18\x04 \x01(\tH\x00\x12\x15\n\x0b\x62ytes_value\x18\x05 \x01(\x0cH\x00\x12\x14\n\nbool_value\x18\x06 \x01(\x08H\x00\x42\x07\n\x05value\"\x87\x03\n\x0eRecordMetadata\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x14\n\x0csample_label\x18\x02 \x03(\x05\x12\x19\n\x11sample_prediction\x18\x03 \x03(\x05\x12=\n\x10sample_last_loss\x18\x04 \x03(\x0b\x32#.RecordMetadata.SampleLastLossEntry\x12\x19\n\x11sample_encounters\x18\x05 \x01(\x05\x12\x18\n\x10sample_discarded\x18\x06 \x01(\x08\x12 \n\x0c\x65xtra_fields\x18\x07 \x03(\x0b\x32\n.TaskField\x12\x16\n\x0eprediction_raw\x18\t \x01(\x0c\x12\x11\n\ttask_type\x18\n \x01(\t\x12\x19\n\x11sample_label_text\x18\x0b \x03(\t\x12\x1e\n\x16sample_prediction_text\x18\x0c \x03(\t\x1a\x35\n\x13SampleLastLossEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x93\x01\n\x10SampleStatistics\x12\x13\n\x06origin\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0csample_count\x18\x07 \x01(\x05H\x01\x88\x01\x01\x12\x11\n\ttask_type\x18\t \x01(\t\x12 \n\x07records\x18\x08 \x03(\x0b\x32\x0f.RecordMetadataB\t\n\x07_originB\x0f\n\r_sample_count\"\xe6\x01\n\x0f\x43ommandResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x33\n\x16hyper_parameters_descs\x18\x03 \x03(\x0b\x32\x13.HyperParameterDesc\x12\x33\n\x15layer_representations\x18\x04 \x03(\x0b\x32\x14.LayerRepresentation\x12\x31\n\x11sample_statistics\x18\x05 \x01(\x0b\x32\x11.SampleStatisticsH\x00\x88\x01\x01\x42\x14\n\x12_sample_statistics\"U\n\rSampleRequest\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_origin\"\xad\x02\n\x15SampleRequestResponse\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05label\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x64\x61ta\x18\x04 \x01(\x0cH\x03\x88\x01\x01\x12\x1a\n\rerror_message\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x15\n\x08raw_data\x18\x06 \x01(\x0cH\x05\x88\x01\x01\x12\x11\n\x04mask\x18\x07 \x01(\x0cH\x06\x88\x01\x01\x12\x17\n\nprediction\x18\x08 \x01(\x0cH\x07\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_originB\x08\n\x06_labelB\x07\n\x05_dataB\x10\n\x0e_error_messageB\x0b\n\t_raw_dataB\x07\n\x05_maskB\r\n\x0b_prediction\"\x92\x01\n\x12\x42\x61tchSampleRequest\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x19\n\x0cresize_width\x18\x03 \x01(\x05H\x00\x88\x01\x01\x12\x1a\n\rresize_height\x18\x04 \x01(\x05H\x01\x88\x01\x01\x42\x0f\n\r_resize_widthB\x10\n\x0e_resize_height\">\n\x13\x42\x61tchSampleResponse\x12\'\n\x07samples\x18\x01 \x03(\x0b\x32\x16.SampleRequestResponse\".\n\x0eWeightsRequest\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\"\x9d\x02\n\x0fWeightsResponse\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x10\n\x08incoming\x18\x04 \x01(\x05\x12\x10\n\x08outgoing\x18\x05 \x01(\x05\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x0f\n\x07weights\x18\x07 \x03(\x02\x12\x0f\n\x07success\x18\x0b \x01(\x08\x12\x1a\n\rerror_message\x18\x0c \x01(\tH\x03\x88\x01\x01\x42\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x0e\n\x0c_kernel_sizeB\x10\n\x0e_error_message\"R\n\x10\x44\x61taQueryRequest\x12\r\n\x05query\x18\x01 \x01(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\x12\x1b\n\x13is_natural_language\x18\x03 \x01(\x08\"5\n\x11\x43\x61tegoricalTagDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\ncategories\x18\x02 \x03(\t\"\xa9\x02\n\x11\x44\x61taQueryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1d\n\x15number_of_all_samples\x18\x03 \x01(\x05\x12%\n\x1dnumber_of_samples_in_the_loop\x18\x04 \x01(\x05\x12#\n\x1bnumber_of_discarded_samples\x18\x05 \x01(\x05\x12\x13\n\x0bunique_tags\x18\x06 \x03(\t\x12+\n\x11\x61gent_intent_type\x18\x07 \x01(\x0e\x32\x10.AgentIntentType\x12\x17\n\x0f\x61nalysis_result\x18\x08 \x01(\t\x12,\n\x10\x63\x61tegorical_tags\x18\t \x03(\x0b\x32\x12.CategoricalTagDef\"\xc2\x01\n\x12\x44\x61taSamplesRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12 \n\x18include_transformed_data\x18\x03 \x01(\x08\x12\x18\n\x10include_raw_data\x18\x04 \x01(\x08\x12\x19\n\x11stats_to_retrieve\x18\x05 \x03(\t\x12\x14\n\x0cresize_width\x18\x06 \x01(\x05\x12\x15\n\rresize_height\x18\x07 \x01(\x05\"m\n\x08\x44\x61taStat\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\r\n\x05shape\x18\x03 \x03(\x05\x12\r\n\x05value\x18\x04 \x03(\x02\x12\x14\n\x0cvalue_string\x18\x05 \x01(\t\x12\x11\n\tthumbnail\x18\x06 \x01(\x0c\">\n\nDataRecord\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x1d\n\ndata_stats\x18\x02 \x03(\x0b\x32\t.DataStat\"Z\n\x13\x44\x61taSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12!\n\x0c\x64\x61ta_records\x18\x03 \x03(\x0b\x32\x0b.DataRecord\"C\n\x0fHistogramSubBar\x12\x0e\n\x06origin\x18\x01 \x01(\t\x12\x11\n\tdiscarded\x18\x02 \x01(\x08\x12\r\n\x05\x63ount\x18\x03 \x01(\x03\"h\n\x0cHistogramBin\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03\x61vg\x18\x03 \x01(\x01\x12\r\n\x05\x63ount\x18\x04 \x01(\x03\x12\"\n\x08sub_bars\x18\x05 \x03(\x0b\x32\x10.HistogramSubBar\"[\n\x17\x43\x61tegoricalHistogramBar\x12\r\n\x05label\x18\x01 \x01(\t\x12\r\n\x05\x63ount\x18\x02 \x01(\x03\x12\"\n\x08sub_bars\x18\x03 \x03(\x0b\x32\x10.HistogramSubBar\"4\n\x10HistogramRequest\x12\x0e\n\x06\x63olumn\x18\x01 \x01(\t\x12\x10\n\x08max_bins\x18\x02 \x01(\x05\"\xb2\x01\n\x11HistogramResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\ntotal_rows\x18\x03 \x01(\x03\x12\x1b\n\x04\x62ins\x18\x04 \x03(\x0b\x32\r.HistogramBin\x12\x16\n\x0eis_categorical\x18\x05 \x01(\x08\x12\x32\n\x10\x63\x61tegorical_bars\x18\x06 \x03(\x0b\x32\x18.CategoricalHistogramBar\"W\n\x12GetMetaDataRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12\x17\n\x0fmodal_sample_id\x18\x03 \x01(\t\"\x99\x01\n\x13GetMetaDataResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1a\n\x12\x61ll_metadata_names\x18\x03 \x03(\t\x12!\n\x0cgrid_records\x18\x04 \x03(\x0b\x32\x0b.DataRecord\x12!\n\x0cmodal_record\x18\x05 \x01(\x0b\x32\x0b.DataRecord\"Y\n\x1aGetSignalTrajectoryRequest\x12\x13\n\x0bsignal_name\x18\x01 \x01(\t\x12\x12\n\nsample_ids\x18\x02 \x03(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"4\n\x10SignalTrajectory\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x03(\x02\"}\n\x1bGetSignalTrajectoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x13\n\x0bsignal_name\x18\x03 \x01(\t\x12\'\n\x0ctrajectories\x18\x04 \x03(\x0b\x32\x11.SignalTrajectory\"J\n\x11PointCloudRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"\xbf\x01\n\x0fPointCloudChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nnum_points\x18\x03 \x01(\x05\x12\x14\n\x0cnum_features\x18\x04 \x01(\x05\x12\x10\n\x08pc_range\x18\x05 \x03(\x02\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\x07 \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x08 \x01(\x05\x12\x15\n\rfeature_names\x18\t \x03(\t\"\xdc\x01\n\x10\x44\x61taEditsRequest\x12\x11\n\tstat_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66loat_value\x18\x02 \x01(\x02\x12\x14\n\x0cstring_value\x18\x03 \x01(\t\x12\x12\n\nbool_value\x18\x04 \x01(\x08\x12\x1d\n\x04type\x18\x05 \x01(\x0e\x32\x0f.SampleEditType\x12\x13\n\x0bsamples_ids\x18\x06 \x03(\t\x12\x16\n\x0esample_origins\x18\x07 \x03(\t\x12\x16\n\x0eis_categorical\x18\x08 \x01(\x08\x12\x12\n\ncategories\x18\t \x03(\t\"5\n\x11\x44\x61taEditsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\":\n\x12\x44\x61taSplitsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x13\n\x0bsplit_names\x18\x02 \x03(\t\"9\n\x13\x41gentHealthResponse\x12\x11\n\tavailable\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"^\n\x16InitializeAgentRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12$\n\x08provider\x18\x02 \x01(\x0e\x32\x12.AgentProviderType\x12\r\n\x05model\x18\x03 \x01(\t\";\n\x17InitializeAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"(\n\x17\x43hangeAgentModelRequest\x12\r\n\x05model\x18\x01 \x01(\t\"<\n\x18\x43hangeAgentModelResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x17\n\x15GetAgentModelsRequest\"J\n\x16GetAgentModelsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0e\n\x06models\x18\x02 \x03(\t\x12\x0f\n\x07message\x18\x03 \x01(\t\"6\n\x12ResetAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"=\n\x19\x43learAgentHistoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"?\n\x1b\x43ompactAgentHistoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"3\n\x18RestoreCheckpointRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\"=\n\x19RestoreCheckpointResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"R\n\x18TriggerEvaluationRequest\x12\x12\n\nsplit_name\x18\x01 \x01(\t\x12\x0c\n\x04tags\x18\x02 \x03(\t\x12\x14\n\x0cuse_full_set\x18\x03 \x01(\x08\"=\n\x19TriggerEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x1c\n\x1aGetEvaluationStatusRequest\"\x81\x01\n\x1bGetEvaluationStatusResponse\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x0f\n\x07\x63urrent\x18\x02 \x01(\x05\x12\r\n\x05total\x18\x03 \x01(\x05\x12\x0f\n\x07message\x18\x04 \x01(\t\x12\r\n\x05\x65rror\x18\x05 \x01(\t\x12\x12\n\nsplit_name\x18\x06 \x01(\t\")\n\x17\x43\x61ncelEvaluationRequest\x12\x0e\n\x06reason\x18\x01 \x01(\t\"<\n\x18\x43\x61ncelEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"7\n\x16RunNotebookCellRequest\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07\x63\x65ll_id\x18\x02 \x01(\t\"2\n\x10NotebookCellDone\x12\x12\n\nexec_count\x18\x01 \x01(\x05\x12\n\n\x02ok\x18\x02 \x01(\x08\"\x1e\n\x1cInterruptNotebookCellRequest\":\n\x1dInterruptNotebookCellResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x65rror\x18\x02 \x01(\t\"\xbd\x01\n\x11NotebookCellChunk\x12\x0f\n\x07\x63\x65ll_id\x18\x01 \x01(\t\x12\x10\n\x06stdout\x18\x02 \x01(\tH\x00\x12\x10\n\x06stderr\x18\x03 \x01(\tH\x00\x12\x15\n\x0bresult_text\x18\x04 \x01(\tH\x00\x12\x13\n\timage_png\x18\x05 \x01(\x0cH\x00\x12\x19\n\x0f\x65rror_traceback\x18\x06 \x01(\tH\x00\x12!\n\x04\x64one\x18\x07 \x01(\x0b\x32\x11.NotebookCellDoneH\x00\x42\t\n\x07payload\"S\n\x10NotebookResponse\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0f\n\x07\x65xisted\x18\x02 \x01(\x08\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"7\n\x13SaveNotebookRequest\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\"M\n\x14SaveNotebookResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"C\n\x1bGenerateNotebookCodeRequest\x12\x0e\n\x06prompt\x18\x01 \x01(\t\x12\x14\n\x0c\x63ontext_code\x18\x02 \x01(\t\"\\\n\x1cGenerateNotebookCodeResponse\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x13\n\x0b\x65xplanation\x18\x02 \x01(\t\x12\n\n\x02ok\x18\x03 \x01(\x08\x12\r\n\x05\x65rror\x18\x04 \x01(\t*d\n\x13WeightOperationType\x12\n\n\x06ZEROFY\x10\x00\x12\x10\n\x0cREINITIALIZE\x10\x01\x12\n\n\x06\x46REEZE\x10\x02\x12\x12\n\x0eREMOVE_NEURONS\x10\t\x12\x0f\n\x0b\x41\x44\x44_NEURONS\x10\n*o\n\x0fZerofyPredicate\x12\x19\n\x15ZEROFY_PREDICATE_NONE\x10\x00\x12 \n\x1cZEROFY_PREDICATE_WITH_FROZEN\x10\x01\x12\x1f\n\x1bZEROFY_PREDICATE_WITH_OLDER\x10\x02*M\n\x0f\x41gentIntentType\x12\x12\n\x0eINTENT_UNKNOWN\x10\x00\x12\x11\n\rINTENT_FILTER\x10\x01\x12\x13\n\x0fINTENT_ANALYSIS\x10\x02*I\n\x0eSampleEditType\x12\x11\n\rEDIT_OVERRIDE\x10\x00\x12\x13\n\x0f\x45\x44IT_ACCUMULATE\x10\x01\x12\x0f\n\x0b\x45\x44IT_REMOVE\x10\x02*C\n\x11\x41gentProviderType\x12\x17\n\x13PROVIDER_OPENROUTER\x10\x00\x12\x15\n\x11PROVIDER_OPENCODE\x10\x01\x32\x93\x0f\n\x11\x45xperimentService\x12P\n\x13GetLatestLoggerData\x12\x1b.GetLatestLoggerDataRequest\x1a\x1c.GetLatestLoggerDataResponse\x12\x36\n\x11\x45xperimentCommand\x12\x0f.TrainerCommand\x1a\x10.CommandResponse\x12H\n\x11ManipulateWeights\x12\x18.WeightsOperationRequest\x1a\x19.WeightsOperationResponse\x12/\n\nGetWeights\x12\x0f.WeightsRequest\x1a\x10.WeightsResponse\x12\x39\n\x0eGetActivations\x12\x12.ActivationRequest\x1a\x13.ActivationResponse\x12\x37\n\nGetSamples\x12\x13.BatchSampleRequest\x1a\x14.BatchSampleResponse\x12\x37\n\x0e\x41pplyDataQuery\x12\x11.DataQueryRequest\x1a\x12.DataQueryResponse\x12;\n\x0eGetDataSamples\x12\x13.DataSamplesRequest\x1a\x14.DataSamplesResponse\x12\x35\n\x0cGetHistogram\x12\x11.HistogramRequest\x1a\x12.HistogramResponse\x12\x38\n\x0bGetMetaData\x12\x13.GetMetaDataRequest\x1a\x14.GetMetaDataResponse\x12P\n\x13GetSignalTrajectory\x12\x1b.GetSignalTrajectoryRequest\x1a\x1c.GetSignalTrajectoryResponse\x12\x37\n\rGetPointCloud\x12\x12.PointCloudRequest\x1a\x10.PointCloudChunk0\x01\x12\x37\n\x0e\x45\x64itDataSample\x12\x11.DataEditsRequest\x1a\x12.DataEditsResponse\x12,\n\rGetDataSplits\x12\x06.Empty\x1a\x13.DataSplitsResponse\x12\x30\n\x10\x43heckAgentHealth\x12\x06.Empty\x1a\x14.AgentHealthResponse\x12\x44\n\x0fInitializeAgent\x12\x17.InitializeAgentRequest\x1a\x18.InitializeAgentResponse\x12G\n\x10\x43hangeAgentModel\x12\x18.ChangeAgentModelRequest\x1a\x19.ChangeAgentModelResponse\x12\x41\n\x0eGetAgentModels\x12\x16.GetAgentModelsRequest\x1a\x17.GetAgentModelsResponse\x12)\n\nResetAgent\x12\x06.Empty\x1a\x13.ResetAgentResponse\x12\x37\n\x11\x43learAgentHistory\x12\x06.Empty\x1a\x1a.ClearAgentHistoryResponse\x12;\n\x13\x43ompactAgentHistory\x12\x06.Empty\x1a\x1c.CompactAgentHistoryResponse\x12@\n\x0fRunNotebookCell\x12\x17.RunNotebookCellRequest\x1a\x12.NotebookCellChunk0\x01\x12V\n\x15InterruptNotebookCell\x12\x1d.InterruptNotebookCellRequest\x1a\x1e.InterruptNotebookCellResponse\x12(\n\x0bGetNotebook\x12\x06.Empty\x1a\x11.NotebookResponse\x12;\n\x0cSaveNotebook\x12\x14.SaveNotebookRequest\x1a\x15.SaveNotebookResponse\x12S\n\x14GenerateNotebookCode\x12\x1c.GenerateNotebookCodeRequest\x1a\x1d.GenerateNotebookCodeResponse\x12J\n\x11RestoreCheckpoint\x12\x19.RestoreCheckpointRequest\x1a\x1a.RestoreCheckpointResponse\x12J\n\x11TriggerEvaluation\x12\x19.TriggerEvaluationRequest\x1a\x1a.TriggerEvaluationResponse\x12P\n\x13GetEvaluationStatus\x12\x1b.GetEvaluationStatusRequest\x1a\x1c.GetEvaluationStatusResponse\x12G\n\x10\x43\x61ncelEvaluation\x12\x18.CancelEvaluationRequest\x1a\x19.CancelEvaluationResponseb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'weightslab.proto.experiment_service_pb2', _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'experiment_service_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: DESCRIPTOR._loaded_options = None _globals['_ANNOTATSTATUS_METADATAENTRY']._loaded_options = None @@ -37,188 +37,192 @@ _globals['_NEURONSTATISTICS_INCOMINGLRENTRY']._serialized_options = b'8\001' _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._loaded_options = None _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._serialized_options = b'8\001' - _globals['_WEIGHTOPERATIONTYPE']._serialized_start=10979 - _globals['_WEIGHTOPERATIONTYPE']._serialized_end=11079 - _globals['_ZEROFYPREDICATE']._serialized_start=11081 - _globals['_ZEROFYPREDICATE']._serialized_end=11192 - _globals['_AGENTINTENTTYPE']._serialized_start=11194 - _globals['_AGENTINTENTTYPE']._serialized_end=11271 - _globals['_SAMPLEEDITTYPE']._serialized_start=11273 - _globals['_SAMPLEEDITTYPE']._serialized_end=11346 - _globals['_AGENTPROVIDERTYPE']._serialized_start=11348 - _globals['_AGENTPROVIDERTYPE']._serialized_end=11392 - _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_start=46 - _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_end=183 - _globals['_LOGGERDATAPOINT']._serialized_start=186 - _globals['_LOGGERDATAPOINT']._serialized_end=443 - _globals['_GETLATESTLOGGERDATARESPONSE']._serialized_start=445 - _globals['_GETLATESTLOGGERDATARESPONSE']._serialized_end=536 - _globals['_EMPTY']._serialized_start=538 - _globals['_EMPTY']._serialized_end=545 - _globals['_NEURONID']._serialized_start=547 - _globals['_NEURONID']._serialized_end=594 - _globals['_WEIGHTOPERATION']._serialized_start=597 - _globals['_WEIGHTOPERATION']._serialized_end=870 - _globals['_WEIGHTSOPERATIONREQUEST']._serialized_start=872 - _globals['_WEIGHTSOPERATIONREQUEST']._serialized_end=967 - _globals['_WEIGHTSOPERATIONRESPONSE']._serialized_start=969 - _globals['_WEIGHTSOPERATIONRESPONSE']._serialized_end=1029 - _globals['_HYPERPARAMETERS']._serialized_start=1032 - _globals['_HYPERPARAMETERS']._serialized_end=1737 - _globals['_METRICSSTATUS']._serialized_start=1739 - _globals['_METRICSSTATUS']._serialized_end=1783 - _globals['_ANNOTATSTATUS']._serialized_start=1785 - _globals['_ANNOTATSTATUS']._serialized_end=1911 - _globals['_ANNOTATSTATUS_METADATAENTRY']._serialized_start=1864 - _globals['_ANNOTATSTATUS_METADATAENTRY']._serialized_end=1911 - _globals['_TRAININGSTATUSEX']._serialized_start=1914 - _globals['_TRAININGSTATUSEX']._serialized_end=2186 - _globals['_HYPERPARAMETERCOMMAND']._serialized_start=2188 - _globals['_HYPERPARAMETERCOMMAND']._serialized_end=2281 - _globals['_DENYSAMPLESOPERATION']._serialized_start=2283 - _globals['_DENYSAMPLESOPERATION']._serialized_end=2345 - _globals['_LOADCHECKPOINTOPERATION']._serialized_start=2347 - _globals['_LOADCHECKPOINTOPERATION']._serialized_end=2395 - _globals['_PLOTNOTEOPERATION']._serialized_start=2397 - _globals['_PLOTNOTEOPERATION']._serialized_end=2495 - _globals['_SAVECHECKPOINTOPERATION']._serialized_start=2497 - _globals['_SAVECHECKPOINTOPERATION']._serialized_end=2573 - _globals['_RESTARTINSTANCEOPERATION']._serialized_start=2575 - _globals['_RESTARTINSTANCEOPERATION']._serialized_end=2601 - _globals['_TRAINERCOMMAND']._serialized_start=2604 - _globals['_TRAINERCOMMAND']._serialized_end=3641 - _globals['_HYPERPARAMETERDESC']._serialized_start=3644 - _globals['_HYPERPARAMETERDESC']._serialized_end=3801 - _globals['_NEURONSTATISTICS']._serialized_start=3804 - _globals['_NEURONSTATISTICS']._serialized_end=4174 - _globals['_NEURONSTATISTICS_INCOMINGLRENTRY']._serialized_start=4033 - _globals['_NEURONSTATISTICS_INCOMINGLRENTRY']._serialized_end=4082 - _globals['_LAYERREPRESENTATION']._serialized_start=4177 - _globals['_LAYERREPRESENTATION']._serialized_end=4545 - _globals['_ACTIVATIONREQUEST']._serialized_start=4547 - _globals['_ACTIVATIONREQUEST']._serialized_end=4619 - _globals['_ACTIVATIONMAP']._serialized_start=4621 - _globals['_ACTIVATIONMAP']._serialized_end=4693 - _globals['_ACTIVATIONRESPONSE']._serialized_start=4695 - _globals['_ACTIVATIONRESPONSE']._serialized_end=4795 - _globals['_TASKFIELD']._serialized_start=4798 - _globals['_TASKFIELD']._serialized_end=4945 - _globals['_RECORDMETADATA']._serialized_start=4948 - _globals['_RECORDMETADATA']._serialized_end=5339 - _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._serialized_start=5286 - _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._serialized_end=5339 - _globals['_SAMPLESTATISTICS']._serialized_start=5342 - _globals['_SAMPLESTATISTICS']._serialized_end=5489 - _globals['_COMMANDRESPONSE']._serialized_start=5492 - _globals['_COMMANDRESPONSE']._serialized_end=5722 - _globals['_SAMPLEREQUEST']._serialized_start=5724 - _globals['_SAMPLEREQUEST']._serialized_end=5809 - _globals['_SAMPLEREQUESTRESPONSE']._serialized_start=5812 - _globals['_SAMPLEREQUESTRESPONSE']._serialized_end=6113 - _globals['_BATCHSAMPLEREQUEST']._serialized_start=6116 - _globals['_BATCHSAMPLEREQUEST']._serialized_end=6262 - _globals['_BATCHSAMPLERESPONSE']._serialized_start=6264 - _globals['_BATCHSAMPLERESPONSE']._serialized_end=6326 - _globals['_WEIGHTSREQUEST']._serialized_start=6328 - _globals['_WEIGHTSREQUEST']._serialized_end=6374 - _globals['_WEIGHTSRESPONSE']._serialized_start=6377 - _globals['_WEIGHTSRESPONSE']._serialized_end=6662 - _globals['_DATAQUERYREQUEST']._serialized_start=6664 - _globals['_DATAQUERYREQUEST']._serialized_end=6746 - _globals['_CATEGORICALTAGDEF']._serialized_start=6748 - _globals['_CATEGORICALTAGDEF']._serialized_end=6801 - _globals['_DATAQUERYRESPONSE']._serialized_start=6804 - _globals['_DATAQUERYRESPONSE']._serialized_end=7101 - _globals['_DATASAMPLESREQUEST']._serialized_start=7104 - _globals['_DATASAMPLESREQUEST']._serialized_end=7298 - _globals['_DATASTAT']._serialized_start=7300 - _globals['_DATASTAT']._serialized_end=7409 - _globals['_DATARECORD']._serialized_start=7411 - _globals['_DATARECORD']._serialized_end=7473 - _globals['_DATASAMPLESRESPONSE']._serialized_start=7475 - _globals['_DATASAMPLESRESPONSE']._serialized_end=7565 - _globals['_HISTOGRAMSUBBAR']._serialized_start=7567 - _globals['_HISTOGRAMSUBBAR']._serialized_end=7634 - _globals['_HISTOGRAMBIN']._serialized_start=7636 - _globals['_HISTOGRAMBIN']._serialized_end=7740 - _globals['_CATEGORICALHISTOGRAMBAR']._serialized_start=7742 - _globals['_CATEGORICALHISTOGRAMBAR']._serialized_end=7833 - _globals['_HISTOGRAMREQUEST']._serialized_start=7835 - _globals['_HISTOGRAMREQUEST']._serialized_end=7887 - _globals['_HISTOGRAMRESPONSE']._serialized_start=7890 - _globals['_HISTOGRAMRESPONSE']._serialized_end=8068 - _globals['_GETMETADATAREQUEST']._serialized_start=8070 - _globals['_GETMETADATAREQUEST']._serialized_end=8157 - _globals['_GETMETADATARESPONSE']._serialized_start=8160 - _globals['_GETMETADATARESPONSE']._serialized_end=8313 - _globals['_GETSIGNALTRAJECTORYREQUEST']._serialized_start=8315 - _globals['_GETSIGNALTRAJECTORYREQUEST']._serialized_end=8404 - _globals['_SIGNALTRAJECTORY']._serialized_start=8406 - _globals['_SIGNALTRAJECTORY']._serialized_end=8458 - _globals['_GETSIGNALTRAJECTORYRESPONSE']._serialized_start=8460 - _globals['_GETSIGNALTRAJECTORYRESPONSE']._serialized_end=8585 - _globals['_POINTCLOUDREQUEST']._serialized_start=8587 - _globals['_POINTCLOUDREQUEST']._serialized_end=8661 - _globals['_POINTCLOUDCHUNK']._serialized_start=8664 - _globals['_POINTCLOUDCHUNK']._serialized_end=8855 - _globals['_DATAEDITSREQUEST']._serialized_start=8858 - _globals['_DATAEDITSREQUEST']._serialized_end=9078 - _globals['_DATAEDITSRESPONSE']._serialized_start=9080 - _globals['_DATAEDITSRESPONSE']._serialized_end=9133 - _globals['_DATASPLITSRESPONSE']._serialized_start=9135 - _globals['_DATASPLITSRESPONSE']._serialized_end=9193 - _globals['_AGENTHEALTHRESPONSE']._serialized_start=9195 - _globals['_AGENTHEALTHRESPONSE']._serialized_end=9252 - _globals['_INITIALIZEAGENTREQUEST']._serialized_start=9254 - _globals['_INITIALIZEAGENTREQUEST']._serialized_end=9348 - _globals['_INITIALIZEAGENTRESPONSE']._serialized_start=9350 - _globals['_INITIALIZEAGENTRESPONSE']._serialized_end=9409 - _globals['_CHANGEAGENTMODELREQUEST']._serialized_start=9411 - _globals['_CHANGEAGENTMODELREQUEST']._serialized_end=9451 - _globals['_CHANGEAGENTMODELRESPONSE']._serialized_start=9453 - _globals['_CHANGEAGENTMODELRESPONSE']._serialized_end=9513 - _globals['_GETAGENTMODELSREQUEST']._serialized_start=9515 - _globals['_GETAGENTMODELSREQUEST']._serialized_end=9538 - _globals['_GETAGENTMODELSRESPONSE']._serialized_start=9540 - _globals['_GETAGENTMODELSRESPONSE']._serialized_end=9614 - _globals['_RESETAGENTRESPONSE']._serialized_start=9616 - _globals['_RESETAGENTRESPONSE']._serialized_end=9670 - _globals['_RESTORECHECKPOINTREQUEST']._serialized_start=9672 - _globals['_RESTORECHECKPOINTREQUEST']._serialized_end=9723 - _globals['_RESTORECHECKPOINTRESPONSE']._serialized_start=9725 - _globals['_RESTORECHECKPOINTRESPONSE']._serialized_end=9786 - _globals['_TRIGGEREVALUATIONREQUEST']._serialized_start=9788 - _globals['_TRIGGEREVALUATIONREQUEST']._serialized_end=9870 - _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_start=9872 - _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_end=9933 - _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_start=9935 - _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_end=9963 - _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_start=9966 - _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_end=10095 - _globals['_CANCELEVALUATIONREQUEST']._serialized_start=10097 - _globals['_CANCELEVALUATIONREQUEST']._serialized_end=10138 - _globals['_CANCELEVALUATIONRESPONSE']._serialized_start=10140 - _globals['_CANCELEVALUATIONRESPONSE']._serialized_end=10200 - _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_start=10202 - _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_end=10257 - _globals['_NOTEBOOKCELLDONE']._serialized_start=10259 - _globals['_NOTEBOOKCELLDONE']._serialized_end=10309 - _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_start=10311 - _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_end=10341 - _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_start=10343 - _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_end=10401 - _globals['_NOTEBOOKCELLCHUNK']._serialized_start=10404 - _globals['_NOTEBOOKCELLCHUNK']._serialized_end=10593 - _globals['_NOTEBOOKRESPONSE']._serialized_start=10595 - _globals['_NOTEBOOKRESPONSE']._serialized_end=10678 - _globals['_SAVENOTEBOOKREQUEST']._serialized_start=10680 - _globals['_SAVENOTEBOOKREQUEST']._serialized_end=10735 - _globals['_SAVENOTEBOOKRESPONSE']._serialized_start=10737 - _globals['_SAVENOTEBOOKRESPONSE']._serialized_end=10814 - _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_start=10816 - _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_end=10883 - _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_start=10885 - _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_end=10977 - _globals['_EXPERIMENTSERVICE']._serialized_start=11395 - _globals['_EXPERIMENTSERVICE']._serialized_end=13216 + _globals['_WEIGHTOPERATIONTYPE']._serialized_start=11090 + _globals['_WEIGHTOPERATIONTYPE']._serialized_end=11190 + _globals['_ZEROFYPREDICATE']._serialized_start=11192 + _globals['_ZEROFYPREDICATE']._serialized_end=11303 + _globals['_AGENTINTENTTYPE']._serialized_start=11305 + _globals['_AGENTINTENTTYPE']._serialized_end=11382 + _globals['_SAMPLEEDITTYPE']._serialized_start=11384 + _globals['_SAMPLEEDITTYPE']._serialized_end=11457 + _globals['_AGENTPROVIDERTYPE']._serialized_start=11459 + _globals['_AGENTPROVIDERTYPE']._serialized_end=11526 + _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_start=29 + _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_end=166 + _globals['_LOGGERDATAPOINT']._serialized_start=169 + _globals['_LOGGERDATAPOINT']._serialized_end=426 + _globals['_GETLATESTLOGGERDATARESPONSE']._serialized_start=428 + _globals['_GETLATESTLOGGERDATARESPONSE']._serialized_end=519 + _globals['_EMPTY']._serialized_start=521 + _globals['_EMPTY']._serialized_end=528 + _globals['_NEURONID']._serialized_start=530 + _globals['_NEURONID']._serialized_end=577 + _globals['_WEIGHTOPERATION']._serialized_start=580 + _globals['_WEIGHTOPERATION']._serialized_end=853 + _globals['_WEIGHTSOPERATIONREQUEST']._serialized_start=855 + _globals['_WEIGHTSOPERATIONREQUEST']._serialized_end=950 + _globals['_WEIGHTSOPERATIONRESPONSE']._serialized_start=952 + _globals['_WEIGHTSOPERATIONRESPONSE']._serialized_end=1012 + _globals['_HYPERPARAMETERS']._serialized_start=1015 + _globals['_HYPERPARAMETERS']._serialized_end=1720 + _globals['_METRICSSTATUS']._serialized_start=1722 + _globals['_METRICSSTATUS']._serialized_end=1766 + _globals['_ANNOTATSTATUS']._serialized_start=1768 + _globals['_ANNOTATSTATUS']._serialized_end=1894 + _globals['_ANNOTATSTATUS_METADATAENTRY']._serialized_start=1847 + _globals['_ANNOTATSTATUS_METADATAENTRY']._serialized_end=1894 + _globals['_TRAININGSTATUSEX']._serialized_start=1897 + _globals['_TRAININGSTATUSEX']._serialized_end=2169 + _globals['_HYPERPARAMETERCOMMAND']._serialized_start=2171 + _globals['_HYPERPARAMETERCOMMAND']._serialized_end=2264 + _globals['_DENYSAMPLESOPERATION']._serialized_start=2266 + _globals['_DENYSAMPLESOPERATION']._serialized_end=2328 + _globals['_LOADCHECKPOINTOPERATION']._serialized_start=2330 + _globals['_LOADCHECKPOINTOPERATION']._serialized_end=2378 + _globals['_PLOTNOTEOPERATION']._serialized_start=2380 + _globals['_PLOTNOTEOPERATION']._serialized_end=2478 + _globals['_SAVECHECKPOINTOPERATION']._serialized_start=2480 + _globals['_SAVECHECKPOINTOPERATION']._serialized_end=2556 + _globals['_RESTARTINSTANCEOPERATION']._serialized_start=2558 + _globals['_RESTARTINSTANCEOPERATION']._serialized_end=2584 + _globals['_TRAINERCOMMAND']._serialized_start=2587 + _globals['_TRAINERCOMMAND']._serialized_end=3624 + _globals['_HYPERPARAMETERDESC']._serialized_start=3627 + _globals['_HYPERPARAMETERDESC']._serialized_end=3784 + _globals['_NEURONSTATISTICS']._serialized_start=3787 + _globals['_NEURONSTATISTICS']._serialized_end=4157 + _globals['_NEURONSTATISTICS_INCOMINGLRENTRY']._serialized_start=4016 + _globals['_NEURONSTATISTICS_INCOMINGLRENTRY']._serialized_end=4065 + _globals['_LAYERREPRESENTATION']._serialized_start=4160 + _globals['_LAYERREPRESENTATION']._serialized_end=4528 + _globals['_ACTIVATIONREQUEST']._serialized_start=4530 + _globals['_ACTIVATIONREQUEST']._serialized_end=4602 + _globals['_ACTIVATIONMAP']._serialized_start=4604 + _globals['_ACTIVATIONMAP']._serialized_end=4676 + _globals['_ACTIVATIONRESPONSE']._serialized_start=4678 + _globals['_ACTIVATIONRESPONSE']._serialized_end=4778 + _globals['_TASKFIELD']._serialized_start=4781 + _globals['_TASKFIELD']._serialized_end=4928 + _globals['_RECORDMETADATA']._serialized_start=4931 + _globals['_RECORDMETADATA']._serialized_end=5322 + _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._serialized_start=5269 + _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._serialized_end=5322 + _globals['_SAMPLESTATISTICS']._serialized_start=5325 + _globals['_SAMPLESTATISTICS']._serialized_end=5472 + _globals['_COMMANDRESPONSE']._serialized_start=5475 + _globals['_COMMANDRESPONSE']._serialized_end=5705 + _globals['_SAMPLEREQUEST']._serialized_start=5707 + _globals['_SAMPLEREQUEST']._serialized_end=5792 + _globals['_SAMPLEREQUESTRESPONSE']._serialized_start=5795 + _globals['_SAMPLEREQUESTRESPONSE']._serialized_end=6096 + _globals['_BATCHSAMPLEREQUEST']._serialized_start=6099 + _globals['_BATCHSAMPLEREQUEST']._serialized_end=6245 + _globals['_BATCHSAMPLERESPONSE']._serialized_start=6247 + _globals['_BATCHSAMPLERESPONSE']._serialized_end=6309 + _globals['_WEIGHTSREQUEST']._serialized_start=6311 + _globals['_WEIGHTSREQUEST']._serialized_end=6357 + _globals['_WEIGHTSRESPONSE']._serialized_start=6360 + _globals['_WEIGHTSRESPONSE']._serialized_end=6645 + _globals['_DATAQUERYREQUEST']._serialized_start=6647 + _globals['_DATAQUERYREQUEST']._serialized_end=6729 + _globals['_CATEGORICALTAGDEF']._serialized_start=6731 + _globals['_CATEGORICALTAGDEF']._serialized_end=6784 + _globals['_DATAQUERYRESPONSE']._serialized_start=6787 + _globals['_DATAQUERYRESPONSE']._serialized_end=7084 + _globals['_DATASAMPLESREQUEST']._serialized_start=7087 + _globals['_DATASAMPLESREQUEST']._serialized_end=7281 + _globals['_DATASTAT']._serialized_start=7283 + _globals['_DATASTAT']._serialized_end=7392 + _globals['_DATARECORD']._serialized_start=7394 + _globals['_DATARECORD']._serialized_end=7456 + _globals['_DATASAMPLESRESPONSE']._serialized_start=7458 + _globals['_DATASAMPLESRESPONSE']._serialized_end=7548 + _globals['_HISTOGRAMSUBBAR']._serialized_start=7550 + _globals['_HISTOGRAMSUBBAR']._serialized_end=7617 + _globals['_HISTOGRAMBIN']._serialized_start=7619 + _globals['_HISTOGRAMBIN']._serialized_end=7723 + _globals['_CATEGORICALHISTOGRAMBAR']._serialized_start=7725 + _globals['_CATEGORICALHISTOGRAMBAR']._serialized_end=7816 + _globals['_HISTOGRAMREQUEST']._serialized_start=7818 + _globals['_HISTOGRAMREQUEST']._serialized_end=7870 + _globals['_HISTOGRAMRESPONSE']._serialized_start=7873 + _globals['_HISTOGRAMRESPONSE']._serialized_end=8051 + _globals['_GETMETADATAREQUEST']._serialized_start=8053 + _globals['_GETMETADATAREQUEST']._serialized_end=8140 + _globals['_GETMETADATARESPONSE']._serialized_start=8143 + _globals['_GETMETADATARESPONSE']._serialized_end=8296 + _globals['_GETSIGNALTRAJECTORYREQUEST']._serialized_start=8298 + _globals['_GETSIGNALTRAJECTORYREQUEST']._serialized_end=8387 + _globals['_SIGNALTRAJECTORY']._serialized_start=8389 + _globals['_SIGNALTRAJECTORY']._serialized_end=8441 + _globals['_GETSIGNALTRAJECTORYRESPONSE']._serialized_start=8443 + _globals['_GETSIGNALTRAJECTORYRESPONSE']._serialized_end=8568 + _globals['_POINTCLOUDREQUEST']._serialized_start=8570 + _globals['_POINTCLOUDREQUEST']._serialized_end=8644 + _globals['_POINTCLOUDCHUNK']._serialized_start=8647 + _globals['_POINTCLOUDCHUNK']._serialized_end=8838 + _globals['_DATAEDITSREQUEST']._serialized_start=8841 + _globals['_DATAEDITSREQUEST']._serialized_end=9061 + _globals['_DATAEDITSRESPONSE']._serialized_start=9063 + _globals['_DATAEDITSRESPONSE']._serialized_end=9116 + _globals['_DATASPLITSRESPONSE']._serialized_start=9118 + _globals['_DATASPLITSRESPONSE']._serialized_end=9176 + _globals['_AGENTHEALTHRESPONSE']._serialized_start=9178 + _globals['_AGENTHEALTHRESPONSE']._serialized_end=9235 + _globals['_INITIALIZEAGENTREQUEST']._serialized_start=9237 + _globals['_INITIALIZEAGENTREQUEST']._serialized_end=9331 + _globals['_INITIALIZEAGENTRESPONSE']._serialized_start=9333 + _globals['_INITIALIZEAGENTRESPONSE']._serialized_end=9392 + _globals['_CHANGEAGENTMODELREQUEST']._serialized_start=9394 + _globals['_CHANGEAGENTMODELREQUEST']._serialized_end=9434 + _globals['_CHANGEAGENTMODELRESPONSE']._serialized_start=9436 + _globals['_CHANGEAGENTMODELRESPONSE']._serialized_end=9496 + _globals['_GETAGENTMODELSREQUEST']._serialized_start=9498 + _globals['_GETAGENTMODELSREQUEST']._serialized_end=9521 + _globals['_GETAGENTMODELSRESPONSE']._serialized_start=9523 + _globals['_GETAGENTMODELSRESPONSE']._serialized_end=9597 + _globals['_RESETAGENTRESPONSE']._serialized_start=9599 + _globals['_RESETAGENTRESPONSE']._serialized_end=9653 + _globals['_CLEARAGENTHISTORYRESPONSE']._serialized_start=9655 + _globals['_CLEARAGENTHISTORYRESPONSE']._serialized_end=9716 + _globals['_COMPACTAGENTHISTORYRESPONSE']._serialized_start=9718 + _globals['_COMPACTAGENTHISTORYRESPONSE']._serialized_end=9781 + _globals['_RESTORECHECKPOINTREQUEST']._serialized_start=9783 + _globals['_RESTORECHECKPOINTREQUEST']._serialized_end=9834 + _globals['_RESTORECHECKPOINTRESPONSE']._serialized_start=9836 + _globals['_RESTORECHECKPOINTRESPONSE']._serialized_end=9897 + _globals['_TRIGGEREVALUATIONREQUEST']._serialized_start=9899 + _globals['_TRIGGEREVALUATIONREQUEST']._serialized_end=9981 + _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_start=9983 + _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_end=10044 + _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_start=10046 + _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_end=10074 + _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_start=10077 + _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_end=10206 + _globals['_CANCELEVALUATIONREQUEST']._serialized_start=10208 + _globals['_CANCELEVALUATIONREQUEST']._serialized_end=10249 + _globals['_CANCELEVALUATIONRESPONSE']._serialized_start=10251 + _globals['_CANCELEVALUATIONRESPONSE']._serialized_end=10311 + _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_start=10313 + _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_end=10368 + _globals['_NOTEBOOKCELLDONE']._serialized_start=10370 + _globals['_NOTEBOOKCELLDONE']._serialized_end=10420 + _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_start=10422 + _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_end=10452 + _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_start=10454 + _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_end=10512 + _globals['_NOTEBOOKCELLCHUNK']._serialized_start=10515 + _globals['_NOTEBOOKCELLCHUNK']._serialized_end=10704 + _globals['_NOTEBOOKRESPONSE']._serialized_start=10706 + _globals['_NOTEBOOKRESPONSE']._serialized_end=10789 + _globals['_SAVENOTEBOOKREQUEST']._serialized_start=10791 + _globals['_SAVENOTEBOOKREQUEST']._serialized_end=10846 + _globals['_SAVENOTEBOOKRESPONSE']._serialized_start=10848 + _globals['_SAVENOTEBOOKRESPONSE']._serialized_end=10925 + _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_start=10927 + _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_end=10994 + _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_start=10996 + _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_end=11088 + _globals['_EXPERIMENTSERVICE']._serialized_start=11529 + _globals['_EXPERIMENTSERVICE']._serialized_end=13468 # @@protoc_insertion_point(module_scope) diff --git a/weightslab/proto/experiment_service_pb2_grpc.py b/weightslab/proto/experiment_service_pb2_grpc.py index 2f9127b2..55830bca 100644 --- a/weightslab/proto/experiment_service_pb2_grpc.py +++ b/weightslab/proto/experiment_service_pb2_grpc.py @@ -3,9 +3,9 @@ import grpc import warnings -from weightslab.proto import experiment_service_pb2 as weightslab_dot_proto_dot_experiment__service__pb2 +from weightslab.proto import experiment_service_pb2 as experiment__service__pb2 -GRPC_GENERATED_VERSION = '1.76.0' +GRPC_GENERATED_VERSION = '1.68.1' GRPC_VERSION = grpc.__version__ _version_not_supported = False @@ -18,7 +18,7 @@ if _version_not_supported: raise RuntimeError( f'The grpc package installed is at version {GRPC_VERSION},' - + ' but the generated code in weightslab/proto/experiment_service_pb2_grpc.py depends on' + + f' but the generated code in experiment_service_pb2_grpc.py depends on' + f' grpcio>={GRPC_GENERATED_VERSION}.' + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' @@ -36,143 +36,153 @@ def __init__(self, channel): """ self.GetLatestLoggerData = channel.unary_unary( '/ExperimentService/GetLatestLoggerData', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetLatestLoggerDataRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetLatestLoggerDataResponse.FromString, + request_serializer=experiment__service__pb2.GetLatestLoggerDataRequest.SerializeToString, + response_deserializer=experiment__service__pb2.GetLatestLoggerDataResponse.FromString, _registered_method=True) self.ExperimentCommand = channel.unary_unary( '/ExperimentService/ExperimentCommand', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.TrainerCommand.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.CommandResponse.FromString, + request_serializer=experiment__service__pb2.TrainerCommand.SerializeToString, + response_deserializer=experiment__service__pb2.CommandResponse.FromString, _registered_method=True) self.ManipulateWeights = channel.unary_unary( '/ExperimentService/ManipulateWeights', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsOperationRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsOperationResponse.FromString, + request_serializer=experiment__service__pb2.WeightsOperationRequest.SerializeToString, + response_deserializer=experiment__service__pb2.WeightsOperationResponse.FromString, _registered_method=True) self.GetWeights = channel.unary_unary( '/ExperimentService/GetWeights', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsResponse.FromString, + request_serializer=experiment__service__pb2.WeightsRequest.SerializeToString, + response_deserializer=experiment__service__pb2.WeightsResponse.FromString, _registered_method=True) self.GetActivations = channel.unary_unary( '/ExperimentService/GetActivations', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ActivationRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ActivationResponse.FromString, + request_serializer=experiment__service__pb2.ActivationRequest.SerializeToString, + response_deserializer=experiment__service__pb2.ActivationResponse.FromString, _registered_method=True) self.GetSamples = channel.unary_unary( '/ExperimentService/GetSamples', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.BatchSampleRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.BatchSampleResponse.FromString, + request_serializer=experiment__service__pb2.BatchSampleRequest.SerializeToString, + response_deserializer=experiment__service__pb2.BatchSampleResponse.FromString, _registered_method=True) self.ApplyDataQuery = channel.unary_unary( '/ExperimentService/ApplyDataQuery', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataQueryRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataQueryResponse.FromString, + request_serializer=experiment__service__pb2.DataQueryRequest.SerializeToString, + response_deserializer=experiment__service__pb2.DataQueryResponse.FromString, _registered_method=True) self.GetDataSamples = channel.unary_unary( '/ExperimentService/GetDataSamples', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataSamplesRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataSamplesResponse.FromString, + request_serializer=experiment__service__pb2.DataSamplesRequest.SerializeToString, + response_deserializer=experiment__service__pb2.DataSamplesResponse.FromString, _registered_method=True) self.GetHistogram = channel.unary_unary( '/ExperimentService/GetHistogram', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.HistogramRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.HistogramResponse.FromString, + request_serializer=experiment__service__pb2.HistogramRequest.SerializeToString, + response_deserializer=experiment__service__pb2.HistogramResponse.FromString, _registered_method=True) self.GetMetaData = channel.unary_unary( '/ExperimentService/GetMetaData', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetMetaDataRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetMetaDataResponse.FromString, + request_serializer=experiment__service__pb2.GetMetaDataRequest.SerializeToString, + response_deserializer=experiment__service__pb2.GetMetaDataResponse.FromString, _registered_method=True) self.GetSignalTrajectory = channel.unary_unary( '/ExperimentService/GetSignalTrajectory', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetSignalTrajectoryRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetSignalTrajectoryResponse.FromString, + request_serializer=experiment__service__pb2.GetSignalTrajectoryRequest.SerializeToString, + response_deserializer=experiment__service__pb2.GetSignalTrajectoryResponse.FromString, _registered_method=True) self.GetPointCloud = channel.unary_stream( '/ExperimentService/GetPointCloud', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudChunk.FromString, + request_serializer=experiment__service__pb2.PointCloudRequest.SerializeToString, + response_deserializer=experiment__service__pb2.PointCloudChunk.FromString, _registered_method=True) self.EditDataSample = channel.unary_unary( '/ExperimentService/EditDataSample', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataEditsRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataEditsResponse.FromString, + request_serializer=experiment__service__pb2.DataEditsRequest.SerializeToString, + response_deserializer=experiment__service__pb2.DataEditsResponse.FromString, _registered_method=True) self.GetDataSplits = channel.unary_unary( '/ExperimentService/GetDataSplits', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataSplitsResponse.FromString, + request_serializer=experiment__service__pb2.Empty.SerializeToString, + response_deserializer=experiment__service__pb2.DataSplitsResponse.FromString, _registered_method=True) self.CheckAgentHealth = channel.unary_unary( '/ExperimentService/CheckAgentHealth', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.AgentHealthResponse.FromString, + request_serializer=experiment__service__pb2.Empty.SerializeToString, + response_deserializer=experiment__service__pb2.AgentHealthResponse.FromString, _registered_method=True) self.InitializeAgent = channel.unary_unary( '/ExperimentService/InitializeAgent', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.InitializeAgentRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.InitializeAgentResponse.FromString, + request_serializer=experiment__service__pb2.InitializeAgentRequest.SerializeToString, + response_deserializer=experiment__service__pb2.InitializeAgentResponse.FromString, _registered_method=True) self.ChangeAgentModel = channel.unary_unary( '/ExperimentService/ChangeAgentModel', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ChangeAgentModelRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ChangeAgentModelResponse.FromString, + request_serializer=experiment__service__pb2.ChangeAgentModelRequest.SerializeToString, + response_deserializer=experiment__service__pb2.ChangeAgentModelResponse.FromString, _registered_method=True) self.GetAgentModels = channel.unary_unary( '/ExperimentService/GetAgentModels', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetAgentModelsRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetAgentModelsResponse.FromString, + request_serializer=experiment__service__pb2.GetAgentModelsRequest.SerializeToString, + response_deserializer=experiment__service__pb2.GetAgentModelsResponse.FromString, _registered_method=True) self.ResetAgent = channel.unary_unary( '/ExperimentService/ResetAgent', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ResetAgentResponse.FromString, + request_serializer=experiment__service__pb2.Empty.SerializeToString, + response_deserializer=experiment__service__pb2.ResetAgentResponse.FromString, + _registered_method=True) + self.ClearAgentHistory = channel.unary_unary( + '/ExperimentService/ClearAgentHistory', + request_serializer=experiment__service__pb2.Empty.SerializeToString, + response_deserializer=experiment__service__pb2.ClearAgentHistoryResponse.FromString, + _registered_method=True) + self.CompactAgentHistory = channel.unary_unary( + '/ExperimentService/CompactAgentHistory', + request_serializer=experiment__service__pb2.Empty.SerializeToString, + response_deserializer=experiment__service__pb2.CompactAgentHistoryResponse.FromString, _registered_method=True) self.RunNotebookCell = channel.unary_stream( '/ExperimentService/RunNotebookCell', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.RunNotebookCellRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.NotebookCellChunk.FromString, + request_serializer=experiment__service__pb2.RunNotebookCellRequest.SerializeToString, + response_deserializer=experiment__service__pb2.NotebookCellChunk.FromString, _registered_method=True) self.InterruptNotebookCell = channel.unary_unary( '/ExperimentService/InterruptNotebookCell', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.InterruptNotebookCellRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.InterruptNotebookCellResponse.FromString, + request_serializer=experiment__service__pb2.InterruptNotebookCellRequest.SerializeToString, + response_deserializer=experiment__service__pb2.InterruptNotebookCellResponse.FromString, _registered_method=True) self.GetNotebook = channel.unary_unary( '/ExperimentService/GetNotebook', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.NotebookResponse.FromString, + request_serializer=experiment__service__pb2.Empty.SerializeToString, + response_deserializer=experiment__service__pb2.NotebookResponse.FromString, _registered_method=True) self.SaveNotebook = channel.unary_unary( '/ExperimentService/SaveNotebook', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.SaveNotebookRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.SaveNotebookResponse.FromString, + request_serializer=experiment__service__pb2.SaveNotebookRequest.SerializeToString, + response_deserializer=experiment__service__pb2.SaveNotebookResponse.FromString, _registered_method=True) self.GenerateNotebookCode = channel.unary_unary( '/ExperimentService/GenerateNotebookCode', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GenerateNotebookCodeRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GenerateNotebookCodeResponse.FromString, + request_serializer=experiment__service__pb2.GenerateNotebookCodeRequest.SerializeToString, + response_deserializer=experiment__service__pb2.GenerateNotebookCodeResponse.FromString, _registered_method=True) self.RestoreCheckpoint = channel.unary_unary( '/ExperimentService/RestoreCheckpoint', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.RestoreCheckpointRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.RestoreCheckpointResponse.FromString, + request_serializer=experiment__service__pb2.RestoreCheckpointRequest.SerializeToString, + response_deserializer=experiment__service__pb2.RestoreCheckpointResponse.FromString, _registered_method=True) self.TriggerEvaluation = channel.unary_unary( '/ExperimentService/TriggerEvaluation', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.TriggerEvaluationRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.TriggerEvaluationResponse.FromString, + request_serializer=experiment__service__pb2.TriggerEvaluationRequest.SerializeToString, + response_deserializer=experiment__service__pb2.TriggerEvaluationResponse.FromString, _registered_method=True) self.GetEvaluationStatus = channel.unary_unary( '/ExperimentService/GetEvaluationStatus', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetEvaluationStatusRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetEvaluationStatusResponse.FromString, + request_serializer=experiment__service__pb2.GetEvaluationStatusRequest.SerializeToString, + response_deserializer=experiment__service__pb2.GetEvaluationStatusResponse.FromString, _registered_method=True) self.CancelEvaluation = channel.unary_unary( '/ExperimentService/CancelEvaluation', - request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.CancelEvaluationRequest.SerializeToString, - response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.CancelEvaluationResponse.FromString, + request_serializer=experiment__service__pb2.CancelEvaluationRequest.SerializeToString, + response_deserializer=experiment__service__pb2.CancelEvaluationResponse.FromString, _registered_method=True) @@ -306,6 +316,23 @@ def ResetAgent(self, request, context): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') + def ClearAgentHistory(self, request, context): + """Wipe the agent's conversation history (self.history) without touching the + provider connection -- distinct from ResetAgent, which drops the connection. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CompactAgentHistory(self, request, context): + """Summarize the agent's conversation history via the active model, replacing + it with the summary. Distinct from OpenCode's own session compaction (which + this does not touch) -- this is the SDK agent's own self.history. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + def RunNotebookCell(self, request, context): """Notebook (shared in-process Python kernel for the studio UI). Cell output is server-streamed in chunks (stdout / stderr / result / image), mirroring @@ -375,143 +402,153 @@ def add_ExperimentServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'GetLatestLoggerData': grpc.unary_unary_rpc_method_handler( servicer.GetLatestLoggerData, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetLatestLoggerDataRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetLatestLoggerDataResponse.SerializeToString, + request_deserializer=experiment__service__pb2.GetLatestLoggerDataRequest.FromString, + response_serializer=experiment__service__pb2.GetLatestLoggerDataResponse.SerializeToString, ), 'ExperimentCommand': grpc.unary_unary_rpc_method_handler( servicer.ExperimentCommand, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.TrainerCommand.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.CommandResponse.SerializeToString, + request_deserializer=experiment__service__pb2.TrainerCommand.FromString, + response_serializer=experiment__service__pb2.CommandResponse.SerializeToString, ), 'ManipulateWeights': grpc.unary_unary_rpc_method_handler( servicer.ManipulateWeights, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsOperationRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsOperationResponse.SerializeToString, + request_deserializer=experiment__service__pb2.WeightsOperationRequest.FromString, + response_serializer=experiment__service__pb2.WeightsOperationResponse.SerializeToString, ), 'GetWeights': grpc.unary_unary_rpc_method_handler( servicer.GetWeights, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.WeightsResponse.SerializeToString, + request_deserializer=experiment__service__pb2.WeightsRequest.FromString, + response_serializer=experiment__service__pb2.WeightsResponse.SerializeToString, ), 'GetActivations': grpc.unary_unary_rpc_method_handler( servicer.GetActivations, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ActivationRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ActivationResponse.SerializeToString, + request_deserializer=experiment__service__pb2.ActivationRequest.FromString, + response_serializer=experiment__service__pb2.ActivationResponse.SerializeToString, ), 'GetSamples': grpc.unary_unary_rpc_method_handler( servicer.GetSamples, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.BatchSampleRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.BatchSampleResponse.SerializeToString, + request_deserializer=experiment__service__pb2.BatchSampleRequest.FromString, + response_serializer=experiment__service__pb2.BatchSampleResponse.SerializeToString, ), 'ApplyDataQuery': grpc.unary_unary_rpc_method_handler( servicer.ApplyDataQuery, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataQueryRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataQueryResponse.SerializeToString, + request_deserializer=experiment__service__pb2.DataQueryRequest.FromString, + response_serializer=experiment__service__pb2.DataQueryResponse.SerializeToString, ), 'GetDataSamples': grpc.unary_unary_rpc_method_handler( servicer.GetDataSamples, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataSamplesRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataSamplesResponse.SerializeToString, + request_deserializer=experiment__service__pb2.DataSamplesRequest.FromString, + response_serializer=experiment__service__pb2.DataSamplesResponse.SerializeToString, ), 'GetHistogram': grpc.unary_unary_rpc_method_handler( servicer.GetHistogram, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.HistogramRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.HistogramResponse.SerializeToString, + request_deserializer=experiment__service__pb2.HistogramRequest.FromString, + response_serializer=experiment__service__pb2.HistogramResponse.SerializeToString, ), 'GetMetaData': grpc.unary_unary_rpc_method_handler( servicer.GetMetaData, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetMetaDataRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetMetaDataResponse.SerializeToString, + request_deserializer=experiment__service__pb2.GetMetaDataRequest.FromString, + response_serializer=experiment__service__pb2.GetMetaDataResponse.SerializeToString, ), 'GetSignalTrajectory': grpc.unary_unary_rpc_method_handler( servicer.GetSignalTrajectory, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetSignalTrajectoryRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetSignalTrajectoryResponse.SerializeToString, + request_deserializer=experiment__service__pb2.GetSignalTrajectoryRequest.FromString, + response_serializer=experiment__service__pb2.GetSignalTrajectoryResponse.SerializeToString, ), 'GetPointCloud': grpc.unary_stream_rpc_method_handler( servicer.GetPointCloud, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudChunk.SerializeToString, + request_deserializer=experiment__service__pb2.PointCloudRequest.FromString, + response_serializer=experiment__service__pb2.PointCloudChunk.SerializeToString, ), 'EditDataSample': grpc.unary_unary_rpc_method_handler( servicer.EditDataSample, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataEditsRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataEditsResponse.SerializeToString, + request_deserializer=experiment__service__pb2.DataEditsRequest.FromString, + response_serializer=experiment__service__pb2.DataEditsResponse.SerializeToString, ), 'GetDataSplits': grpc.unary_unary_rpc_method_handler( servicer.GetDataSplits, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataSplitsResponse.SerializeToString, + request_deserializer=experiment__service__pb2.Empty.FromString, + response_serializer=experiment__service__pb2.DataSplitsResponse.SerializeToString, ), 'CheckAgentHealth': grpc.unary_unary_rpc_method_handler( servicer.CheckAgentHealth, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.AgentHealthResponse.SerializeToString, + request_deserializer=experiment__service__pb2.Empty.FromString, + response_serializer=experiment__service__pb2.AgentHealthResponse.SerializeToString, ), 'InitializeAgent': grpc.unary_unary_rpc_method_handler( servicer.InitializeAgent, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.InitializeAgentRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.InitializeAgentResponse.SerializeToString, + request_deserializer=experiment__service__pb2.InitializeAgentRequest.FromString, + response_serializer=experiment__service__pb2.InitializeAgentResponse.SerializeToString, ), 'ChangeAgentModel': grpc.unary_unary_rpc_method_handler( servicer.ChangeAgentModel, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ChangeAgentModelRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ChangeAgentModelResponse.SerializeToString, + request_deserializer=experiment__service__pb2.ChangeAgentModelRequest.FromString, + response_serializer=experiment__service__pb2.ChangeAgentModelResponse.SerializeToString, ), 'GetAgentModels': grpc.unary_unary_rpc_method_handler( servicer.GetAgentModels, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetAgentModelsRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetAgentModelsResponse.SerializeToString, + request_deserializer=experiment__service__pb2.GetAgentModelsRequest.FromString, + response_serializer=experiment__service__pb2.GetAgentModelsResponse.SerializeToString, ), 'ResetAgent': grpc.unary_unary_rpc_method_handler( servicer.ResetAgent, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ResetAgentResponse.SerializeToString, + request_deserializer=experiment__service__pb2.Empty.FromString, + response_serializer=experiment__service__pb2.ResetAgentResponse.SerializeToString, + ), + 'ClearAgentHistory': grpc.unary_unary_rpc_method_handler( + servicer.ClearAgentHistory, + request_deserializer=experiment__service__pb2.Empty.FromString, + response_serializer=experiment__service__pb2.ClearAgentHistoryResponse.SerializeToString, + ), + 'CompactAgentHistory': grpc.unary_unary_rpc_method_handler( + servicer.CompactAgentHistory, + request_deserializer=experiment__service__pb2.Empty.FromString, + response_serializer=experiment__service__pb2.CompactAgentHistoryResponse.SerializeToString, ), 'RunNotebookCell': grpc.unary_stream_rpc_method_handler( servicer.RunNotebookCell, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.RunNotebookCellRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.NotebookCellChunk.SerializeToString, + request_deserializer=experiment__service__pb2.RunNotebookCellRequest.FromString, + response_serializer=experiment__service__pb2.NotebookCellChunk.SerializeToString, ), 'InterruptNotebookCell': grpc.unary_unary_rpc_method_handler( servicer.InterruptNotebookCell, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.InterruptNotebookCellRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.InterruptNotebookCellResponse.SerializeToString, + request_deserializer=experiment__service__pb2.InterruptNotebookCellRequest.FromString, + response_serializer=experiment__service__pb2.InterruptNotebookCellResponse.SerializeToString, ), 'GetNotebook': grpc.unary_unary_rpc_method_handler( servicer.GetNotebook, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.Empty.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.NotebookResponse.SerializeToString, + request_deserializer=experiment__service__pb2.Empty.FromString, + response_serializer=experiment__service__pb2.NotebookResponse.SerializeToString, ), 'SaveNotebook': grpc.unary_unary_rpc_method_handler( servicer.SaveNotebook, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.SaveNotebookRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.SaveNotebookResponse.SerializeToString, + request_deserializer=experiment__service__pb2.SaveNotebookRequest.FromString, + response_serializer=experiment__service__pb2.SaveNotebookResponse.SerializeToString, ), 'GenerateNotebookCode': grpc.unary_unary_rpc_method_handler( servicer.GenerateNotebookCode, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GenerateNotebookCodeRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GenerateNotebookCodeResponse.SerializeToString, + request_deserializer=experiment__service__pb2.GenerateNotebookCodeRequest.FromString, + response_serializer=experiment__service__pb2.GenerateNotebookCodeResponse.SerializeToString, ), 'RestoreCheckpoint': grpc.unary_unary_rpc_method_handler( servicer.RestoreCheckpoint, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.RestoreCheckpointRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.RestoreCheckpointResponse.SerializeToString, + request_deserializer=experiment__service__pb2.RestoreCheckpointRequest.FromString, + response_serializer=experiment__service__pb2.RestoreCheckpointResponse.SerializeToString, ), 'TriggerEvaluation': grpc.unary_unary_rpc_method_handler( servicer.TriggerEvaluation, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.TriggerEvaluationRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.TriggerEvaluationResponse.SerializeToString, + request_deserializer=experiment__service__pb2.TriggerEvaluationRequest.FromString, + response_serializer=experiment__service__pb2.TriggerEvaluationResponse.SerializeToString, ), 'GetEvaluationStatus': grpc.unary_unary_rpc_method_handler( servicer.GetEvaluationStatus, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.GetEvaluationStatusRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.GetEvaluationStatusResponse.SerializeToString, + request_deserializer=experiment__service__pb2.GetEvaluationStatusRequest.FromString, + response_serializer=experiment__service__pb2.GetEvaluationStatusResponse.SerializeToString, ), 'CancelEvaluation': grpc.unary_unary_rpc_method_handler( servicer.CancelEvaluation, - request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.CancelEvaluationRequest.FromString, - response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.CancelEvaluationResponse.SerializeToString, + request_deserializer=experiment__service__pb2.CancelEvaluationRequest.FromString, + response_serializer=experiment__service__pb2.CancelEvaluationResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( @@ -539,8 +576,8 @@ def GetLatestLoggerData(request, request, target, '/ExperimentService/GetLatestLoggerData', - weightslab_dot_proto_dot_experiment__service__pb2.GetLatestLoggerDataRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.GetLatestLoggerDataResponse.FromString, + experiment__service__pb2.GetLatestLoggerDataRequest.SerializeToString, + experiment__service__pb2.GetLatestLoggerDataResponse.FromString, options, channel_credentials, insecure, @@ -566,8 +603,8 @@ def ExperimentCommand(request, request, target, '/ExperimentService/ExperimentCommand', - weightslab_dot_proto_dot_experiment__service__pb2.TrainerCommand.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.CommandResponse.FromString, + experiment__service__pb2.TrainerCommand.SerializeToString, + experiment__service__pb2.CommandResponse.FromString, options, channel_credentials, insecure, @@ -593,8 +630,8 @@ def ManipulateWeights(request, request, target, '/ExperimentService/ManipulateWeights', - weightslab_dot_proto_dot_experiment__service__pb2.WeightsOperationRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.WeightsOperationResponse.FromString, + experiment__service__pb2.WeightsOperationRequest.SerializeToString, + experiment__service__pb2.WeightsOperationResponse.FromString, options, channel_credentials, insecure, @@ -620,8 +657,8 @@ def GetWeights(request, request, target, '/ExperimentService/GetWeights', - weightslab_dot_proto_dot_experiment__service__pb2.WeightsRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.WeightsResponse.FromString, + experiment__service__pb2.WeightsRequest.SerializeToString, + experiment__service__pb2.WeightsResponse.FromString, options, channel_credentials, insecure, @@ -647,8 +684,8 @@ def GetActivations(request, request, target, '/ExperimentService/GetActivations', - weightslab_dot_proto_dot_experiment__service__pb2.ActivationRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.ActivationResponse.FromString, + experiment__service__pb2.ActivationRequest.SerializeToString, + experiment__service__pb2.ActivationResponse.FromString, options, channel_credentials, insecure, @@ -674,8 +711,8 @@ def GetSamples(request, request, target, '/ExperimentService/GetSamples', - weightslab_dot_proto_dot_experiment__service__pb2.BatchSampleRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.BatchSampleResponse.FromString, + experiment__service__pb2.BatchSampleRequest.SerializeToString, + experiment__service__pb2.BatchSampleResponse.FromString, options, channel_credentials, insecure, @@ -701,8 +738,8 @@ def ApplyDataQuery(request, request, target, '/ExperimentService/ApplyDataQuery', - weightslab_dot_proto_dot_experiment__service__pb2.DataQueryRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.DataQueryResponse.FromString, + experiment__service__pb2.DataQueryRequest.SerializeToString, + experiment__service__pb2.DataQueryResponse.FromString, options, channel_credentials, insecure, @@ -728,8 +765,8 @@ def GetDataSamples(request, request, target, '/ExperimentService/GetDataSamples', - weightslab_dot_proto_dot_experiment__service__pb2.DataSamplesRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.DataSamplesResponse.FromString, + experiment__service__pb2.DataSamplesRequest.SerializeToString, + experiment__service__pb2.DataSamplesResponse.FromString, options, channel_credentials, insecure, @@ -755,8 +792,8 @@ def GetHistogram(request, request, target, '/ExperimentService/GetHistogram', - weightslab_dot_proto_dot_experiment__service__pb2.HistogramRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.HistogramResponse.FromString, + experiment__service__pb2.HistogramRequest.SerializeToString, + experiment__service__pb2.HistogramResponse.FromString, options, channel_credentials, insecure, @@ -782,8 +819,8 @@ def GetMetaData(request, request, target, '/ExperimentService/GetMetaData', - weightslab_dot_proto_dot_experiment__service__pb2.GetMetaDataRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.GetMetaDataResponse.FromString, + experiment__service__pb2.GetMetaDataRequest.SerializeToString, + experiment__service__pb2.GetMetaDataResponse.FromString, options, channel_credentials, insecure, @@ -809,8 +846,8 @@ def GetSignalTrajectory(request, request, target, '/ExperimentService/GetSignalTrajectory', - weightslab_dot_proto_dot_experiment__service__pb2.GetSignalTrajectoryRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.GetSignalTrajectoryResponse.FromString, + experiment__service__pb2.GetSignalTrajectoryRequest.SerializeToString, + experiment__service__pb2.GetSignalTrajectoryResponse.FromString, options, channel_credentials, insecure, @@ -836,8 +873,8 @@ def GetPointCloud(request, request, target, '/ExperimentService/GetPointCloud', - weightslab_dot_proto_dot_experiment__service__pb2.PointCloudRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.PointCloudChunk.FromString, + experiment__service__pb2.PointCloudRequest.SerializeToString, + experiment__service__pb2.PointCloudChunk.FromString, options, channel_credentials, insecure, @@ -863,8 +900,8 @@ def EditDataSample(request, request, target, '/ExperimentService/EditDataSample', - weightslab_dot_proto_dot_experiment__service__pb2.DataEditsRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.DataEditsResponse.FromString, + experiment__service__pb2.DataEditsRequest.SerializeToString, + experiment__service__pb2.DataEditsResponse.FromString, options, channel_credentials, insecure, @@ -890,8 +927,8 @@ def GetDataSplits(request, request, target, '/ExperimentService/GetDataSplits', - weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.DataSplitsResponse.FromString, + experiment__service__pb2.Empty.SerializeToString, + experiment__service__pb2.DataSplitsResponse.FromString, options, channel_credentials, insecure, @@ -917,8 +954,8 @@ def CheckAgentHealth(request, request, target, '/ExperimentService/CheckAgentHealth', - weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.AgentHealthResponse.FromString, + experiment__service__pb2.Empty.SerializeToString, + experiment__service__pb2.AgentHealthResponse.FromString, options, channel_credentials, insecure, @@ -944,8 +981,8 @@ def InitializeAgent(request, request, target, '/ExperimentService/InitializeAgent', - weightslab_dot_proto_dot_experiment__service__pb2.InitializeAgentRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.InitializeAgentResponse.FromString, + experiment__service__pb2.InitializeAgentRequest.SerializeToString, + experiment__service__pb2.InitializeAgentResponse.FromString, options, channel_credentials, insecure, @@ -971,8 +1008,8 @@ def ChangeAgentModel(request, request, target, '/ExperimentService/ChangeAgentModel', - weightslab_dot_proto_dot_experiment__service__pb2.ChangeAgentModelRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.ChangeAgentModelResponse.FromString, + experiment__service__pb2.ChangeAgentModelRequest.SerializeToString, + experiment__service__pb2.ChangeAgentModelResponse.FromString, options, channel_credentials, insecure, @@ -998,8 +1035,8 @@ def GetAgentModels(request, request, target, '/ExperimentService/GetAgentModels', - weightslab_dot_proto_dot_experiment__service__pb2.GetAgentModelsRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.GetAgentModelsResponse.FromString, + experiment__service__pb2.GetAgentModelsRequest.SerializeToString, + experiment__service__pb2.GetAgentModelsResponse.FromString, options, channel_credentials, insecure, @@ -1025,8 +1062,62 @@ def ResetAgent(request, request, target, '/ExperimentService/ResetAgent', - weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.ResetAgentResponse.FromString, + experiment__service__pb2.Empty.SerializeToString, + experiment__service__pb2.ResetAgentResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def ClearAgentHistory(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/ExperimentService/ClearAgentHistory', + experiment__service__pb2.Empty.SerializeToString, + experiment__service__pb2.ClearAgentHistoryResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def CompactAgentHistory(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/ExperimentService/CompactAgentHistory', + experiment__service__pb2.Empty.SerializeToString, + experiment__service__pb2.CompactAgentHistoryResponse.FromString, options, channel_credentials, insecure, @@ -1052,8 +1143,8 @@ def RunNotebookCell(request, request, target, '/ExperimentService/RunNotebookCell', - weightslab_dot_proto_dot_experiment__service__pb2.RunNotebookCellRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.NotebookCellChunk.FromString, + experiment__service__pb2.RunNotebookCellRequest.SerializeToString, + experiment__service__pb2.NotebookCellChunk.FromString, options, channel_credentials, insecure, @@ -1079,8 +1170,8 @@ def InterruptNotebookCell(request, request, target, '/ExperimentService/InterruptNotebookCell', - weightslab_dot_proto_dot_experiment__service__pb2.InterruptNotebookCellRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.InterruptNotebookCellResponse.FromString, + experiment__service__pb2.InterruptNotebookCellRequest.SerializeToString, + experiment__service__pb2.InterruptNotebookCellResponse.FromString, options, channel_credentials, insecure, @@ -1106,8 +1197,8 @@ def GetNotebook(request, request, target, '/ExperimentService/GetNotebook', - weightslab_dot_proto_dot_experiment__service__pb2.Empty.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.NotebookResponse.FromString, + experiment__service__pb2.Empty.SerializeToString, + experiment__service__pb2.NotebookResponse.FromString, options, channel_credentials, insecure, @@ -1133,8 +1224,8 @@ def SaveNotebook(request, request, target, '/ExperimentService/SaveNotebook', - weightslab_dot_proto_dot_experiment__service__pb2.SaveNotebookRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.SaveNotebookResponse.FromString, + experiment__service__pb2.SaveNotebookRequest.SerializeToString, + experiment__service__pb2.SaveNotebookResponse.FromString, options, channel_credentials, insecure, @@ -1160,8 +1251,8 @@ def GenerateNotebookCode(request, request, target, '/ExperimentService/GenerateNotebookCode', - weightslab_dot_proto_dot_experiment__service__pb2.GenerateNotebookCodeRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.GenerateNotebookCodeResponse.FromString, + experiment__service__pb2.GenerateNotebookCodeRequest.SerializeToString, + experiment__service__pb2.GenerateNotebookCodeResponse.FromString, options, channel_credentials, insecure, @@ -1187,8 +1278,8 @@ def RestoreCheckpoint(request, request, target, '/ExperimentService/RestoreCheckpoint', - weightslab_dot_proto_dot_experiment__service__pb2.RestoreCheckpointRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.RestoreCheckpointResponse.FromString, + experiment__service__pb2.RestoreCheckpointRequest.SerializeToString, + experiment__service__pb2.RestoreCheckpointResponse.FromString, options, channel_credentials, insecure, @@ -1214,8 +1305,8 @@ def TriggerEvaluation(request, request, target, '/ExperimentService/TriggerEvaluation', - weightslab_dot_proto_dot_experiment__service__pb2.TriggerEvaluationRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.TriggerEvaluationResponse.FromString, + experiment__service__pb2.TriggerEvaluationRequest.SerializeToString, + experiment__service__pb2.TriggerEvaluationResponse.FromString, options, channel_credentials, insecure, @@ -1241,8 +1332,8 @@ def GetEvaluationStatus(request, request, target, '/ExperimentService/GetEvaluationStatus', - weightslab_dot_proto_dot_experiment__service__pb2.GetEvaluationStatusRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.GetEvaluationStatusResponse.FromString, + experiment__service__pb2.GetEvaluationStatusRequest.SerializeToString, + experiment__service__pb2.GetEvaluationStatusResponse.FromString, options, channel_credentials, insecure, @@ -1268,8 +1359,8 @@ def CancelEvaluation(request, request, target, '/ExperimentService/CancelEvaluation', - weightslab_dot_proto_dot_experiment__service__pb2.CancelEvaluationRequest.SerializeToString, - weightslab_dot_proto_dot_experiment__service__pb2.CancelEvaluationResponse.FromString, + experiment__service__pb2.CancelEvaluationRequest.SerializeToString, + experiment__service__pb2.CancelEvaluationResponse.FromString, options, channel_credentials, insecure, diff --git a/weightslab/trainer/services/agent/agent.py b/weightslab/trainer/services/agent/agent.py index 8866d580..3dac1912 100644 --- a/weightslab/trainer/services/agent/agent.py +++ b/weightslab/trainer/services/agent/agent.py @@ -7,22 +7,12 @@ import logging import threading import pandas as pd -from urllib.parse import urlparse, urlunparse from abc import ABC, abstractmethod from typing import Optional, List, Union, Literal, Callable, Dict, Any from dotenv import load_dotenv from pathlib import Path -try: - from langchain_ollama import ChatOllama -except ImportError: - ChatOllama = None - -try: - from langchain_openai import ChatOpenAI -except ImportError: - ChatOpenAI = None from langchain_core.prompts import ChatPromptTemplate from pydantic import BaseModel, Field @@ -30,6 +20,7 @@ from .intent_prompt import INTENT_PROMPT from .notebook_prompt import NOTEBOOK_CODE_PROMPT from .report_prompt import REPORT_ANALYSIS_PROMPT +from .opencode_chat import OpenCodeChat from weightslab.data.sample_stats import SampleStatsEx from weightslab.trainer.trainer_tools import get_layer_representations @@ -397,7 +388,6 @@ def __init__(self, context): self._build_column_index() self._load_config() self._setup_providers() - self._verify_startup_providers() self.history = [] # --- HANDLER REGISTRY --- @@ -722,41 +712,16 @@ def _format_layers_table(self, df: Optional[pd.DataFrame] = None) -> str: return "\n".join(lines) def _load_config(self): - self.preferred_provider = os.environ.get("PREFERRED_PROVIDER", "openrouter") # Default to OpenRouter if API key is provided, otherwise fallback to local Ollama. This can be overridden by config file or env variable. - - # Cloud provider settings with sensible defaults. OpenRouter is the default cloud provider if API key is provided. - # Default to a fast flash-class model: the intent-planning task is - # simple JSON generation, and a 70B model added ~15-30s of latency for - # no accuracy benefit (see plan Phase B.1). Override with OPENROUTER_MODEL - # (or agent_config.yaml) to use a larger model. - self.openrouter_model = os.environ.get("OPENROUTER_MODEL", "~google/gemini-flash-latest") - self.openrouter_base_url = os.environ.get("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1") - self.openrouter_api_key = os.environ.get("OPENROUTER_API_KEY", None) - self.openrouter_request_timeout = float(os.environ.get("OPENROUTER_REQUEST_TIMEOUT", "15.0")) - # Cap the completion length. OpenRouter pre-authorizes - # `max_tokens * completion_price` against the key's remaining budget - # BEFORE generating; with no cap the client requests the model's full - # output window (tens of thousands of tokens), which 402s on a - # credit/weekly-limited key even though the actual intent-planning - # response is only a few hundred tokens. Keep this modest. - self.openrouter_max_tokens = int(os.environ.get("OPENROUTER_MAX_TOKENS", "2048")) - # Bias OpenRouter's upstream provider selection ("throughput"/"latency"/ - # "price"); empty string lets OpenRouter choose freely (see Phase B.2). - self.openrouter_provider_sort = os.environ.get("OPENROUTER_PROVIDER_SORT", "throughput") - # Ask the model for a schema-validated Intent object directly instead of - # free-form JSON + regex repair. More reliable, but only works on models - # whose OpenRouter route supports structured/JSON-schema output (e.g. - # Gemini, GPT-4o). Default OFF so an unsupported model can't break; flip - # on with OPENROUTER_STRUCTURED_OUTPUT=1 or the config key. - self.openrouter_structured_output = os.environ.get( - "OPENROUTER_STRUCTURED_OUTPUT", "" - ).strip().lower() in ("1", "true", "yes", "on") - - # Local fallback if no cloud (OpenRouter) is available or if the user prefers it. Ollama is the default local provider. - self.fallback_to_local = True # Default to allowing fallback to local Ollama if OpenRouter fails - self.ollama_host = "localhost" - self.ollama_port = "11435" - self.ollama_model = "llama3.2:3b" + # OpenCode is the only supported agent backend: a local OpenCode + # server (opencode.ai) backs every LLM call. No API key here -- the + # credential lives in OpenCode's own config, entered once via + # `opencode auth login` or the Weights Studio landing page's login + # modal. OPENCODE_URL is the SAME shared root env var the frontend + # reads (forwarded to window.WS_OPENCODE_URL by weightslab/ui/server.py's + # _UI_ENV_GLOBALS) -- set it once and both sides point at one server. + self.preferred_provider = "opencode" + self.opencode_url = os.environ.get("OPENCODE_URL", "http://127.0.0.1:4096") + self.opencode_model = os.environ.get("OPENCODE_MODEL", "") repo_root = Path(__file__).resolve().parents[4] # weightslab/ root inner_pkg = Path(__file__).resolve().parents[3] @@ -783,23 +748,9 @@ def _load_config(self): if not cfg or "agent" not in cfg: continue a_cfg = cfg["agent"] - # Agents settings - self.preferred_provider = a_cfg.get("provider", self.preferred_provider).lower() - self.fallback_to_local = a_cfg.get("fallback_to_local", self.fallback_to_local) - - # OPENROUTER - self.openrouter_model = a_cfg.get("openrouter_model", self.openrouter_model) - self.openrouter_base_url = a_cfg.get("openrouter_base_url", self.openrouter_base_url) - self.openrouter_api_key = a_cfg.get("openrouter_api_key", self.openrouter_api_key) - self.openrouter_request_timeout = float(a_cfg.get("openrouter_request_timeout", self.openrouter_request_timeout)) - self.openrouter_max_tokens = int(a_cfg.get("openrouter_max_tokens", self.openrouter_max_tokens)) - self.openrouter_provider_sort = a_cfg.get("openrouter_provider_sort", self.openrouter_provider_sort) - self.openrouter_structured_output = bool(a_cfg.get("openrouter_structured_output", self.openrouter_structured_output)) - - # OLLAMA - self.ollama_host = a_cfg.get("ollama_host", self.ollama_host) - self.ollama_port = a_cfg.get("ollama_port", self.ollama_port) - self.ollama_model = a_cfg.get("ollama_model", self.ollama_model) + # OPENCODE + self.opencode_url = a_cfg.get("opencode_url", self.opencode_url) + self.opencode_model = a_cfg.get("opencode_model", self.opencode_model) _LOGGER.info(f"Applied agent configuration from {path}") _LOGGER.debug(f"Agent Config: {cfg}") @@ -813,228 +764,67 @@ def _load_config(self): "\n# #######################################" + "\n" + "# #######################################" + "\n" + f"Agent initialized from configuration {path}: " + "\n" + - f"\tFinal Agent Configuration: Preferred Provider={self.preferred_provider}, " + "\n" + - f"\tFallback to Local={self.fallback_to_local}, " + "\n" + - f"\tOpenRouter Model={self.openrouter_model} with:" + "\n" + - f"\t\tAPI Key={f'{self.openrouter_api_key[:4]}****{self.openrouter_api_key[-4:]}' if self.openrouter_api_key else 'None'}" + "\n" + - f"\t\tBase URL={self.openrouter_base_url}, " + "\n" + - f"\tOllama Model={self.ollama_model}" + "\n" + + f"\tOpenCode URL={self.opencode_url}, Model={self.opencode_model or '(server default)'}" + "\n" + "# #######################################" + "\n" + "# #######################################" + "\n" + "" ) - @staticmethod - def _effective_http_port(parsed_url, explicit_port: Optional[str]) -> int: - if explicit_port and explicit_port.isdigit(): - return int(explicit_port) - if parsed_url.port is not None: - return parsed_url.port - return 443 if parsed_url.scheme == "https" else 80 - - @staticmethod - def _normalize_openrouter_base_url(raw_url: str, explicit_port: Optional[str]) -> str: - url = (raw_url or "https://openrouter.ai/api/v1").strip() - parsed = urlparse(url) - if not parsed.scheme: - parsed = urlparse(f"https://{url}") - - host = parsed.hostname or "openrouter.ai" - port = DataManipulationAgent._effective_http_port(parsed, explicit_port) - netloc = host if ((parsed.scheme == "https" and port == 443) or (parsed.scheme == "http" and port == 80)) else f"{host}:{port}" - path = parsed.path if parsed.path else "/api/v1" - - return urlunparse((parsed.scheme, netloc, path, "", "", "")) - def _setup_providers(self): - self.chain_ollama = None - self.chain_openrouter = None + self.chain_opencode = None initialized = False - # Determine which providers to initialize - active_providers = {self.preferred_provider} - if self.fallback_to_local: - active_providers.add("ollama") - - # OPEN ROUTER - if "openrouter" in active_providers and self.openrouter_api_key: - _LOGGER.info(f"Setting up OpenRouter with model {self.openrouter_model}") - try: - if ChatOpenAI is None: - _LOGGER.warning("langchain_openai is not installed, skipping OpenRouter provider") - else: - explicit_openrouter_port = os.environ.get("OPENROUTER_PORT", "").strip() - openrouter_base_url = self._normalize_openrouter_base_url(self.openrouter_base_url, explicit_openrouter_port) - parsed = urlparse(openrouter_base_url) - effective_port = self._effective_http_port(parsed, explicit_openrouter_port) - - # Bias OpenRouter's upstream routing so it doesn't pick a - # slow/cold provider for the model — a major source of the - # query tail latency (see plan Phase B.2). Configurable via - # `openrouter_provider_sort` ("throughput"/"latency"/"price"); - # set to "" to let OpenRouter choose freely. - extra_body = None - sort = getattr(self, "openrouter_provider_sort", "throughput") - if sort: - extra_body = {"provider": {"sort": sort}} - - llm = ChatOpenAI( - model=self.openrouter_model, temperature=0, - api_key=self.openrouter_api_key, - base_url=openrouter_base_url, - streaming=False, max_retries=1, request_timeout=self.openrouter_request_timeout, - max_tokens=self.openrouter_max_tokens, - extra_body=extra_body, - ) - self.chain_openrouter = llm - initialized = True - _LOGGER.info( - f"[Agent] OpenRouter enabled: {self.openrouter_model} via {parsed.hostname}:{effective_port} " - f"(provider sort={sort or 'default'})" - ) - except Exception as e: _LOGGER.error(f"OpenRouter error: {e}") - - # LOCAL - if "ollama" in active_providers: - try: - if ChatOllama is None: - _LOGGER.warning("langchain_ollama is not installed, skipping Ollama provider") - else: - _LOGGER.info(f"Setting up Ollama with model {self.ollama_model}") - host = self.ollama_host.split(':')[0] - port = self.ollama_port - llm = ChatOllama(base_url=f"http://{host}:{port}", model=self.ollama_model, temperature=0, timeout=15) - self.chain_ollama = llm - initialized = True - _LOGGER.info(f"[Agent] Ollama enabled: {self.ollama_model}") - except Exception as e: _LOGGER.error(f"Ollama error: {e}") - - return initialized - - def _verify_startup_providers(self) -> None: - """ - A provider configured at construction time (agent_config.yaml / env - vars, e.g. OPENROUTER_API_KEY) never goes through the connectivity - check the `/init` UI flow runs (see `initialize_with_cloud_key`) -- - `_setup_providers()` only builds a client object from whatever key - string it was given, it never confirms the key is actually accepted. - `is_available()` treats "a client object exists" as "ready", so a - bad startup key would otherwise report available=True indefinitely - (CheckAgentHealth says "ready to help you") until the first real - query 401s. Probe once here instead, so that mismatch can't happen. - - Only runs for the OpenRouter chain built in `__init__`: Ollama's - `is_available()` already does a live reachability check every call, - and `initialize_with_cloud_key`/`change_model` already run their own - explicit post-`_setup_providers()` check. - """ - if self.chain_openrouter is None: - return - ok, message = self._check_chat_provider("openrouter") - if not ok: - _LOGGER.warning(f"[Agent] Startup OpenRouter connectivity check failed, disabling it: {message}") - self.chain_openrouter = None - - def _check_chat_provider(self, provider: str) -> "tuple[bool, str]": - """Run a minimal chat request to verify the provider is actually usable.""" - chain = getattr(self, f"chain_{provider}", None) - if chain is None: - return False, f"{provider} client was not initialized." - try: - # Keep the probe's reservation tiny. OpenRouter pre-authorizes - # `max_tokens * price` before generating, so a large (or unset) - # cap makes the health check 402 on a budget-limited key even - # though the model is perfectly usable for real (short) requests. - probe = chain - if provider == "openrouter" and hasattr(chain, "bind"): - try: - probe = chain.bind(max_tokens=16) - except Exception: - probe = chain - response = probe.invoke("Reply with OK.") + self.chain_opencode = OpenCodeChat(self.opencode_url, self.opencode_model).as_runnable() + initialized = True + _LOGGER.info( + f"[Agent] OpenCode enabled: {self.opencode_url} " + f"(model={self.opencode_model or 'server default'})" + ) except Exception as e: - _LOGGER.warning(f"[{provider}] connectivity check failed: {e}") - return False, f"{provider} connectivity check failed: {e}" - - text = response.content if hasattr(response, "content") else str(response) - if not str(text).strip(): - return False, f"{provider} connectivity check returned an empty response." - - return True, f"{provider} connectivity check succeeded." + _LOGGER.error(f"OpenCode error: {e}") - def is_ollama_available(self) -> bool: - return self.chain_ollama is not None + return initialized def is_available(self) -> bool: - """ - Return True if any LLM provider is actually ready to serve requests. - - OpenRouter is considered ready as soon as its chain is set up. - - Ollama requires an active HTTP connection check because the ChatOllama - constructor succeeds even when the daemon is not running. - """ - if self.chain_openrouter is not None: - return True - if self.chain_ollama is not None: - return self._is_ollama_reachable() - return False - - def _is_ollama_reachable(self) -> bool: - """Ping the Ollama HTTP endpoint to verify the daemon is actually running.""" - try: - import urllib.request as _ur - host = self.ollama_host.split(':')[0] - url = f"http://{host}:{self.ollama_port}/api/version" - with _ur.urlopen(_ur.Request(url), timeout=2) as resp: - return resp.status == 200 - except Exception: - return False + """Return True if the OpenCode provider is ready to serve requests.""" + return self.chain_opencode is not None def initialize_with_cloud_key(self, api_key: str, provider: str, model: Optional[str] = None) -> "tuple[bool, str]": """ - Initialize (or reinitialize) the OpenRouter cloud provider. + Initialize (or reinitialize) the OpenCode provider. Args: - api_key: The API key obtained from the provider's website. - provider: Must be ``"openrouter"``. - model: OpenRouter model identifier chosen by the user. + api_key: Unused -- OpenCode's credential lives in the OpenCode + server's own config (``opencode auth login`` or the Weights + Studio landing page's login modal), never in this agent. + provider: Must be ``"opencode"`` -- the only supported backend. + model: OpenCode's ``"providerID/modelID"`` model identifier. Returns: ``(True, success_message)`` or ``(False, error_message)``. """ - if not api_key or not api_key.strip(): - return False, "API key cannot be empty." - - if provider.lower() != "openrouter": - return False, "Only OpenRouter cloud onboarding is supported." + provider = (provider or "").lower() + if provider != "opencode": + return False, "Only OpenCode is supported as the agent backend." if model is not None and not model.strip(): return False, "Model cannot be empty." - self.openrouter_api_key = api_key.strip() - self.openrouter_base_url = "https://openrouter.ai/api/v1" - self.openrouter_model = model.strip() if model and model.strip() else self.openrouter_model - self.preferred_provider = "openrouter" + self.opencode_model = model.strip() if model and model.strip() else self.opencode_model + self.preferred_provider = "opencode" success = self._setup_providers() + if self.chain_opencode is None or not success: + return False, "Could not reach the OpenCode server. Please verify OPENCODE_URL and that it is running." - if self.chain_openrouter is None or not success: - return False, "Provider client initialization failed. Please verify your API key, base URL, and model." - - chat_ok, chat_message = self._check_chat_provider("openrouter") - if not chat_ok: - self.chain_openrouter = None - return False, chat_message - - return True, "Agent initialized successfully. Ready to help you." + return True, "Agent initialized successfully via OpenCode. Ready to help you." def change_model(self, model: str) -> "tuple[bool, str]": """ - Switch the active OpenRouter model without re-entering the API key. + Switch the active OpenCode model, without re-entering credentials. Args: - model: OpenRouter model identifier (e.g. ``"openai/gpt-4o"``). + model: OpenCode's ``"providerID/modelID"`` model identifier. Returns: ``(True, success_message)`` or ``(False, error_message)``. @@ -1042,25 +832,17 @@ def change_model(self, model: str) -> "tuple[bool, str]": if not model or not model.strip(): return False, "Model cannot be empty." - if not getattr(self, "openrouter_api_key", None): - return False, "No API key configured. Please initialize the agent first (/init)." - - self.openrouter_model = model.strip() + self.opencode_model = model.strip() success = self._setup_providers() - - if self.chain_openrouter is None or not success: - return False, "Provider reinitialization failed. Please verify the model name." - - chat_ok, chat_message = self._check_chat_provider("openrouter") - if not chat_ok: - self.chain_openrouter = None - return False, chat_message - - return True, f"Model switched to {self.openrouter_model}. Ready to help you." + if self.chain_opencode is None or not success: + return False, "Could not reach the OpenCode server. Please verify OPENCODE_URL and that it is running." + return True, f"Model switched to {self.opencode_model}. Ready to help you." def get_available_models(self) -> "tuple[bool, list[str], str]": """ - Fetch the list of models available via the configured OpenRouter API key. + Fetch every provider/model OpenCode is authenticated for, flattened + into "providerID/modelID" strings (same convention used throughout + the frontend). Returns: ``(True, model_ids, "")`` on success, or ``(False, [], error_message)``. @@ -1068,31 +850,77 @@ def get_available_models(self) -> "tuple[bool, list[str], str]": import urllib.request as _ur import json as _json - api_key = getattr(self, "openrouter_api_key", None) - if not api_key: - return False, [], "No API key configured. Please initialize the agent first (/init)." - try: - url = "https://openrouter.ai/api/v1/models" - req = _ur.Request(url, headers={"Authorization": f"Bearer {api_key}"}) - with _ur.urlopen(req, timeout=10) as resp: + url = f"{self.opencode_url.rstrip('/')}/config/providers" + with _ur.urlopen(_ur.Request(url), timeout=10) as resp: data = _json.loads(resp.read().decode()) - models = sorted(entry["id"] for entry in data.get("data", []) if "id" in entry) - return True, models, "" + models = [] + for provider in data.get("providers", []): + provider_id = provider.get("id") + if not provider_id: + continue + provider_models = provider.get("models", {}) + model_ids = provider_models.keys() if isinstance(provider_models, dict) else \ + (m.get("id") for m in provider_models if isinstance(m, dict)) + for model_id in model_ids: + if model_id: + models.append(f"{provider_id}/{model_id}") + return True, sorted(models), "" except Exception as exc: - _LOGGER.warning("get_available_models error: %s", exc) - return False, [], f"Could not fetch models: {exc}" + _LOGGER.warning("get_available_models (opencode) error: %s", exc) + return False, [], f"Could not reach the OpenCode server: {exc}" def reset_connection(self) -> "tuple[bool, str]": - """Clear the active cloud connection and revert the agent to the uninitialized state.""" - self.chain_ollama = None - self.chain_openrouter = None - self.openrouter_api_key = None - self.openrouter_model = os.environ.get("OPENROUTER_MODEL", "~google/gemini-flash-latest") - self.preferred_provider = "openrouter" + """Clear the active OpenCode connection and revert the agent to the + uninitialized state.""" + self.chain_opencode = None + self.opencode_model = os.environ.get("OPENCODE_MODEL", "") + self.preferred_provider = "opencode" return True, "Agent connection reset. Type /init to set up again." + def clear_history(self) -> "tuple[bool, str]": + """Wipe the conversation history (self.history) without touching the + provider connection -- distinct from reset_connection, which drops the + connection but leaves history untouched. Only affects this SDK agent's + own history; OpenCode-backed sessions (the landing page, /loop jobs) + manage their own, separate context and are untouched by this.""" + count = len(self.history) + self.history = [] + return True, f"Cleared {count} history entries." + + def compact_history(self) -> "tuple[bool, str]": + """Summarize self.history via OpenCode, replacing it with a single + entry -- a real compaction (LLM summary), not a harder truncation of + the existing self.history[-5:] read-time cap.""" + if not self.history: + return True, "History is already empty; nothing to compact." + + if not self.chain_opencode: + return False, "Could not compact history -- no provider available." + + transcript = "\n".join(self.history) + summarize_prompt = ChatPromptTemplate.from_messages([ + ("system", "Summarize the following agent conversation log into a short, " + "dense paragraph that preserves the concrete actions taken and " + "any user preferences stated. Output only the summary text, no " + "preamble."), + ("human", "{transcript}"), + ]) + + try: + response = (summarize_prompt | self.chain_opencode).invoke({"transcript": transcript}) + summary = response.content if hasattr(response, "content") else str(response) + summary = summary.strip() + if summary: + before = len(self.history) + self.history = [f"Summary of {before} prior entries: {summary}"] + return True, f"Compacted {before} entries into one summary." + except Exception as e: + _LOGGER.error(f"OpenCode failed to compact history: {e}") + + return False, "Could not compact history -- no provider available." + def _build_column_index(self): """Builds normalized token indexes and lightweight synonyms for column resolution.""" self._cols = list(self.df_schema['columns']) @@ -1702,24 +1530,12 @@ def _query_langchain(self, name: str, chain, instruction: str, system_prompt: st prompt = ChatPromptTemplate.from_messages([("system", escaped_sys), ("human", "{instruction}")]) - # Optionally bind schema-validated structured output (Intent) so the - # model returns a parsed object directly, skipping the free-form JSON - # + regex-repair path in _parse_intent_from_response. Gated because - # not every OpenRouter route supports it (see _load_config). - runnable_chain = chain - if name == "openrouter" and getattr(self, "openrouter_structured_output", False): - try: - runnable_chain = chain.with_structured_output(Intent) - except Exception as e: - _LOGGER.warning(f"[{name}] structured output unavailable, falling back to free-form JSON: {e}") - runnable_chain = chain - # Time the pure LLM round-trip (see plan Phase A). This is the # dominant cost of a query; ~chars/4 is a rough token estimate. prompt_chars = len(escaped_sys) + len(instruction) - model_name = getattr(self, f"{name}_model", None) or getattr(self, "openrouter_model", name) + model_name = getattr(self, f"{name}_model", None) or name _t0 = time.perf_counter() - response = (prompt | runnable_chain).invoke({"instruction": instruction}) + response = (prompt | chain).invoke({"instruction": instruction}) _llm_elapsed = time.perf_counter() - _t0 _LOGGER.info( @@ -1774,7 +1590,7 @@ def _is_auth_error(error) -> bool: ) def _try_query_provider(self, provider: str, instruction: str, system_prompt: str) -> Optional[List[dict]]: - # 1. Dynamically find the chain (chain_openrouter, chain_ollama) + # 1. Dynamically find the chain (chain_opencode) chain = getattr(self, f"chain_{provider}", None) # 2. If it exists, use the standard LangChain method @@ -1852,8 +1668,6 @@ def query(self, instruction: str, abort_event: Optional[threading.Event] = None, ) order = [self.preferred_provider] - if self.fallback_to_local and self.preferred_provider != "ollama": - order.append("ollama") for provider in order: if abort_event and abort_event.is_set(): return [] @@ -1872,19 +1686,18 @@ def query(self, instruction: str, abort_event: Optional[threading.Event] = None, # If we get here, all providers failed error_msg = "Internal Agent Error: Failed to generate a plan." if self._is_auth_error(self._last_query_error): - # The provider rejected our credentials — the agent isn't - # connected. Invalidate the cached client so is_available()/ - # CheckAgentHealth immediately stop reporting "available" for a - # connection that just proved broken, instead of leaving the - # health check permanently stale until the process restarts. - self.chain_openrouter = None + # OpenCode rejected the request -- the agent isn't connected. + # Invalidate the cached client so is_available()/CheckAgentHealth + # immediately stop reporting "available" for a connection that + # just proved broken, instead of leaving the health check + # permanently stale until the process restarts. + self.chain_opencode = None error_msg = ( - "Agent not connected: the LLM provider rejected the request " - "(401 Unauthorized). Check your API key and re-initialize the " - "agent with /init." + "Agent not connected: OpenCode rejected the request " + "(401 Unauthorized). Re-initialize the agent with /init." ) - elif not self.is_ollama_available() and not os.environ.get("OPENROUTER_API_KEY"): - error_msg = "No LLM providers configured. Please check your API keys or local Ollama setup. Initialize the agent with /init." + elif not self.is_available(): + error_msg = "No LLM provider configured. Please check your OpenCode server (OPENCODE_URL) and initialize the agent with /init." _LOGGER.info(f"[Agent] Query total wall time: {time.perf_counter() - _query_t0:.2f}s (no provider succeeded)") return [{"function": "out_of_scope", "params": {"reason": error_msg}}] @@ -1957,14 +1770,14 @@ def generate_code(self, prompt: str, context_code: str = ""): """Propose Python for a notebook cell from a natural-language ``prompt``. This does NOT execute anything and does NOT touch the intent pipeline; it - only asks the active LLM provider for runnable code. Returns a - ``(code, explanation)`` tuple. Raises RuntimeError when no provider is - configured so the service can report a clean error. + only asks OpenCode for runnable code. Returns a ``(code, explanation)`` + tuple. Raises RuntimeError when no provider is configured so the + service can report a clean error. """ if not self.is_available(): raise RuntimeError( - "Agent not configured. Initialize a provider with /init (or run a " - "local Ollama server) before using \">\" notebook cells." + "Agent not configured. Initialize OpenCode with /init before " + "using \">\" notebook cells." ) system_prompt = NOTEBOOK_CODE_PROMPT.format( @@ -1978,33 +1791,17 @@ def generate_code(self, prompt: str, context_code: str = ""): [("system", escaped_sys), ("human", "{instruction}")] ) - order = [self.preferred_provider] - if self.fallback_to_local and self.preferred_provider != "ollama": - order.append("ollama") - - last_error = None - for provider in order: - chain = getattr(self, f"chain_{provider}", None) - if chain is None: - continue - try: - response = (chat_prompt | chain).invoke({"instruction": prompt}) - text = getattr(response, "content", None) - if text is None: - text = str(response) - code, explanation = self._extract_code_and_explanation(text) - self.history.append(f"User (notebook): {prompt}") - self.history.append("Action: proposed notebook code") - return code, explanation - except Exception as exc: - last_error = exc - _LOGGER.warning("[notebook code-gen] provider %s failed: %s", provider, exc) - continue - - raise RuntimeError( - f"Code generation failed: {last_error}" if last_error - else "Code generation failed: no provider produced a response." - ) + try: + response = (chat_prompt | self.chain_opencode).invoke({"instruction": prompt}) + text = getattr(response, "content", None) + if text is None: + text = str(response) + code, explanation = self._extract_code_and_explanation(text) + self.history.append(f"User (notebook): {prompt}") + self.history.append("Action: proposed notebook code") + return code, explanation + except Exception as exc: + raise RuntimeError(f"Code generation failed: {exc}") # ------------------------------------------------------------------ # Experiment report narrative @@ -2017,15 +1814,15 @@ def generate_report_narrative(self, stats_summary: str) -> str: sees ``stats_summary`` (signal shapes + dataframe stats), never raw history, so it can comment on the numbers but not invent new ones. - Mirrors ``generate_code``'s provider-fallback shape, but returns plain - text (no code-fence extraction) and does not touch ``self.history`` -- - this is a side-channel report artifact, not a step in the - NL-to-data-operation conversation. + Mirrors ``generate_code``, but returns plain text (no code-fence + extraction) and does not touch ``self.history`` -- this is a + side-channel report artifact, not a step in the NL-to-data-operation + conversation. """ if not self.is_available(): raise RuntimeError( - "Agent not configured. Initialize a provider with /init (or run a " - "local Ollama server) before generating a report." + "Agent not configured. Initialize OpenCode with /init before " + "generating a report." ) system_prompt = REPORT_ANALYSIS_PROMPT.format( @@ -2036,25 +1833,9 @@ def generate_report_narrative(self, stats_summary: str) -> str: [("system", escaped_sys), ("human", "Write the analysis section.")] ) - order = [self.preferred_provider] - if self.fallback_to_local and self.preferred_provider != "ollama": - order.append("ollama") - - last_error = None - for provider in order: - chain = getattr(self, f"chain_{provider}", None) - if chain is None: - continue - try: - response = (chat_prompt | chain).invoke({}) - text = getattr(response, "content", None) - return text.strip() if text else str(response).strip() - except Exception as exc: - last_error = exc - _LOGGER.warning("[report narrative] provider %s failed: %s", provider, exc) - continue - - raise RuntimeError( - f"Report narrative generation failed: {last_error}" if last_error - else "Report narrative generation failed: no provider produced a response." - ) + try: + response = (chat_prompt | self.chain_opencode).invoke({}) + text = getattr(response, "content", None) + return text.strip() if text else str(response).strip() + except Exception as exc: + raise RuntimeError(f"Report narrative generation failed: {exc}") diff --git a/weightslab/trainer/services/agent/opencode_chat.py b/weightslab/trainer/services/agent/opencode_chat.py new file mode 100644 index 00000000..80904a29 --- /dev/null +++ b/weightslab/trainer/services/agent/opencode_chat.py @@ -0,0 +1,229 @@ +"""LangChain-`Runnable`-compatible wrapper around a local OpenCode server. + +OpenCode (https://opencode.ai) is a session-based, tool-using coding-agent +server -- not a plain chat-completions endpoint like OpenRouter/Ollama. This +module lets `DataManipulationAgent` (agent.py) use one as a third provider +without changing any of its three existing call sites (`query()`, +`generate_code()`, `generate_report_narrative()`), all of which do +`(prompt | chain).invoke(...)` and read `.content` off the result. + +Mirrors the wire protocol already implemented in +weights_studio/src/landing/agent/opencodeClient.ts: create a session, send a +message, and collect the assistant's text from the global SSE event stream +ending on `session.idle` -- a bare POST .../message response is not a +reliable completion/content signal on its own (weights_studio's TS client +treats it the same way, with its own idle-driven finish as the source of +truth). The stream must be opened BEFORE the message is sent, or events +emitted in between are lost (same "stream-first" rule the TS client follows). + +A fresh session is created per call rather than reused across calls: +`self.history` on `DataManipulationAgent` already carries cross-call context +via `INTENT_PROMPT`'s own `history=` placeholder, so OpenCode's own +multi-turn session memory isn't needed here, and a fresh session per call +avoids an ever-growing OpenCode-side history for a long-running experiment. + +Every mutating tool (write/edit/patch/bash) is explicitly disabled on the +outgoing message: this call wants a text/JSON reply for the SDK agent to +parse and act on itself, not file writes as a side effect. That is the +opposite default from the Weights Studio landing-page chat, which +deliberately runs OpenCode's full toolset. +""" + +from __future__ import annotations + +import json +import logging +import threading +import time +import urllib.error +import urllib.request +from typing import Optional + +_LOGGER = logging.getLogger(__name__) + +# Mutating tools, disabled on every message this wrapper sends -- see module +# docstring. Named explicitly (not an allowlist) so a newly-added read-only +# tool on the OpenCode side keeps working without a change here. +_MUTATING_TOOLS = ("write", "edit", "patch", "bash") + + +class OpenCodeError(RuntimeError): + """Raised when a call to the OpenCode server fails outright (not just an + empty/partial reply -- those degrade to whatever text was collected).""" + + +class OpenCodeChat: + """One instance per configured `(base_url, model)` pair; safe to reuse + across calls (`_call` is the only state-touching method, and it is + self-contained per invocation).""" + + def __init__(self, base_url: str, model: Optional[str] = None, timeout: float = 60.0): + self.base_url = (base_url or "http://127.0.0.1:4096").rstrip("/") + self.model = model + self.timeout = timeout + + # -- wire helpers --------------------------------------------------- # + + def _model_ref(self) -> Optional[dict]: + """OpenCode identifies a model as {providerID, modelID}; our config + carries it as one "providerID/modelID" string (matching the exact + convention used throughout the frontend -- see opencodeClient.ts's + formatModelValue/parseModelValue). Split on the FIRST slash only: + model IDs themselves often contain slashes (e.g. OpenRouter's + "anthropic/claude-opus-4.6"), so a naive split would truncate it.""" + if not self.model or "/" not in self.model: + return None + provider_id, model_id = self.model.split("/", 1) + return {"providerID": provider_id, "modelID": model_id} + + def _request(self, path: str, method: str = "GET", body: Optional[dict] = None, headers: Optional[dict] = None): + data = json.dumps(body).encode("utf-8") if body is not None else None + all_headers = {"Content-Type": "application/json"} if body is not None else {} + all_headers.update(headers or {}) + req = urllib.request.Request(f"{self.base_url}{path}", data=data, headers=all_headers, method=method) + return urllib.request.urlopen(req, timeout=self.timeout) + + def _create_session(self) -> str: + try: + with self._request("/session", method="POST", body={"title": "weightslab-sdk-agent"}) as resp: + data = json.loads(resp.read().decode("utf-8")) + except (urllib.error.URLError, ValueError) as exc: + raise OpenCodeError(f"Could not create an OpenCode session at {self.base_url}: {exc}") from exc + session_id = data.get("id") + if not session_id: + raise OpenCodeError(f"OpenCode did not return a session id: {data!r}") + return session_id + + def _send_message(self, session_id: str, text: str) -> None: + body = { + "parts": [{"type": "text", "text": text}], + "tools": {name: False for name in _MUTATING_TOOLS}, + } + model_ref = self._model_ref() + if model_ref: + body["model"] = model_ref + try: + self._request(f"/session/{session_id}/message", method="POST", body=body).read() + except urllib.error.URLError as exc: + raise OpenCodeError(f"OpenCode rejected the prompt: {exc}") from exc + + @staticmethod + def _handle_event(payload: str, session_id: str, assistant_message_ids: set, text_parts: dict) -> Optional[str]: + """Returns None while the turn is still in progress, or a string + ("idle"/"error") once this session's turn has finished. Field names + are read defensively -- OpenCode's event shapes are not part of its + published docs, so a minor server change should degrade to "no text + collected" rather than raising inside the stream loop.""" + try: + event = json.loads(payload) + except (ValueError, TypeError): + return None + event_type = event.get("type") + props = event.get("properties") or {} + + if event_type == "message.updated": + msg = props.get("info") or props.get("message") or props + msg_session = str(msg.get("sessionID") or "") + if msg_session and msg_session != session_id: + return None + if str(msg.get("role") or "") == "assistant" and msg.get("id"): + assistant_message_ids.add(str(msg["id"])) + return None + + if event_type == "message.part.updated": + part = props.get("part") or props + part_session = str(part.get("sessionID") or "") + if part_session and part_session != session_id: + return None + message_id = str(part.get("messageID") or "") + if message_id not in assistant_message_ids: + return None + if part.get("type") == "text" and isinstance(part.get("text"), str): + part_id = str(part.get("id") or message_id) + text_parts[part_id] = part["text"] + return None + + if event_type == "session.idle" and str(props.get("sessionID") or "") == session_id: + return "idle" + + if event_type == "session.error" and str(props.get("sessionID") or "") == session_id: + return "error" + + return None + + def _collect_reply(self, session_id: str, text: str) -> str: + """Open the SSE stream, THEN send the message on a background thread, + THEN read the stream until this session goes idle/errors. Ordering + matters: opening the stream first means it starts buffering before the + message is sent, so nothing emitted in the gap between "create + session" and "start reading" is lost (see module docstring).""" + text_parts: dict = {} + assistant_message_ids: set = set() + + try: + stream = self._request("/event", headers={"Accept": "text/event-stream"}) + except urllib.error.URLError as exc: + raise OpenCodeError(f"Could not open the OpenCode event stream at {self.base_url}: {exc}") from exc + + send_errors: list = [] + + def _send() -> None: + try: + self._send_message(session_id, text) + except Exception as exc: # noqa: BLE001 - surfaced via send_errors + send_errors.append(exc) + + sender = threading.Thread(target=_send, daemon=True) + + try: + sender.start() + data_lines: list = [] + deadline = time.monotonic() + self.timeout + for raw_line in stream: + if time.monotonic() > deadline: + _LOGGER.warning("OpenCode event stream timed out waiting for session %s", session_id) + break + line = raw_line.decode("utf-8", errors="replace").rstrip("\r\n") + if line == "": + if data_lines: + payload = "\n".join(data_lines) + data_lines = [] + outcome = self._handle_event(payload, session_id, assistant_message_ids, text_parts) + if outcome is not None: + break + continue + if line.startswith(":"): + continue + if line.startswith("data:"): + data_lines.append(line[5:].lstrip(" ")) + except (urllib.error.URLError, TimeoutError, OSError) as exc: + # A dropped stream degrades to whatever text was collected so far, + # matching the browser client's "stream dropped -> don't leave the + # turn hanging" behavior, rather than losing a partial reply. + _LOGGER.warning("OpenCode event stream dropped for session %s: %s", session_id, exc) + finally: + stream.close() + + sender.join(timeout=1.0) + if send_errors: + raise send_errors[0] + + return "".join(text_parts[key] for key in text_parts) + + # -- public surface --------------------------------------------------- # + + def _call(self, prompt_value): + from langchain_core.messages import AIMessage + + text = prompt_value.to_string() if hasattr(prompt_value, "to_string") else str(prompt_value) + session_id = self._create_session() + reply = self._collect_reply(session_id, text) + return AIMessage(content=reply) + + def as_runnable(self): + """A `Runnable` usable exactly like `ChatOpenAI`/`ChatOllama` in + `(prompt | chain).invoke(...)` — the one integration point + `_query_langchain`/`generate_code`/`generate_report_narrative` need.""" + from langchain_core.runnables import RunnableLambda + + return RunnableLambda(self._call) diff --git a/weightslab/trainer/services/agent_service.py b/weightslab/trainer/services/agent_service.py index e39998b0..6d9ab191 100644 --- a/weightslab/trainer/services/agent_service.py +++ b/weightslab/trainer/services/agent_service.py @@ -77,10 +77,15 @@ def CheckAgentHealth(self, request, context): @safe_grpc(lambda msg: pb2.InitializeAgentResponse(success=False, message=msg)) def InitializeAgent(self, request, context): """ - Initialize the OpenRouter cloud provider with a user-supplied API key. + Initialize the OpenCode agent backend. - Supported provider (AgentProviderType enum): - PROVIDER_OPENROUTER (0) — openrouter.ai + Supported providers (AgentProviderType enum): + PROVIDER_OPENCODE (1) — a local OpenCode server; api_key is + ignored, the credential lives in OpenCode's own config. + PROVIDER_OPENROUTER (0) is a legacy enum value kept for wire + compatibility with older frontends; requesting it is rejected + below rather than removed from the .proto, so an old client + sending it gets a clean error instead of a decode failure. Returns: InitializeAgentResponse { success: bool, message: str } @@ -93,10 +98,10 @@ def InitializeAgent(self, request, context): ) provider_name = AGENT_PROVIDER_MAP.get(request.provider) - if provider_name != "openrouter": + if provider_name != "opencode": return pb2.InitializeAgentResponse( success=False, - message="Only OpenRouter cloud onboarding is supported.", + message="Only OpenCode is supported as the agent backend.", ) logger.debug( @@ -112,7 +117,7 @@ def InitializeAgent(self, request, context): @safe_grpc(lambda msg: pb2.ChangeAgentModelResponse(success=False, message=msg)) def ChangeAgentModel(self, request, context): """ - Switch the active OpenRouter model without re-entering the API key. + Switch the active OpenCode model without re-entering credentials. Returns: ChangeAgentModelResponse { success: bool, message: str } @@ -131,7 +136,7 @@ def ChangeAgentModel(self, request, context): @safe_grpc(lambda msg: pb2.GetAgentModelsResponse(success=False, models=[], message=msg)) def GetAgentModels(self, request, context): """ - Return the list of models available via the stored OpenRouter API key. + Return every provider/model the OpenCode server is authenticated for. Returns: GetAgentModelsResponse { success: bool, models: [str], message: str } @@ -160,3 +165,33 @@ def ResetAgent(self, request, context): logger.debug("ResetAgent") success, message = agent.reset_connection() return pb2.ResetAgentResponse(success=success, message=message) + + @safe_grpc(lambda msg: pb2.ClearAgentHistoryResponse(success=False, message=msg)) + def ClearAgentHistory(self, request, context): + """Wipe the agent's conversation history without touching the provider + connection -- distinct from ResetAgent.""" + agent = self._agent + if agent is None: + return pb2.ClearAgentHistoryResponse( + success=False, + message="Agent backend is not running.", + ) + + logger.debug("ClearAgentHistory") + success, message = agent.clear_history() + return pb2.ClearAgentHistoryResponse(success=success, message=message) + + @safe_grpc(lambda msg: pb2.CompactAgentHistoryResponse(success=False, message=msg)) + def CompactAgentHistory(self, request, context): + """Summarize the agent's conversation history via the active provider, + replacing it with the summary.""" + agent = self._agent + if agent is None: + return pb2.CompactAgentHistoryResponse( + success=False, + message="Agent backend is not running.", + ) + + logger.debug("CompactAgentHistory") + success, message = agent.compact_history() + return pb2.CompactAgentHistoryResponse(success=success, message=message) diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py index 64192cfd..f233e3ac 100755 --- a/weightslab/trainer/services/data_service.py +++ b/weightslab/trainer/services/data_service.py @@ -904,7 +904,7 @@ def _log_audit( def _is_agent_available(self) -> bool: """ - Check if the agent (Ollama) is available for natural language queries. + Check if the agent (OpenCode) is available for natural language queries. Returns: bool: True if agent is available, False otherwise diff --git a/weightslab/trainer/services/utils/tools.py b/weightslab/trainer/services/utils/tools.py index 11148009..f46f907b 100644 --- a/weightslab/trainer/services/utils/tools.py +++ b/weightslab/trainer/services/utils/tools.py @@ -22,9 +22,14 @@ # Maps the AgentProviderType proto enum integer values to the internal # provider names understood by DataManipulationAgent.initialize_with_cloud_key. -# Cloud onboarding is currently limited to OpenRouter. +# OpenCode (1) is the only agent backend weightslab actually supports; +# PROVIDER_OPENROUTER (0) is kept here (rather than removed) purely for wire +# compatibility with older frontends, so a client that still sends it gets a +# clean "not supported" rejection from AgentService.InitializeAgent instead of +# an unrecognized-enum-value decode failure. AGENT_PROVIDER_MAP: dict[int, str] = { 0: "openrouter", + 1: "opencode", } diff --git a/weightslab/ui/server.py b/weightslab/ui/server.py index 8a7fcd4f..4a5d4bc4 100644 --- a/weightslab/ui/server.py +++ b/weightslab/ui/server.py @@ -36,6 +36,7 @@ import subprocess import sys import threading +import time import webbrowser from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -342,6 +343,556 @@ def _kill_process_tree(process: "subprocess.Popen") -> None: atexit.register(_jupyter_session.shutdown) +# --------------------------------------------------------------------------- # +# OpenCode agent server (backs the landing-page agent chat) +# --------------------------------------------------------------------------- # + +# Generous: a cold `npx` run downloads the package before the server binds. +_OPENCODE_START_TIMEOUT = 45.0 + + +def _free_port() -> int: + """Reserve an unused loopback port by binding and releasing it. + + We pick the port ourselves rather than parsing it out of the child's stdout: + that gives us a URL to health-poll immediately, and avoids depending on the + exact wording of a startup log line we do not control. + """ + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _resolve_opencode_argv() -> Optional[list]: + """Locate a way to run OpenCode, preferring an already-installed binary. + + Falls back to ``npx --yes``, which fetches the package into the npx cache on + first use. That is deliberately *not* ``npm install -g``: a global install may + need elevated permissions and mutates the user's toolchain behind their back, + while the npx path needs neither and is equally automatic. + """ + exe = shutil.which("opencode") + if exe: + return [exe] + npx = shutil.which("npx") + if npx: + return [npx, "--yes", "opencode-ai@latest"] + return None + + +def _opencode_healthy(base_url: str, timeout: float = 1.5) -> bool: + import urllib.error + import urllib.request + try: + with urllib.request.urlopen(base_url + "/global/health", timeout=timeout) as resp: + return 200 <= int(resp.status) < 300 + except Exception: + return False + + +def _cors_origin_variants(origin: Optional[str]) -> list: + """Expand one origin into the set a browser might actually send. + + ``localhost`` and ``127.0.0.1`` are *different* origins to the CORS check, and + getting that wrong produces the single most confusing failure in this whole + feature: every request is blocked, and from the page it is indistinguishable + from the server being down. So allow both spellings of whichever we were given. + """ + origins: list = [] + + def add(value: str) -> None: + if value and value not in origins: + origins.append(value) + + if origin: + add(origin) + if "localhost" in origin: + add(origin.replace("localhost", "127.0.0.1")) + elif "127.0.0.1" in origin: + add(origin.replace("127.0.0.1", "localhost")) + return origins + + +class _OpencodeSession: + """Tracks the single OpenCode server process this UI server has launched. + + The browser cannot start a process, so the landing-page agent asks us to do it + (``POST /agent-server/start``). The child is rooted at the experiment + directory, which becomes the agent's workspace, and is torn down with this + UI server so it never outlives the session that spawned it. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._process: Optional[subprocess.Popen] = None + self._port: Optional[int] = None + self._workspace: Optional[str] = None + self._error: Optional[str] = None + self._log: list = [] + # Set when OPENCODE_URL points at an already-running server we adopted + # instead of spawning our own -- see ensure(). Mutually exclusive with + # self._process; only one of the two is ever active at a time. + self._external_url: Optional[str] = None + + def _running_locked(self) -> bool: + if self._external_url is not None: + return True + return self._process is not None and self._process.poll() is None + + def _url_locked(self) -> Optional[str]: + if self._external_url is not None: + return self._external_url + return f"http://127.0.0.1:{self._port}" if self._port else None + + def _drain_output(self, process: "subprocess.Popen[str]") -> None: + # Must be drained for the life of the process or the pipe buffer fills and + # blocks the child. We keep only a short tail, purely so a failed start can + # report why instead of a bare timeout. + try: + for line in process.stdout: # type: ignore[union-attr] + self._log.append(line.rstrip()) + del self._log[:-15] + except Exception: # pragma: no cover - best-effort log draining + pass + + def ensure(self, workspace_dir: str, origin: Optional[str]) -> dict: + """Start the agent server if it is not already running, and wait until it + answers a health check. Idempotent: a second call while alive is a no-op. + + If OPENCODE_URL is set and healthy, adopt it directly instead of + spawning a child -- this is what makes the shared root config in + agent.py's _load_config actually converge: point both the SDK agent + and this UI server at one already-running server via one env var, + rather than each independently spawning its own.""" + with self._lock: + if self._running_locked(): + return {"ok": True, "url": self._url_locked(), "reused": True, + "workspace": self._workspace} + + external_url = os.environ.get("OPENCODE_URL", "").strip() + if external_url and _opencode_healthy(external_url): + with self._lock: + self._external_url = external_url + self._workspace = workspace_dir + self._error = None + return {"ok": True, "url": external_url, "reused": False, "workspace": workspace_dir} + + with self._lock: + argv = _resolve_opencode_argv() + if argv is None: + self._error = ( + "Could not find `opencode` or `npx`. Install Node.js 20+ " + "(which provides npx), or `npm i -g opencode-ai`." + ) + return {"ok": False, "error": self._error} + + port = _free_port() + cmd = argv + ["serve", "--hostname", "127.0.0.1", "--port", str(port)] + for value in _cors_origin_variants(origin): + cmd += ["--cors", value] + + self._log = [] + self._error = None + try: + process = subprocess.Popen( + cmd, + cwd=workspace_dir, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + # Process-group leader so shutdown() can kill the whole tree -- + # npx spawns the real binary as a child, so killing only this + # PID would orphan the server. + start_new_session=True, + ) + except Exception as exc: # pragma: no cover - defensive + self._error = str(exc) + return {"ok": False, "error": self._error} + + self._process = process + self._port = port + self._workspace = workspace_dir + threading.Thread(target=self._drain_output, args=(process,), daemon=True).start() + + # Health-poll outside the lock so status() stays responsive while a cold + # npx download runs. + base_url = f"http://127.0.0.1:{port}" + deadline = time.monotonic() + _OPENCODE_START_TIMEOUT + while time.monotonic() < deadline: + if process.poll() is not None: + break + if _opencode_healthy(base_url): + return {"ok": True, "url": base_url, "reused": False, + "workspace": workspace_dir} + time.sleep(0.4) + + tail = " / ".join(self._log[-4:]) or "no output" + self._error = ( + f"The agent server did not come up within {int(_OPENCODE_START_TIMEOUT)}s. " + f"Last output: {tail}" + ) + _kill_process_tree(process) + with self._lock: + self._port = None + return {"ok": False, "error": self._error} + + def status(self) -> dict: + with self._lock: + process = self._process + url = self._url_locked() + workspace = self._workspace + error = self._error + external = self._external_url is not None + if external: + return {"state": "running", "url": url, "workspace": workspace} + if process is None: + return {"state": "none", "error": error} + if process.poll() is None: + return {"state": "running", "url": url, "workspace": workspace} + return {"state": "killed", "workspace": workspace, "error": error} + + def shutdown(self) -> None: + with self._lock: + process = self._process + # An externally-provided OPENCODE_URL server is not ours to kill -- + # just drop the reference so a later ensure() re-evaluates it fresh. + self._external_url = None + if process is None or process.poll() is not None: + return + _kill_process_tree(process) + + +_opencode_session = _OpencodeSession() +atexit.register(_opencode_session.shutdown) + + +# --------------------------------------------------------------------------- # +# /loop -- recurring OpenCode-backed monitoring jobs +# --------------------------------------------------------------------------- # + +# Minimum interval a loop can be scheduled at -- guards against a typo'd +# "/loop 30s ..." hammering the model every few seconds. +_LOOP_MIN_INTERVAL_SECONDS = 60.0 + +# Loops run against the live training process with a broad toolset (bash, +# file edits, pause/discard/restart) -- an unbounded number of them is an +# unbounded number of concurrent interventions. Shared across both chat +# surfaces since _loop_registry is a single process-wide instance. +_LOOP_MAX_CONCURRENT = 3 + +# Every mutating capability this job needs already exists as a local CLI verb +# (weightslab.cli:main, backed by weightslab/backend/cli.py's localhost TCP +# command server to the live training process) -- no new bridge/tool is built +# for this. The loop's OpenCode session gets the SAME full toolset the landing +# page chat runs with (bash/read/write/edit/patch, no restriction), so it can +# invoke these directly. +_LOOP_SYSTEM_PREAMBLE = ( + "You are a recurring monitoring agent for a live WeightsLab training run, " + "checking in periodically. Your workspace is the experiment directory; you " + "have bash, read, write, edit, and patch tools rooted there.\n\n" + "A local `weightslab` CLI is already available via bash for controlling " + "the live training run, not just editing files:\n" + " - `weightslab pause` / `weightslab resume` -- freeze/resume weight " + "updates (a safe pause; skips the optimizer step, does not kill the process)\n" + " - `weightslab discard ` -- discard one sample by id\n" + " - `weightslab agent query \"\"` -- ask the data-" + "manipulation agent to do something in plain English, e.g. " + "`weightslab agent query \"discard samples where loss > 5\"` or " + "`weightslab agent query \"tag mislabeled samples with tag:review\"` -- " + "this goes through the normal LLM-driven intent pipeline, not just exact ids\n" + " - `weightslab status` -- a snapshot of hyperparameters/model/training state\n\n" + "You may also read and edit training code directly (e.g. fix a broken " + "signal function) and, if training has crashed or stalled, attempt to " + "restart it via bash -- this is best-effort: look for the process " + "(`ps`/`pgrep`), stop it if still running, and re-launch the training " + "command if you can determine what it was from shell history, a run " + "script, or logs in this directory. There is no dedicated restart " + "command, so use your judgement and explain what you did.\n\n" + "Your monitoring task for this check-in:\n{prompt}" +) + + +def _opencode_json_request(base_url: str, path: str, method: str = "GET", body: Optional[dict] = None, timeout: float = 30.0): + import urllib.request + data = json.dumps(body).encode("utf-8") if body is not None else None + headers = {"Content-Type": "application/json"} if body is not None else {} + req = urllib.request.Request(f"{base_url.rstrip('/')}{path}", data=data, headers=headers, method=method) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def _opencode_send_and_collect(base_url: str, session_id: str, text: str, timeout: float = 600.0) -> str: + """Open the SSE event stream, THEN send the message, THEN read the stream + until this session goes idle -- same stream-first ordering and event + parsing as weights_studio/src/landing/agent/opencodeClient.ts and + weightslab/trainer/services/agent/opencode_chat.py's _collect_reply. + Duplicated in miniature here (rather than importing opencode_chat.py) + because this module is deliberately stdlib-only -- see its own docstring -- + and this loop wants the opposite tool policy (full toolset, no + restriction) from that module's SDK-agent use case anyway. + + A loop check-in can legitimately run long (the agent may read logs, edit + files, run training-control commands) -- default timeout is 10 minutes, + generous relative to the loop's own interval (minimum 1 minute), and a + slow check-in simply delays that job's next tick rather than blocking + anything else (each job's timer callback runs independently).""" + import urllib.error + import urllib.request + + stream_req = urllib.request.Request(f"{base_url.rstrip('/')}/event", headers={"Accept": "text/event-stream"}) + stream = urllib.request.urlopen(stream_req, timeout=timeout) + + send_errors: list = [] + + def _send() -> None: + try: + _opencode_json_request( + base_url, f"/session/{session_id}/message", method="POST", + body={"parts": [{"type": "text", "text": text}]}, timeout=timeout, + ) + except Exception as exc: # noqa: BLE001 - surfaced via send_errors + send_errors.append(exc) + + sender = threading.Thread(target=_send, daemon=True) + text_parts: dict = {} + assistant_message_ids: set = set() + + try: + sender.start() + data_lines: list = [] + deadline = time.monotonic() + timeout + for raw_line in stream: + if time.monotonic() > deadline: + break + line = raw_line.decode("utf-8", errors="replace").rstrip("\r\n") + if line == "": + if data_lines: + payload = "\n".join(data_lines) + data_lines = [] + try: + event = json.loads(payload) + except (ValueError, TypeError): + continue + event_type = event.get("type") + props = event.get("properties") or {} + if event_type == "message.updated": + msg = props.get("info") or props.get("message") or props + msg_session = str(msg.get("sessionID") or "") + if (not msg_session or msg_session == session_id) and \ + str(msg.get("role") or "") == "assistant" and msg.get("id"): + assistant_message_ids.add(str(msg["id"])) + elif event_type == "message.part.updated": + part = props.get("part") or props + part_session = str(part.get("sessionID") or "") + message_id = str(part.get("messageID") or "") + if (not part_session or part_session == session_id) and \ + message_id in assistant_message_ids and part.get("type") == "text" \ + and isinstance(part.get("text"), str): + text_parts[str(part.get("id") or message_id)] = part["text"] + elif event_type in ("session.idle", "session.error") and \ + str(props.get("sessionID") or "") == session_id: + break + continue + if line.startswith(":"): + continue + if line.startswith("data:"): + data_lines.append(line[5:].lstrip(" ")) + except (urllib.error.URLError, TimeoutError, OSError): + pass # degrade to whatever text was collected so far + finally: + stream.close() + + sender.join(timeout=1.0) + if send_errors and not text_parts: + raise send_errors[0] + + return "".join(text_parts[key] for key in text_parts) + + +class _LoopJob: + def __init__(self, job_id: str, prompt: str, interval_seconds: float, workspace: str) -> None: + self.id = job_id + self.prompt = prompt + self.interval_seconds = interval_seconds + self.workspace = workspace + self.session_id: Optional[str] = None + self.base_url: Optional[str] = None + self.next_run_at: Optional[float] = None + self.last_result: Optional[str] = None + self.last_error: Optional[str] = None + self.timer: Optional[threading.Timer] = None + self.stopped = False + + +class _LoopRegistry: + """Recurring OpenCode-backed monitoring jobs, e.g. `/loop 30m `. + + Lives here (this UI server process) rather than in the browser tab or the + training backend: the browser can't run a persistent timer that survives + a page reload, and the training backend doesn't need to be touched at all + since every intervention the loop needs is already reachable through the + local `weightslab` CLI over bash (see _LOOP_SYSTEM_PREAMBLE). A job + survives a page reload/tab close (tied to this process, not the tab) but + not a full `weightslab start` restart -- no persistence beyond that, + matching every other piece of state in this server (Jupyter/OpenCode + sessions). + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._jobs: dict = {} + self._next_id = 1 + + def start(self, prompt: str, interval_seconds: float, workspace_dir: str, origin: Optional[str]) -> dict: + prompt = (prompt or "").strip() + if not prompt: + return {"ok": False, "error": "A monitoring prompt is required."} + if interval_seconds < _LOOP_MIN_INTERVAL_SECONDS: + return {"ok": False, "error": f"Minimum loop interval is {int(_LOOP_MIN_INTERVAL_SECONDS)}s."} + + with self._lock: + if len(self._jobs) >= _LOOP_MAX_CONCURRENT: + return { + "ok": False, + "error": f"{_LOOP_MAX_CONCURRENT} loops already running -- " + f"stop one first (\"/loop stop \").", + } + + ensured = _opencode_session.ensure(workspace_dir, origin) + if not ensured.get("ok"): + return {"ok": False, "error": ensured.get("error") or "Could not start the agent server."} + base_url = ensured["url"] + + with self._lock: + if len(self._jobs) >= _LOOP_MAX_CONCURRENT: + return { + "ok": False, + "error": f"{_LOOP_MAX_CONCURRENT} loops already running -- " + f"stop one first (\"/loop stop \").", + } + job_id = str(self._next_id) + self._next_id += 1 + job = _LoopJob(job_id, prompt, interval_seconds, workspace_dir) + job.base_url = base_url + self._jobs[job_id] = job + + threading.Thread(target=self._fire, args=(job_id, base_url), daemon=True).start() + return {"ok": True, "id": job_id, "intervalSeconds": interval_seconds} + + def _fire(self, job_id: str, base_url: str) -> None: + with self._lock: + job = self._jobs.get(job_id) + if job is None or job.stopped: + return # stopped before this tick ran + + try: + if job.session_id is None: + created = _opencode_json_request( + base_url, "/session", method="POST", + body={"title": f"weightslab-loop-{job_id}"}, + ) + job.session_id = created["id"] + text = _LOOP_SYSTEM_PREAMBLE.format(prompt=job.prompt) + else: + text = job.prompt + result = _opencode_send_and_collect(base_url, job.session_id, text) + with self._lock: + if job_id in self._jobs: + job.last_result = result + job.last_error = None + except Exception as exc: + with self._lock: + if job_id in self._jobs: + job.last_error = str(exc) + + with self._lock: + job = self._jobs.get(job_id) + if job is None or job.stopped: + return # stopped while this tick was running + job.next_run_at = time.time() + job.interval_seconds + timer = threading.Timer(job.interval_seconds, self._fire, args=(job_id, base_url)) + timer.daemon = True + job.timer = timer + timer.start() + + def stop(self, job_id: str) -> dict: + with self._lock: + job = self._jobs.pop(job_id, None) + if job is not None: + job.stopped = True + if job is None: + return {"ok": False, "error": f"No loop job {job_id}."} + if job.timer is not None: + job.timer.cancel() + return {"ok": True} + + def update(self, job_id: str, prompt: Optional[str] = None, interval_seconds: Optional[float] = None) -> dict: + """Change a running job's prompt and/or interval in place. An interval + change reschedules from now rather than waiting out the old timer, so + the change is felt immediately instead of on the tick after next.""" + if prompt is not None and not prompt.strip(): + return {"ok": False, "error": "The monitoring prompt cannot be empty."} + if interval_seconds is not None and interval_seconds < _LOOP_MIN_INTERVAL_SECONDS: + return {"ok": False, "error": f"Minimum loop interval is {int(_LOOP_MIN_INTERVAL_SECONDS)}s."} + + with self._lock: + job = self._jobs.get(job_id) + if job is None: + return {"ok": False, "error": f"No loop job {job_id}."} + + if prompt is not None: + job.prompt = prompt.strip() + + reschedule = interval_seconds is not None and interval_seconds != job.interval_seconds + if interval_seconds is not None: + job.interval_seconds = interval_seconds + + if reschedule and job.timer is not None: + job.timer.cancel() + job.next_run_at = time.time() + job.interval_seconds + timer = threading.Timer(job.interval_seconds, self._fire, args=(job_id, job.base_url)) + timer.daemon = True + job.timer = timer + timer.start() + + return { + "ok": True, + "id": job.id, + "prompt": job.prompt, + "intervalSeconds": job.interval_seconds, + "nextRunAt": job.next_run_at, + } + + def list(self) -> list: + with self._lock: + jobs = list(self._jobs.values()) + return [ + { + "id": j.id, + "prompt": j.prompt, + "intervalSeconds": j.interval_seconds, + "nextRunAt": j.next_run_at, + "lastResult": j.last_result, + "lastError": j.last_error, + } + for j in jobs + ] + + def shutdown(self) -> None: + with self._lock: + jobs = list(self._jobs.values()) + self._jobs.clear() + for job in jobs: + job.stopped = True + if job.timer is not None: + job.timer.cancel() + + +_loop_registry = _LoopRegistry() +atexit.register(_loop_registry.shutdown) + + # --------------------------------------------------------------------------- # # Request handler # --------------------------------------------------------------------------- # @@ -391,6 +942,12 @@ def do_GET(self): # noqa: N802 if path == "/local-notebook/status": self._send_json(HTTPStatus.OK, _jupyter_session.status()) return + if path == "/agent-server/status": + self._send_json(HTTPStatus.OK, _opencode_session.status()) + return + if path == "/agent-server/loop/list": + self._send_json(HTTPStatus.OK, {"loops": _loop_registry.list()}) + return if path == "/local-notebook/list": self._list_local_notebooks() return @@ -409,6 +966,16 @@ def do_POST(self): # noqa: N802 path = self.path.split("?", 1)[0] if path == "/local-notebook": self._start_local_notebook() + elif path == "/agent-server/start": + self._start_agent_server() + elif path == "/agent-server/loop/start": + self._start_loop() + elif path.startswith("/agent-server/loop/") and path.endswith("/stop"): + loop_id = path[len("/agent-server/loop/"):-len("/stop")] + self._stop_loop(loop_id) + elif path.startswith("/agent-server/loop/") and path.endswith("/update"): + loop_id = path[len("/agent-server/loop/"):-len("/update")] + self._update_loop(loop_id) elif path.startswith(self.api_prefix + "/") or path == self.api_prefix: self._proxy_grpc_web(path) else: @@ -524,6 +1091,96 @@ def _list_local_notebooks(self): }) self._send_json(HTTPStatus.OK, {"notebooks": entries}) + def _start_agent_server(self): + """Start (or reuse) the OpenCode server backing the landing-page agent. + + Same shape and same reasoning as ``_start_local_notebook`` below: spawning + a process is an OS-level action, so the browser asks us to do it. Rooted at + the experiment directory, which becomes the agent's workspace. Loopback-only, + like every other local-machine action in this server -- this one starts a + process with filesystem access, so it must never be reachable off-host. + """ + if self.client_address[0] not in _LOOPBACK_ADDRESSES: + self._send_json(HTTPStatus.FORBIDDEN, + {"ok": False, "error": "Only reachable from localhost."}) + return + + workspace = self.experiment_dir or os.getcwd() + os.makedirs(workspace, exist_ok=True) + + # Prefer the browser's own Origin header: it is the exact string the CORS + # check will compare against. Fall back to reconstructing it from Host. + origin = self.headers.get("Origin") + if not origin: + host = self.headers.get("Host") or "localhost" + scheme = "https" if isinstance(getattr(self, "connection", None), ssl.SSLSocket) else "http" + origin = f"{scheme}://{host}" + + result = _opencode_session.ensure(workspace, origin) + status = HTTPStatus.OK if result.get("ok") else HTTPStatus.INTERNAL_SERVER_ERROR + self._send_json(status, result) + + def _start_loop(self): + """Start a recurring OpenCode-backed monitoring job (the `/loop` + command in the connected-experiment agent bar). Loopback-only, same + reasoning as _start_agent_server -- this also starts/reuses that same + process.""" + if self.client_address[0] not in _LOOPBACK_ADDRESSES: + self._send_json(HTTPStatus.FORBIDDEN, + {"ok": False, "error": "Only reachable from localhost."}) + return + + body = self._read_json_body() + prompt = str(body.get("prompt") or "") + try: + interval_minutes = float(body.get("intervalMinutes") or 0) + except (TypeError, ValueError): + interval_minutes = 0 + + workspace = self.experiment_dir or os.getcwd() + os.makedirs(workspace, exist_ok=True) + + origin = self.headers.get("Origin") + if not origin: + host = self.headers.get("Host") or "localhost" + scheme = "https" if isinstance(getattr(self, "connection", None), ssl.SSLSocket) else "http" + origin = f"{scheme}://{host}" + + result = _loop_registry.start(prompt, interval_minutes * 60.0, workspace, origin) + status = HTTPStatus.OK if result.get("ok") else HTTPStatus.BAD_REQUEST + self._send_json(status, result) + + def _stop_loop(self, loop_id: str): + if self.client_address[0] not in _LOOPBACK_ADDRESSES: + self._send_json(HTTPStatus.FORBIDDEN, + {"ok": False, "error": "Only reachable from localhost."}) + return + result = _loop_registry.stop(loop_id) + status = HTTPStatus.OK if result.get("ok") else HTTPStatus.NOT_FOUND + self._send_json(status, result) + + def _update_loop(self, loop_id: str): + """Change a running loop's prompt and/or interval (the panel's Edit + action) without stopping and restarting the job.""" + if self.client_address[0] not in _LOOPBACK_ADDRESSES: + self._send_json(HTTPStatus.FORBIDDEN, + {"ok": False, "error": "Only reachable from localhost."}) + return + + body = self._read_json_body() + prompt = body.get("prompt") + interval_seconds = None + if body.get("intervalMinutes") is not None: + try: + interval_seconds = float(body["intervalMinutes"]) * 60.0 + except (TypeError, ValueError): + self._send_json(HTTPStatus.BAD_REQUEST, {"ok": False, "error": "intervalMinutes must be a number."}) + return + + result = _loop_registry.update(loop_id, prompt=prompt, interval_seconds=interval_seconds) + status = HTTPStatus.OK if result.get("ok") else HTTPStatus.BAD_REQUEST + self._send_json(status, result) + def _start_local_notebook(self): # Copying a file and spawning a process are OS-level actions the # browser can't do itself; this is the one piece of local-only @@ -779,6 +1436,11 @@ def _js_str(value: Optional[str]) -> str: ("ENABLE_HYPERPARAMETERS_OPTIMIZATION", "WS_ENABLE_HYPERPARAMETERS_OPTIMIZATION", "VITE_ENABLE_HYPERPARAMETERS_OPTIMIZATION")), ("WS_ENABLE_AGENT", ("ENABLE_AGENT", "WS_ENABLE_AGENT", "VITE_ENABLE_AGENT")), + # Shared with the SDK agent's own OPENCODE_URL config (agent.py's + # _load_config) -- setting OPENCODE_URL once configures both sides, and + # is also what _OpencodeSession.ensure() prefers over spawning its own + # server (see below). + ("WS_OPENCODE_URL", ("OPENCODE_URL", "WS_OPENCODE_URL", "VITE_OPENCODE_URL")), ] From 4406830e04051afd474e9c3c81c4dfb617b493a4 Mon Sep 17 00:00:00 2001 From: GuillaumePELLUET Date: Fri, 7 Aug 2026 17:10:50 +0200 Subject: [PATCH 02/11] Give each /loop job its own OpenCode session created up front (not lazily on the first tick) and guarded by a lock, and add message/messages endpoints so a loop's monitoring chat can be read from and written to directly -- backing weights_studio's new per-loop chat tabs, where a running check-in and a manual message must never race on the same session. Co-Authored-By: Claude Sonnet 5 --- AGENTS.md | 280 -------------- pyproject.toml | 5 + .../services/test_agent_opencode_provider.py | 89 +++++ .../services/test_agent_service_unit.py | 33 ++ tests/trainer/services/test_opencode_chat.py | 58 ++- tests/ui/test_server_agent.py | 294 ++++++++++++++ weightslab/AGENTS.md | 362 ++++++++++++++++++ weightslab/proto/experiment_service.proto | 19 + weightslab/proto/experiment_service_pb2.py | 100 ++--- .../proto/experiment_service_pb2_grpc.py | 49 +++ weightslab/trainer/services/agent/agent.py | 75 +++- .../trainer/services/agent/opencode_chat.py | 36 +- weightslab/trainer/services/agent_service.py | 26 ++ weightslab/ui/server.py | 263 ++++++++++++- 14 files changed, 1339 insertions(+), 350 deletions(-) delete mode 100644 AGENTS.md create mode 100644 weightslab/AGENTS.md diff --git a/AGENTS.md b/AGENTS.md deleted file mode 100644 index 9629fab8..00000000 --- a/AGENTS.md +++ /dev/null @@ -1,280 +0,0 @@ -# WeightsLab — agent context for users & debugging - -This file is a **portable context for AI coding agents** (Claude Code, etc.) and -the humans driving them. Its job is to let you — or an agent helping you — -**install, configure, run, and debug WeightsLab and Weights Studio** without -having to reverse-engineer the system first. - -It deliberately covers only the two shipped repositories: - -- **weightslab** — the Python backend / core (training instrumentation, data - ledger, gRPC service, the shared proto). -- **weights_studio** — the browser frontend (the studio UI that inspects and - edits a *running* experiment). - -> File/line references drift as the code evolves — treat them as starting points -> and verify against the current source before relying on them. Environment -> variable names and defaults are the most stable thing here; when in doubt the -> authoritative reference is `weightslab/docs/configuration.rst`. - ---- - -## 0. How to load this guide into Claude Code - -So an agent actually *has* this context when you ask it for help: - -- **Working inside a checkout of the repo** (`git clone`): this guide is - committed as `AGENTS.md`; the repo keeps a gitignored `CLAUDE.md` copy of it at - the root so Claude Code auto-loads it every session. Nothing to do. (Claude - Code also loads `~/.claude/CLAUDE.md` global memory and any parent-dir - `CLAUDE.md`.) -- **You only ran `pip install weightslab`** (no checkout — the package lives in - `site-packages`): absolute `@import` paths are fragile because the path - changes per venv/OS. The robust pattern is a small **skill** that locates the - installed file at runtime. Create `~/.claude/skills/weightslab/SKILL.md`: - - ```yaml - --- - name: weightslab - description: Load the WeightsLab debugging & configuration guide when helping with weightslab or weights_studio problems (connection, TLS, env vars, training hangs, rendering). - --- - !`python -c "import weightslab, os; print(open(os.path.join(os.path.dirname(weightslab.__file__), 'AGENTS.md')).read())"` - - Use the guide above to diagnose the user's weightslab / weights_studio issue. - ``` - - Then run `/weightslab` (or let Claude auto-invoke it). This requires the guide - to be **shipped as package data** inside the installed package (see §7); the - copy at the repo root is for contributors working in a checkout. -- **Quick-and-dirty:** copy this file to `~/.claude/WEIGHTSLAB.md` and add - `@~/.claude/WEIGHTSLAB.md` to your `~/.claude/CLAUDE.md`. - ---- - -## 1. What it is and how the pieces connect - -A user wraps their own PyTorch training script with WeightsLab so a running -experiment becomes inspectable/editable; Weights Studio is the UI for that. - -**Wire path (the thing that breaks most often):** - -``` -Browser → weightslab start :8080 (grpc-web → grpc proxy) → Python gRPC servicer → training loop -``` - -- `weightslab start` is a pure-Python HTTP server that serves the bundled SPA - and translates grpc-web (browser) to raw gRPC (backend). No Docker, no Envoy. - If `weightslab start` is not running, the browser has no UI to load. -- The gRPC servicer and the training loop run in the **same process, different - threads**, coordinated by locks in - `weightslab/weightslab/components/global_monitoring.py`. -- One proto is the single source of truth: - `weightslab/weightslab/proto/experiment_service.proto`. - ---- - -## 2. Install & run (the happy path) - -```bash -pip install weightslab -``` - -In your training script: - -```python -import weightslab as wl -# wrap your objects so the studio can see/edit them (see §3), then: -wl.serve(serving_grpc=True, serving_cli=True) # background threads, same process -# ... your training loop ... -wl.keep_serving() # keep the process alive for the UI -``` - -Then start the UI in another terminal and open it in a browser: - -```bash -weightslab start # serves at http://localhost:8080 by default -``` - -Working starting points live in -`weightslab/weightslab/examples/{PyTorch,Lightning,Usecases}//` -(each is a `main.py` + `config.yaml`) — find the closest example and mirror it. - -UI deployment details (port, TLS, certs) are documented in -`weightslab/docs/weights_studio.rst`. TLS is opt-in: run `weightslab se` once, -then `weightslab start --certs`. - ---- - -## 3. The integration API (`import weightslab as wl`) - -How a user's script plugs in. Wrap each training object with -`wl.watch_or_edit(obj, flag=...)`; the returned tracked proxy is registered in -the global ledger (`weightslab/weightslab/backend/ledgers.py`, -`GLOBAL_LEDGER` — the hub everything reads/mutates through). - -- `flag="hyperparameters"` (dict), `flag="model"` (nn.Module, `device=…`), - `flag="optimizer"`, `flag="data"` (Dataset → tracked DataLoader: `loader_name`, - `batch_size`, `is_training`, `collate_fn`, …), `flag="loss"` (a - `reduction="none"` criterion, called with `(preds_raw, targets, batch_ids=ids, - preds=preds)`), `flag="metric"`. - -Conventions that matter for correctness: - -- Wrap the train step in `with guard_training_context:` and eval in - `with guard_testing_context:` (from - `weightslab.components.global_monitoring`). This is how pause/resume and - train/test separation work — **skip it and pause/resume or stats will misbehave.** -- Use `model.get_age()` (steps actually trained; survives checkpoint reloads), - not the raw loop counter. -- `task_type` on the dataset/model selects rendering: `classification`, - `segmentation`, `detection`, `detection_pointcloud`. -- **Hyperparameter handle access:** the registered hyperparameters proxy - supports both `hp.get("lr")` and `hp["lr"]` (subscript == `.get`), and stays - live — reads reflect in-place updates and re-registration. - ---- - -## 4. Configuration (environment variables) - -WeightsLab and Weights Studio are configured almost entirely through env vars. -**Authoritative reference: `weightslab/docs/configuration.rst`.** The high-signal -ones when debugging: - -**Backend (Python):** - -| Variable | Default | Why you touch it | -|---|---|---| -| `WEIGHTSLAB_LOG_LEVEL` | `INFO` | Set `DEBUG` to see what's happening. (`WATCHDOG` level sits between WARNING/ERROR.) | -| `GRPC_BACKEND_HOST` / `GRPC_BACKEND_PORT` | `0.0.0.0` / `50051` | Backend gRPC bind address. | -| `GRPC_TLS_ENABLED` | `0` | TLS on the gRPC socket. Set `1` with `weightslab start --certs`. | -| `GRPC_TLS_REQUIRE_CLIENT_AUTH` | `0` | mTLS. Must match what `weightslab start --certs` presents. | -| `WEIGHTSLAB_CERTS_DIR` | `~/.weightslab-certs` | Where cert files are looked up (single source of truth). | -| `GRPC_AUTH_TOKEN` | *(unset)* | Optional metadata-token auth on top of mTLS. | -| `GRPC_MAX_MESSAGE_BYTES` | `268435456` (256 MB) | Raise it if large tensors/image batches fail. | -| `WEIGHTSLAB_DISABLE_WATCHDOGS` | `0` | Set `1` when debugging with breakpoints (see §5). | -| `GRPC_WATCHDOG_STUCK_SECONDS` | `60` | Lock/RPC stuck threshold + lock-acquire timeout. | - -**Frontend (Weights Studio) — runtime-injected `window.*` globals:** - -| Variable | Default | Why you touch it | -|---|---|---| -| `WS_SERVER_HOST` / `WS_SERVER_PORT` / `WS_SERVER_PROTOCOL` | `localhost` / `8080` / `http` | How the browser reaches the `weightslab start` server. The #1 connection-issue knob. | -| `WS_HISTOGRAM_MAX_BINS` | `512` | Cap on metadata histogram bars. | -| `BB_THUMB_RENDER` | `10` | Max bounding boxes drawn per **thumbnail**, per overlay (GT and PRED capped independently). | -| `BB_MODAL_RENDER` | `100` | Max bounding boxes drawn per **modal** image, per overlay. A `?` button in the modal shows the active limit. | -| `ENABLE_PLOTS` | `1` | `0`/`false` removes the plots board + Signals card and stops plot auto-refresh. | -| `ENABLE_DATA_EXPLORATION` | `1` | `0`/`false` removes the data grid + metadata/details panel and stops the data/metadata auto-refresh. | -| `ENABLE_HYPERPARAMETERS_OPTIMIZATION` | `1` | `0`/`false` removes the Hyperparameters section, makes HP inputs read-only, and stops the HP poll. | -| `ENABLE_AGENT` | `1` | `0`/`false` removes the agent chat bar + history panel and stops the agent health poll. | -| `ENABLE_NOTEBOOK` | `1` | `0`/`false` removes the notebook button (left of the logo) + notebook window. The notebook runs Python in a shared in-process kernel against the live experiment (`df`, model, checkpoints), persisted as `notebook.ipynb` under `root_log_dir`; `>`-prefixed cells ask the agent to propose code. | - -> **VITE_ vs WS_/BB_/ENABLE_:** `VITE_*` variables are baked at **build time** -> (changing them needs a frontend rebuild). `WS_*` / `BB_*` / `ENABLE_*` are -> injected into `config.js` at `weightslab start` time and read as `window.*` -> globals — changing them needs only a restart + browser reload. Each `ENABLE_*` -> defaults to on; set it to `0`/`false`/`no`/`off` to disable. Full reference: -> `weightslab/docs/configuration.rst` (“Feature toggles”). - ---- - -## 5. Troubleshooting — symptom → cause → fix - -This is the core of the guide. Each entry is a real failure mode (several are -distilled from issues hit in development). - -**UI loads but the sample grid is empty / "failed to fetch" / gRPC errors.** -The wire path (§1) is broken somewhere. Check in order: (1) backend actually -serving on `0.0.0.0:50051`; (2) `weightslab start` is running and the browser -can reach it on `:8080`; (3) **TLS mismatch** if using `--certs` — run -`weightslab se` first and export `WEIGHTSLAB_CERTS_DIR`. For local debugging -drop TLS entirely (omit `--certs`; `GRPC_TLS_ENABLED=0`). - -**Changed an env var, restarted, but the UI still uses the old value.** -- `VITE_*` is build-time → you must **rebuild** the frontend, not just restart. -- `WS_*` / `BB_*` / `ENABLE_*` are injected at `weightslab start` time → you - must **restart `weightslab start`** then reload the tab. - -**Sample grid flashes empty cells when auto-refresh fires.** -An auto-refresh (timer or manual) that lands while a `GetDataSamples` grid fetch -is still in flight used to clear the cache mid-render. The fix in -`weights_studio/src/grid_data/gridDataManager.ts` is `isFetchInProgress()`: -refreshes are skipped while a grid fetch is ongoing. If you see this, confirm -you're on a build that has that guard. - -**Detection overlays are slow or unreadably cluttered.** -Dense detection samples can carry hundreds of boxes. Cap rendering with -`BB_THUMB_RENDER` (thumbnails) and `BB_MODAL_RENDER` (modal); each is applied -separately to GT and to predictions. Render-only — no sample data is dropped. - -**Training appears hung; RPCs return `RESOURCE_EXHAUSTED`; server "restarts".** -A watchdog monitors the global rlock and in-flight RPCs. If a lock/RPC is held -longer than `GRPC_WATCHDOG_STUCK_SECONDS` (60s) it's flagged; locks get -interrupted, and after `GRPC_WATCHDOG_RESTART_THRESHOLD` unhealthy polls the -gRPC server restarts. When **debugging with breakpoints** that intentionally -pause longer than that, set `WEIGHTSLAB_DISABLE_WATCHDOGS=1`. If RPCs fail with -`RESOURCE_EXHAUSTED`, a handler couldn't acquire the lock within the window — -something else is holding it; check for a long/blocking train or eval step. - -**Pause/resume doesn't work, or train vs test stats are mixed up.** -The train step isn't wrapped in `guard_training_context` (or eval in -`guard_testing_context`). See §3 — these context managers are how the system -gates and separates phases. - -**Large weights/images fail to transfer.** Raise `GRPC_MAX_MESSAGE_BYTES`. - -**The agent bar says it's unconfigured.** The LLM agent is backed entirely by -a local **OpenCode** server (`OPENCODE_URL`, default `http://127.0.0.1:4096`) -— WeightsLab starts one for you on first use. Initialize from the UI via -`/init` (then `/model` to switch, `/reset` to clear). See -`weightslab/docs/agent.rst` and `weightslab/docs/weights_studio.rst`. - ---- - -## 6. Where things live (for deeper digging) - -**Backend (`weightslab/weightslab/`):** -- `src.py` — the public verbs (`watch_or_edit`, `serve`, `keep_serving`, - `tag_samples`, `query_*`, decorators) re-exported from `__init__.py`. -- `trainer/services/` — `experiment_service.py` (gRPC servicer) delegating to - `{model,data,agent}_service.py`; `data_image_utils.py` (preview/mask encoding). -- `components/` — `global_monitoring.py` (locks, `guard_*` contexts, pause), - `checkpoint_manager.py`, `evaluation_controller.py`. -- `data/` — `dataframe_manager.py`, `data_samples_with_ops.py`, `sample_stats.py`, - H5 storage (`h5_dataframe_store.py`, `h5_array_store.py`, `array_proxy.py`). -- `backend/` — `ledgers.py` (`GLOBAL_LEDGER`), `logger.py`, `audit_logger.py`, `cli.py`. -- `security/` (`CertAuthManager`), `proto/`, `examples/`, `docs/`. - -**Frontend (`weights_studio/src/`):** -- `main.ts` — bootstrap; builds the grpc-web transport from `WS_SERVER_*`. -- `experiment_service.client.ts` / `experiment_service.ts` — generated client - (regenerate with `npm run generate-proto:data`; do not hand-edit). -- `grid_data/` — grid + modal rendering (`GridCell.ts`, `DataImageService.ts`, - `gridDataManager.ts`, `BboxRenderer.ts`, `SegmentationRenderer.ts`, - `PointCloudViewer.ts`). -- `ui/` — `server.py` (pure-Python HTTP + gRPC-Web proxy), `static/` (bundled SPA), - `utils/` (cert-generation scripts, sync-frontend helper). - -**Docs:** `weightslab/docs/` (Sphinx) — `configuration.rst` (all env vars), -`weights_studio.rst` (studio deploy + agent), `quickstart.rst`, `grpc/`. - ---- - -## 7. For contributors (working in a checkout) - -- **The two repos must sit side by side** (`…/weightslab`, `…/weights_studio`); - proto codegen scripts reach across by relative path. -- **Editing the proto is cross-repo** — do all of: edit - `experiment_service.proto`; regenerate Python stubs from the repo root; run - `npm run generate-proto:data` in weights_studio. Editing one side only leaves - the build broken. -- **Tests:** backend `python -m pytest weightslab/tests/...`; frontend unit - `npm run test` (vitest); E2E/user-simulation Playwright lives in - **weights_studio** (`test:realtime:*`, `test:e2e:*`), not here. -- **CI on a custom branch:** pushes to non-`main`/`dev` branches only run CI when - the commit message contains `[force ci]` (both repos). -- **TLS/auth in the bundled UI** is decided by cert presence under - `WEIGHTSLAB_CERTS_DIR` (single source of truth) — don't hardcode secure/insecure. -- **To make this guide available to pip users**, ship it as package data inside - the installed package (e.g. as `weightslab/weightslab/AGENTS.md`) so the §0 - skill can locate it; keep the root `AGENTS.md` (mirrored as the gitignored - `CLAUDE.md`) as the contributor-facing source. diff --git a/pyproject.toml b/pyproject.toml index 3a741a79..c4f99f08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -154,6 +154,11 @@ weightslab = [ "ui/static/**/*", # Logo used to brand generated experiment reports (weightslab/reporting.py). "assets/**/*", + # weightslab-integration grounding for the landing chat's preset prompts + # and for agents dropped directly into a workspace (see server.py's + # _ensure_workspace_agents_md) -- lives inside the package, not just at the + # repo root, so a `pip install weightslab` ships it too. + "AGENTS.md", ] [tool.setuptools.exclude-package-data] diff --git a/tests/trainer/services/test_agent_opencode_provider.py b/tests/trainer/services/test_agent_opencode_provider.py index d1fd4bff..e7c8f403 100644 --- a/tests/trainer/services/test_agent_opencode_provider.py +++ b/tests/trainer/services/test_agent_opencode_provider.py @@ -11,6 +11,7 @@ """ import importlib +import json import sys import types import unittest @@ -206,6 +207,94 @@ def test_reports_failure_when_opencode_server_unreachable(self): self.assertIn("OpenCode", message) +class TestGetContextUsage(unittest.TestCase): + """DataManipulationAgent.get_context_usage() -- backs the /context + command. Combines OpenCodeChat.last_usage (mocked directly here; its own + population from real SSE events is covered by test_opencode_chat.py) with + a context-window lookup via /config/providers (mocked urllib, same + _FakeResp pattern as TestGetAvailableModelsOpenCode above).""" + + class _FakeResp: + def __init__(self, payload): + self._payload = payload + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return json.dumps(self._payload).encode() + + def test_not_configured_when_opencode_chat_is_none(self): + _, agent = _make_agent() + agent._opencode_chat = None + + ok, usage, message = agent.get_context_usage() + + self.assertFalse(ok) + self.assertEqual(usage, {}) + self.assertIn("/init", message) + + def test_reports_no_turns_yet_when_last_usage_is_none(self): + _, agent = _make_agent() + agent.opencode_model = "" + agent._opencode_chat = MagicMock(last_usage=None) + + ok, usage, message = agent.get_context_usage() + + self.assertTrue(ok) + self.assertEqual(usage["context_window"], 0) + self.assertIn("No agent turns yet", message) + + def test_combines_last_usage_with_the_models_context_window(self): + _, agent = _make_agent() + agent.opencode_url = "http://127.0.0.1:4096" + agent.opencode_model = "openrouter/anthropic/claude-opus-4.6" + agent._opencode_chat = MagicMock(last_usage={ + "input": 100, "output": 20, "reasoning": 5, "cache_read": 60, "cache_write": 10, + }) + + payload = { + "providers": [ + {"id": "openrouter", "models": { + "anthropic/claude-opus-4.6": {"limit": {"context": 200000, "output": 8192}}, + }}, + ], + } + with mock.patch("urllib.request.urlopen", return_value=self._FakeResp(payload)): + ok, usage, message = agent.get_context_usage() + + self.assertTrue(ok) + self.assertEqual(message, "") + self.assertEqual(usage, { + "model": "openrouter/anthropic/claude-opus-4.6", + "context_window": 200000, + "input_tokens": 100, + "output_tokens": 20, + "reasoning_tokens": 5, + "cache_read_tokens": 60, + "cache_write_tokens": 10, + }) + + def test_context_window_defaults_to_zero_when_the_server_is_unreachable(self): + """A failed /config/providers lookup must not sink the whole command -- + usage numbers are still worth showing without a window/percentage.""" + _, agent = _make_agent() + agent.opencode_model = "openrouter/anthropic/claude-opus-4.6" + agent._opencode_chat = MagicMock(last_usage={ + "input": 10, "output": 2, "reasoning": 0, "cache_read": 0, "cache_write": 0, + }) + + with mock.patch("urllib.request.urlopen", side_effect=OSError("refused")): + ok, usage, message = agent.get_context_usage() + + self.assertTrue(ok) + self.assertEqual(usage["context_window"], 0) + self.assertEqual(usage["input_tokens"], 10) + + class TestResetConnection(unittest.TestCase): def test_reset_clears_opencode_chain_and_model(self): agent_mod, agent = _make_agent() diff --git a/tests/trainer/services/test_agent_service_unit.py b/tests/trainer/services/test_agent_service_unit.py index 36b8f5aa..c3d394ed 100644 --- a/tests/trainer/services/test_agent_service_unit.py +++ b/tests/trainer/services/test_agent_service_unit.py @@ -147,6 +147,39 @@ def test_clear_and_compact_history_fail_cleanly_when_agent_backend_missing(self) self.assertFalse(compact_response.success) self.assertIn('not running', clear_response.message) + def test_get_agent_context_usage_delegates_to_agent(self): + agent = MagicMock() + agent.get_context_usage.return_value = (True, { + 'model': 'openrouter/anthropic/claude-opus-4.6', + 'context_window': 200000, + 'input_tokens': 100, + 'output_tokens': 20, + 'reasoning_tokens': 5, + 'cache_read_tokens': 60, + 'cache_write_tokens': 10, + }, '') + service, _ = self._make_service(agent=agent) + + response = service.GetAgentContextUsage(pb2.Empty(), None) + + agent.get_context_usage.assert_called_once_with() + self.assertTrue(response.success) + self.assertEqual(response.model, 'openrouter/anthropic/claude-opus-4.6') + self.assertEqual(response.context_window, 200000) + self.assertEqual(response.input_tokens, 100) + self.assertEqual(response.output_tokens, 20) + self.assertEqual(response.reasoning_tokens, 5) + self.assertEqual(response.cache_read_tokens, 60) + self.assertEqual(response.cache_write_tokens, 10) + + def test_get_agent_context_usage_fails_cleanly_when_agent_backend_missing(self): + service, _ = self._make_service(agent=None) + + response = service.GetAgentContextUsage(pb2.Empty(), None) + + self.assertFalse(response.success) + self.assertIn('not running', response.message) + if __name__ == '__main__': unittest.main() diff --git a/tests/trainer/services/test_opencode_chat.py b/tests/trainer/services/test_opencode_chat.py index 52840b2c..e16ccb04 100644 --- a/tests/trainer/services/test_opencode_chat.py +++ b/tests/trainer/services/test_opencode_chat.py @@ -86,7 +86,10 @@ def emit(event: dict) -> None: while not self.server.recorded_messages and time.monotonic() < deadline: time.sleep(0.01) - emit({"type": "message.updated", "properties": {"info": {"id": "msg_1", "role": "assistant", "sessionID": session_id}}}) + info = {"id": "msg_1", "role": "assistant", "sessionID": session_id} + if self.server.reply_tokens is not None: + info["tokens"] = self.server.reply_tokens + emit({"type": "message.updated", "properties": {"info": info}}) for delta in self.server.reply_deltas: emit({"type": "message.part.updated", "properties": {"part": { "id": f"prt_{delta[:4]}", "type": "text", "text": delta, "messageID": "msg_1", "sessionID": session_id, @@ -103,19 +106,21 @@ def emit(event: dict) -> None: class _FakeOpenCodeServer(ThreadingHTTPServer): daemon_threads = True - def __init__(self, *args, reply_deltas, emit_error=False, **kwargs): + def __init__(self, *args, reply_deltas, emit_error=False, reply_tokens=None, **kwargs): super().__init__(*args, **kwargs) self.session_id = "ses_test1" self.reply_deltas = reply_deltas self.emit_error = emit_error + self.reply_tokens = reply_tokens self.recorded_messages = [] self.recorded_session_titles = [] class _ServerTestCase(unittest.TestCase): - def _start_server(self, reply_deltas, emit_error=False): + def _start_server(self, reply_deltas, emit_error=False, reply_tokens=None): self.httpd = _FakeOpenCodeServer( - ("127.0.0.1", 0), _FakeOpenCodeHandler, reply_deltas=reply_deltas, emit_error=emit_error, + ("127.0.0.1", 0), _FakeOpenCodeHandler, + reply_deltas=reply_deltas, emit_error=emit_error, reply_tokens=reply_tokens, ) self.port = self.httpd.server_address[1] self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) @@ -194,6 +199,29 @@ def test_raises_opencode_error_when_the_server_is_unreachable(self): with self.assertRaises(OpenCodeError): chat._call("hello") + def test_populates_last_usage_from_the_assistant_messages_tokens(self): + self._start_server( + reply_deltas=["ok"], + reply_tokens={"input": 120, "output": 30, "reasoning": 5, "cache": {"read": 80, "write": 10}}, + ) + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + self.assertIsNone(chat.last_usage) + chat._call("hello") + + self.assertEqual( + chat.last_usage, + {"input": 120, "output": 30, "reasoning": 5, "cache_read": 80, "cache_write": 10}, + ) + + def test_last_usage_is_none_when_the_reply_carries_no_tokens_field(self): + self._start_server(reply_deltas=["ok"]) # reply_tokens defaults to None + chat = OpenCodeChat(f"http://127.0.0.1:{self.port}", model=None, timeout=10) + + chat._call("hello") + + self.assertIsNone(chat.last_usage) + class TestModelRefParsing(unittest.TestCase): def test_splits_on_first_slash_only(self): @@ -245,6 +273,28 @@ def test_session_idle_signals_completion(self): ) self.assertEqual(outcome, "idle") + def test_populates_the_usage_dict_when_the_assistant_message_carries_tokens(self): + usage = {} + payload = json.dumps({ + "type": "message.updated", + "properties": {"info": { + "id": "msg_1", "role": "assistant", "sessionID": "s1", + "tokens": {"input": 10, "output": 2, "reasoning": 0, "cache": {"read": 5, "write": 1}}, + }}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", set(), {}, usage) + self.assertIsNone(outcome) + self.assertEqual(usage, {"input": 10, "output": 2, "reasoning": 0, "cache_read": 5, "cache_write": 1}) + + def test_usage_param_is_optional_and_ignored_when_omitted(self): + # Existing call sites that predate the usage tracking must keep working. + payload = json.dumps({ + "type": "message.updated", + "properties": {"info": {"id": "msg_1", "role": "assistant", "sessionID": "s1", "tokens": {"input": 1}}}, + }) + outcome = OpenCodeChat._handle_event(payload, "s1", set(), {}) + self.assertIsNone(outcome) + if __name__ == "__main__": unittest.main() diff --git a/tests/ui/test_server_agent.py b/tests/ui/test_server_agent.py index 1d40e35a..6fe89ad4 100644 --- a/tests/ui/test_server_agent.py +++ b/tests/ui/test_server_agent.py @@ -129,6 +129,67 @@ def test_shutdown_stops_a_running_process(self): self.session.shutdown() self.assertIsNotNone(process.poll()) # exited + def test_ensure_drops_agents_md_into_a_fresh_workspace(self): + # This test process runs from an actual repo checkout, so AGENTS.md + # resolves for real via _repo_doc_path -- no mocking needed to prove + # ensure() actually reaches the workspace with it. + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + copied = os.path.join(self.tmp, "AGENTS.md") + self.assertTrue(os.path.isfile(copied)) + with open(copied, encoding="utf-8") as fh: + self.assertTrue(fh.read().strip()) + + def test_ensure_never_overwrites_a_workspace_own_agents_md(self): + # A workspace's own AGENTS.md might be the USER's project + # instructions -- ensure() must never clobber it with ours. + own_path = os.path.join(self.tmp, "AGENTS.md") + with open(own_path, "w", encoding="utf-8") as fh: + fh.write("this workspace's own instructions") + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + with open(own_path, encoding="utf-8") as fh: + self.assertEqual(fh.read(), "this workspace's own instructions") + + def test_ensure_copy_is_best_effort_when_no_source_agents_md_exists(self): + # No source to copy from (e.g. a stripped-down install) must not + # fail ensure() outright -- the agent server should still start. + with patch.object(ui_server, "_repo_doc_path", return_value=None): + with patch.object(ui_server, "_resolve_opencode_argv", return_value=_FAKE_ARGV): + result = self.session.ensure(self.tmp, "http://localhost:5173") + self.assertTrue(result["ok"], result) + self.assertFalse(os.path.isfile(os.path.join(self.tmp, "AGENTS.md"))) + + +class TestEnsureWorkspaceAgentsMd(unittest.TestCase): + """_ensure_workspace_agents_md directly -- no OpenCode process involved.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + + def test_copies_from_the_installed_package_location(self): + ui_server._ensure_workspace_agents_md(self.tmp) + target = os.path.join(self.tmp, "AGENTS.md") + self.assertTrue(os.path.isfile(target)) + with open(target, encoding="utf-8") as fh: + content = fh.read() + self.assertEqual(content, ui_server._read_repo_doc("AGENTS.md")) + + def test_is_a_no_op_when_the_workspace_already_has_one(self): + target = os.path.join(self.tmp, "AGENTS.md") + with open(target, "w", encoding="utf-8") as fh: + fh.write("mine") + ui_server._ensure_workspace_agents_md(self.tmp) + with open(target, encoding="utf-8") as fh: + self.assertEqual(fh.read(), "mine") + + def test_does_nothing_when_no_source_is_found(self): + with patch.object(ui_server, "_repo_doc_path", return_value=None): + ui_server._ensure_workspace_agents_md(self.tmp) + self.assertFalse(os.path.isfile(os.path.join(self.tmp, "AGENTS.md"))) + class TestCorsOriginVariants(unittest.TestCase): """The localhost <-> 127.0.0.1 expansion is the #1 way this feature goes @@ -197,6 +258,14 @@ def _post(self, path, origin=None): req.add_header("Origin", origin) return urllib.request.urlopen(req, timeout=10) + def _post_json(self, path, body): + req = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", method="POST", + data=json.dumps(body).encode("utf-8"), + headers={"Content-Type": "application/json"}, + ) + return urllib.request.urlopen(req, timeout=10) + class TestAgentServerEndpoint(_ServerTestCase): @@ -234,5 +303,230 @@ def test_falls_back_to_reconstructing_origin_from_host_header(self): self.assertTrue(data["ok"], data) +class TestAgentDocsEndpoint(_ServerTestCase): + """GET /agent-server/docs[?example=] -- AGENTS.md, plus + optionally one matching PyTorch usecase example, for the landing chat's + preset prompts to attach (see agentChat.ts's PRESET_PROMPTS). README.md + was dropped from this endpoint on purpose -- AGENTS.md alone carries the + weightslab integration pattern the presets need.""" + + def test_returns_agents_md_when_present_in_a_repo_checkout(self): + # This test process runs from an actual repo checkout, so it should + # resolve via _read_repo_doc. + with self._get("/agent-server/docs") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, {"AGENTS.md"}) + for f in data["files"]: + self.assertTrue(f["content"].strip()) + + def test_omits_the_doc_when_it_cannot_be_found_instead_of_erroring(self): + with patch.object(ui_server, "_read_repo_doc", return_value=None): + with self._get("/agent-server/docs") as r: + data = json.loads(r.read().decode()) + self.assertEqual(data["files"], []) + + def test_example_query_param_additionally_attaches_that_usecases_main_py(self): + with self._get("/agent-server/docs?example=wl-classification") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, {"AGENTS.md", "examples/PyTorch/wl-classification/main.py"}) + + def test_example_query_param_is_repeatable_for_multiple_usecases(self): + with self._get("/agent-server/docs?example=wl-detection&example=wl-segmentation") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, { + "AGENTS.md", + "examples/PyTorch/wl-detection/main.py", + "examples/PyTorch/wl-segmentation/main.py", + }) + + def test_example_query_param_is_ignored_when_not_a_known_usecase(self): + # Client-supplied -- must be checked against the allowlist, never + # trusted as a path component (e.g. "../../../etc/passwd"). + with self._get("/agent-server/docs?example=../../../etc/passwd") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertEqual(names, {"AGENTS.md"}) + + def test_no_example_query_param_means_no_example_file(self): + with self._get("/agent-server/docs") as r: + data = json.loads(r.read().decode()) + names = {f["name"] for f in data["files"]} + self.assertNotIn("examples/PyTorch/wl-classification/main.py", names) + + +class TestLoopRegistryUnit(unittest.TestCase): + """_LoopRegistry directly -- no HTTP server, no real OpenCode process. + Mocks the module-level _opencode_json_request/_opencode_send_and_collect/ + _opencode_get_messages functions the registry itself calls, so these + exercise its own eager-session-creation/locking/preamble-once logic in + isolation. A fresh registry per test, not the module-level singleton.""" + + def setUp(self): + self.registry = ui_server._LoopRegistry() + + def tearDown(self): + for job in self.registry._jobs.values(): + if job.timer is not None: + job.timer.cancel() + + def test_start_creates_the_session_eagerly_not_on_first_tick(self): + with patch.object(ui_server, "_opencode_session") as mock_session, \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "sess-1"}) as mock_req, \ + patch.object(ui_server, "_opencode_send_and_collect", return_value="ok"): + mock_session.ensure.return_value = {"ok": True, "url": "http://fake"} + result = self.registry.start("monitor training", 60.0, "/tmp/ws", "http://localhost:5173") + + self.assertTrue(result["ok"], result) + job = self.registry._jobs[result["id"]] + # Eager: already set by start() itself, not left for _fire's first + # tick (which runs in a background thread started right after). + self.assertEqual(job.session_id, "sess-1") + mock_req.assert_any_call("http://fake", "/session", method="POST", body=unittest.mock.ANY) + + def test_rejected_after_session_creation_deletes_the_orphaned_session(self): + def _ensure_and_fill_concurrently(*_args, **_kwargs): + # Simulates 3 OTHER starts winning the race while this one's + # ensure() call was in flight -- by the time start() re-checks + # under the lock, the cap has already been hit by them. + for i in range(3): + self.registry._jobs[str(i)] = ui_server._LoopJob(str(i), "p", 60.0, "/tmp") + return {"ok": True, "url": "http://fake"} + + with patch.object(ui_server, "_opencode_session") as mock_session, \ + patch.object(ui_server, "_opencode_json_request", return_value={"id": "sess-orphan"}) as mock_req: + mock_session.ensure.side_effect = _ensure_and_fill_concurrently + result = self.registry.start("monitor training", 60.0, "/tmp/ws", "http://localhost:5173") + + self.assertFalse(result["ok"]) + self.assertIn("already running", result["error"]) + mock_req.assert_any_call("http://fake", "/session/sess-orphan", method="DELETE") + + def test_fire_sends_the_preamble_once_then_plain_prompt_on_later_ticks(self): + job = ui_server._LoopJob("1", "check the loss", 60.0, "/tmp") + job.session_id, job.base_url = "sess-1", "http://fake" + self.registry._jobs["1"] = job + + sent_texts = [] + + def _fake_send(_base_url, _session_id, text, timeout=600.0): # noqa: ARG001 + sent_texts.append(text) + return "tick result" + + with patch.object(ui_server, "_opencode_send_and_collect", side_effect=_fake_send): + self.registry._fire("1", "http://fake") + self.assertTrue(job.preamble_sent) + self.assertIn("recurring monitoring agent", sent_texts[0]) + self.assertIn("check the loss", sent_texts[0]) + job.timer.cancel() + + self.registry._fire("1", "http://fake") + self.assertEqual(sent_texts[1], "check the loss") + job.timer.cancel() + + def test_fire_skips_the_tick_when_a_manual_message_is_in_flight(self): + job = ui_server._LoopJob("1", "check the loss", 60.0, "/tmp") + job.session_id, job.base_url = "sess-1", "http://fake" + job.preamble_sent = True # isolate the busy-skip path from preamble logic + self.registry._jobs["1"] = job + + job.lock.acquire() # simulate a manual send from the loop's chat tab + try: + with patch.object(ui_server, "_opencode_send_and_collect") as mock_send: + self.registry._fire("1", "http://fake") + mock_send.assert_not_called() + finally: + job.lock.release() + + self.assertIn("already in progress", job.last_error) + self.assertTrue(job.preamble_sent) # untouched -- the send never ran + self.assertIsNotNone(job.next_run_at) # rescheduled despite the skip + job.timer.cancel() + + def test_send_message_rejects_empty_text(self): + job = ui_server._LoopJob("1", "p", 60.0, "/tmp") + self.registry._jobs["1"] = job + result = self.registry.send_message("1", " ") + self.assertFalse(result["ok"]) + + def test_send_message_unknown_job(self): + result = self.registry.send_message("nope", "hi") + self.assertFalse(result["ok"]) + self.assertIn("No loop job", result["error"]) + + def test_send_message_success_updates_last_result(self): + job = ui_server._LoopJob("1", "p", 60.0, "/tmp") + job.session_id, job.base_url = "sess-1", "http://fake" + self.registry._jobs["1"] = job + with patch.object(ui_server, "_opencode_send_and_collect", return_value="the reply"): + result = self.registry.send_message("1", "hello") + self.assertTrue(result["ok"], result) + self.assertEqual(result["result"], "the reply") + self.assertEqual(job.last_result, "the reply") + + def test_get_messages_success(self): + job = ui_server._LoopJob("1", "p", 60.0, "/tmp") + job.session_id, job.base_url = "sess-1", "http://fake" + self.registry._jobs["1"] = job + canned = [{"info": {"role": "user"}, "parts": [{"type": "text", "text": "hi"}]}] + with patch.object(ui_server, "_opencode_get_messages", return_value=canned): + result = self.registry.get_messages("1") + self.assertTrue(result["ok"], result) + self.assertEqual(result["messages"], canned) + + def test_get_messages_unknown_job(self): + result = self.registry.get_messages("nope") + self.assertFalse(result["ok"]) + + +class TestLoopMessageEndpoints(_ServerTestCase): + """GET/POST /agent-server/loop//messages|message -- a loop tab's + scrollback + manual-send, proxied through the module-level _loop_registry + singleton to a job's own OpenCode session (mocked here; no real OpenCode + process). Loopback-gating itself mirrors the existing loop routes + (_stop_loop et al.), which have no dedicated test for the negative case + either -- a real test client is always loopback.""" + + def _seed_job(self): + job = ui_server._LoopJob("1", "check the loss", 60.0, self.tmp) + job.session_id, job.base_url = "sess-1", "http://fake" + ui_server._loop_registry._jobs["1"] = job + return job + + def tearDown(self): + ui_server._loop_registry._jobs.clear() + super().tearDown() + + def test_get_messages_returns_the_sessions_history(self): + self._seed_job() + canned = [{"info": {"role": "assistant"}, "parts": [{"type": "text", "text": "hi"}]}] + with patch.object(ui_server, "_opencode_get_messages", return_value=canned): + with self._get("/agent-server/loop/1/messages") as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + self.assertEqual(data["messages"], canned) + + def test_get_messages_404s_for_an_unknown_loop(self): + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._get("/agent-server/loop/999/messages") + self.assertEqual(ctx.exception.code, 404) + + def test_post_message_sends_and_returns_the_result(self): + self._seed_job() + with patch.object(ui_server, "_opencode_send_and_collect", return_value="the reply"): + with self._post_json("/agent-server/loop/1/message", {"text": "hello"}) as r: + data = json.loads(r.read().decode()) + self.assertTrue(data["ok"], data) + self.assertEqual(data["result"], "the reply") + + def test_post_message_rejects_empty_text(self): + self._seed_job() + with self.assertRaises(urllib.error.HTTPError) as ctx: + self._post_json("/agent-server/loop/1/message", {"text": ""}) + self.assertEqual(ctx.exception.code, 400) + + if __name__ == "__main__": unittest.main() diff --git a/weightslab/AGENTS.md b/weightslab/AGENTS.md new file mode 100644 index 00000000..57ed52cf --- /dev/null +++ b/weightslab/AGENTS.md @@ -0,0 +1,362 @@ +# WeightsLab — agent context for users & debugging + +Portable context for AI coding agents (and their humans) to **install, run, +integrate, and debug WeightsLab / Weights Studio** without reverse-engineering +the system first. Covers two repos: **weightslab** (Python backend — training +instrumentation, data ledger, gRPC service, shared proto) and **weights_studio** +(browser frontend that inspects/edits a *running* experiment). + +> File/line refs drift — verify against current source. Env var names/defaults +> are stable; authoritative reference is `weightslab/docs/configuration.rst`. + +--- + +## 0. Loading this guide into Claude Code + +- **Repo checkout:** committed as `AGENTS.md`; a gitignored `CLAUDE.md` copy at + the root gets auto-loaded every session. Nothing to do. +- **`pip install weightslab` only** (no checkout): absolute paths are fragile + across venvs/OS. Use a skill that locates the installed copy at runtime — + `~/.claude/skills/weightslab/SKILL.md`: + + ```yaml + --- + name: weightslab + description: Load the WeightsLab debugging & integration guide for weightslab/weights_studio problems (connection, TLS, env vars, training hangs, rendering, wl.* integration). + --- + !`python -c "import weightslab, os; print(open(os.path.join(os.path.dirname(weightslab.__file__), 'AGENTS.md')).read())"` + + Use the guide above to diagnose or implement the user's request. + ``` + + Requires the guide shipped as package data (`weightslab/weightslab/AGENTS.md` — see §7). +- **Quick-and-dirty:** copy this file to `~/.claude/WEIGHTSLAB.md`, `@`-import it + from `~/.claude/CLAUDE.md`. + +--- + +## 1. What it is, how the pieces connect + +A user wraps their PyTorch training script with WeightsLab so a running +experiment becomes inspectable/editable; Weights Studio is the UI. + +``` +Browser → weightslab start :8080 (grpc-web → grpc proxy) → Python gRPC servicer → training loop +``` + +- `weightslab start`: pure-Python HTTP server, serves the bundled SPA and + translates grpc-web↔gRPC. No Docker, no Envoy. Not running ⇒ no UI to load. +- gRPC servicer and training loop share **one process, different threads**, + coordinated by locks in `weightslab/weightslab/components/global_monitoring.py`. +- Proto is the single source of truth: `weightslab/weightslab/proto/experiment_service.proto`. + +--- + +## 2. Install & run + +```bash +pip install weightslab +``` + +```python +import weightslab as wl +# wrap objects so the studio can see/edit them (§3), then: +wl.serve(serving_grpc=True, serving_cli=True) +# ... training loop ... +wl.keep_serving() # keep process alive for the UI +``` + +```bash +weightslab start # http://localhost:8080 by default +``` + +For a new script, pick the closest match in +`weightslab/weightslab/examples/{PyTorch,Lightning,Ultralytics,Usecases}//main.py` +via the decision table in §3.9, then copy its `wl.*` calls — §3 documents that +whole API surface (reactive signals, group signals, the Ultralytics mixin, +etc. aren't in the `.rst` docs; the examples are the primary source). + +TLS/UI deploy details: `weightslab/docs/weights_studio.rst`. TLS is opt-in: +`weightslab se` once, then `weightslab start --certs`. + +--- + +## 3. The integration API (`import weightslab as wl`) + +How to wire a new training script correctly with no docs access — every verb +and kwarg here is real, taken from a shipping example under +`weightslab/weightslab/examples/` and checked against `weightslab/src.py`. + +### 3.1 Lifecycle + +```python +import weightslab as wl + +wl.watch_or_edit(..., flag=...) # register objects (§3.2) +wl.serve(serving_grpc=True, serving_cli=True) # background threads, same process +wl.start_training(timeout=3) # let UI/CLI attach before stepping +# ... training loop, guarded (§3.5) ... +wl.keep_serving() # block so the process/UI survives +``` + +- Register every object with `watch_or_edit` **before** `wl.serve`, using the parameter flag to define which object category it is. +- `timeout=0` skips the pre-start wait entirely. +- Skip `keep_serving()` for a script that should exit after writing a report + (`Usecases/*signals*` examples); include it otherwise. +- Tabular examples (`wl-fraud-detection`, `wl-ads-recommendation`) pass only + `serving_grpc=` to `serve` — no `serving_cli`. + +### 3.2 `wl.watch_or_edit(obj, flag=..., **kwargs)` + +Registers/wraps `obj` in the global ledger (`backend/ledgers.py`, +`GLOBAL_LEDGER`) and returns a live proxy. `flag` matches by substring +(case-insensitive). + +| flag | wraps | key kwargs | +|---|---|---| +| `"hyperparameters"` | plain `dict` | `defaults=parameters`, `poll_interval=1.0`, optional `name=` | +| `"model"` | `nn.Module` | `device=`; `compute_dependencies=False` (skip arch-op dependency graph when not editing architecture); `forced_model_wrapping=True` (Ultralytics only — load current object, not a checkpoint) | +| `"optimizer"` | `torch.optim.Optimizer` | none typically; build from the **watched** model's `.parameters()` | +| `"data"` | `Dataset` → tracked `DataLoader` | `loader_name=`, `batch_size=`, `shuffle=`, `is_training=`, `compute_hash=False`, `collate_fn=`, `preload_labels=`, `preload_metadata=` (both `True` for tabular — §3.4), `enable_h5_persistence=`, `num_workers=`; point-cloud/array data adds `array_autoload_arrays=`, `array_return_proxies=`, `array_use_cache=` | +| `"loss"` | `reduction="none"` criterion | `signal_name=`/`name=` (aliases), `log=True`, `per_sample=True` (one value/sample), `per_instance=True` (one value per `(sample_id, annotation_id)` — multi-box/mask samples); called as `criterion(preds_raw, targets, batch_ids=ids, preds=preds)` | +| `"metric"` | anything with `.compute()`/`.forward()` | same as `"loss"` | + +Registering `flag="loss"` auto-enrolls the signal for the background +loss-shape classifier (§3.6) unless overridden via `@wl.signal_classifier`. + +Objects need a `__name__` — set `obj.__name__ = "..."` manually if missing +(plain callables/custom loss modules). + +Hyperparameter proxy supports both `hp.get("lr")` and `hp["lr"]`, and stays +live (reflects later edits/re-registration). + +### 3.3 Per-sample / per-instance / grouped logging + +Watched loss/metric objects call `save_signals` internally on every +forward/compute — call these yourself only for derived values or anything not +from a watched object: + +- `wl.save_signals(batch_ids=ids, signals={...}, preds_raw=, targets=, preds=, log=True)` — + one value per sample id. `log=False` → stored as metadata, not a plotted signal. +- `wl.save_instance_signals(...)` — internal use by `per_instance=True`; rarely called directly. +- `wl.save_group_signals(signals={...}, group_ids=[...], origin="train_loader")` — + one row per group, for pairwise values (e.g. contrastive loss) that can't map + to a single sample. Needs a dataset that emits a `group_id` in its metadata + (`PyTorch/wl-generation`). +- `wl.trajectory_stats(values)` / `wl.classify_loss_shape(values)` — building + blocks behind the loss-shape tag (§3.6); call directly only for a custom classifier. + +### 3.4 `task_type` + +Set `self.task_type = "..."` on **both** model and dataset before +`watch_or_edit`. Confirmed values: + +| `task_type` | renders | set in | +|---|---|---| +| *(unset)* | classification (default) — also clustering, tabular, signal-tagging use cases | most examples | +| `"detection"` | 2D bounding boxes | `PyTorch/wl-detection/utils/{model,data}.py` | +| `"segmentation"` | instance/semantic masks | `PyTorch/wl-segmentation/utils/{model,data}.py` | +| `"detection_pointcloud"` | LiDAR point clouds, 2D or 3D (box column count disambiguates) | `Usecases/wl-{2d,3d}-lidar-detection/utils/{model,data}.py` | + +No `task_type="tabular"` exists — tabular rendering comes from the dataset +exposing feature values as sample **metadata** (`preload_labels=True, +preload_metadata=True`); mirror `PyTorch/wl-fraud-detection`, not the image +classification example. + +A dataset can implement `render_thumbnail_2d(...)` / `project_boxes_2d(...)` +for custom thumbnails — picked up automatically, no registration +(`Usecases/wl-3d-lidar-detection`, `CustomLidarDataset`). + +### 3.5 Guard contexts (required) + +```python +from weightslab import guard_training_context, guard_testing_context + +with guard_training_context: + ... # one training step +with guard_testing_context, torch.no_grad(): + ... # one eval step +``` + +Skip this and pause/resume and train/test stat separation break. Framework +variants: +- **Lightning:** wrap the body of `training_step`/`validation_step` — no manual loop. +- **Ultralytics mixin:** entered/exited manually (`guard.__enter__()`/`__exit__(None,None,None)`) + across `on_train_batch_start`/`_end` callback pairs (§3.9). + +Use `model.get_age()` (steps actually trained, survives checkpoints) for +step-based cadence, not a raw loop counter. + +### 3.6 Reactive signals & loss-shape classification + +```python +@wl.signal(name="sig/entropy", subscribe_to="loss_sample", batched=True) +def entropy(b): ... # b.logits, etc. → per-sample values + +@wl.signal(name="sig/hardness", inputs=["loss_sample", "sig/entropy"], batched=True) +def hardness(loss_vals, entropy_vals): ... +``` + +- `subscribe_to=` fires reactively when that signal saves (push); `inputs=[...]` + pulls named signals as args and can chain off other `@wl.signal` outputs. +- Define before `wl.serve()`/`wl.start_training()` — module scope or inline in `main()`. +- `@wl.signal_classifier(signal=)` overrides the built-in 6-way + loss-shape classifier (`monotonic/plateaued/Flat_high/high_variance/U_Shape/Spiked`) + for that signal, surfaced as categorical `tag:loss_shape` — no manual tagging + needed. Rebind at runtime: `wl.signal_classifier(signal=name)(fn)`. +- No custom classifier needed? Pass `loss_shape_signal=` to + `wl.write_dataframe(...)` (§3.7) to compute the built-in tag at dump time. + +### 3.7 Persisting & inspecting history + +- `wl.write_history()` / `wl.write_dataframe(path=, format="csv", columns=[...], loss_shape_signal=)` — + dump the ledger; `columns=` filters groups (e.g. `["signals","tags"]`). Call + periodically in long loops and once at the end. +- `wl.drain_signals()` — force-flush async signals before reading them back + (dataframe export, or a `GetDataSamples` call right after training). +- `wl.query_signal_history(...)` / `query_sample_history(...)` / `query_instance_history(...)` — + programmatic readback. + +### 3.8 Tagging & filtering samples + +`wl.tag_samples(...)`, `wl.register_categorical_tag(...)`/`set_categorical_tag(...)` +(multi-value, predefined categories — boolean tags are separate), +`wl.discard_samples(...)`, `wl.get_samples_by_tag(...)`, `wl.get_discarded_samples(...)`. +The automatic `tag:loss_shape` tag (§3.6) uses these same primitives. + +### 3.9 Which example to copy + +| Integrating... | Mirror | Notes | +|---|---|---| +| Plain PyTorch loop (classification) | `PyTorch/wl-classification` | Simplest pattern, manual loop in `main()`. | +| Detection (2D boxes) | `PyTorch/wl-detection` | `task_type="detection"`, `per_sample`/`per_instance`, custom `collate_fn`, decoded preds passed for overlays. | +| Segmentation | `PyTorch/wl-segmentation` | `task_type="segmentation"`, masks as list-of-tensors. | +| LiDAR detection (2D/3D) | `Usecases/wl-{2d,3d}-lidar-detection` | `task_type="detection_pointcloud"`; 3D adds `render_thumbnail_2d`. | +| Tabular / feature vectors | `PyTorch/wl-fraud-detection` | No `task_type`; `preload_labels=True, preload_metadata=True`; see §3.10 for a headless verification script. | +| Embedding / clustering | `PyTorch/wl-clustering` (+ `face/model.py`) | `watch_or_edit` calls live inside the model wrapper, not `main.py`; open-ended loop. | +| Paired/contrastive samples, group-level signals | `PyTorch/wl-generation` | `wl.save_group_signals`; dataset emits 2 rows per item via a `uids` metadata key. | +| Reactive signals / custom loss-shape tagging | `Usecases/wl-classification-signals_shape_classification`, `Usecases/ws-signals-mnist` | §3.6; the latter is the minimal variant with no custom classifier. | +| PyTorch Lightning | `Lightning/wl-classification` | Same `watch_or_edit` calls as plain PyTorch; guards wrap `training_step`/`validation_step` bodies; `Trainer(log_every_n_steps=0, enable_checkpointing=False, logger=False)`. | +| Ultralytics YOLO (detect/segment) | `Ultralytics/wl-detection` | Don't call `watch_or_edit` for model/optimizer/data/loss/metric — pass `trainer=WLAwareTrainer` (or `WLAwareSegmentationTrainer`) from `weightslab.integrations.ultralytics` to `YOLO(...).train(...)`. It wires everything via UL callbacks; you only watch the run config as `flag="hyperparameters"`. | + +### 3.10 Verifying an integration headlessly + +Watch model/optimizer/data/loss/metrics, `wl.serve(serving_grpc=True, grpc_port=...)`, +`wl.start_training()`, run real steps inside the guard contexts, +`wl.drain_signals()`, then assert on `wl.write_dataframe(..., format="csv")` +columns or a raw gRPC `GetDataSamples` call's `raw_data.type` (e.g. `"vector"` +for tabular). Plain script, not pytest — `python verify_integration.py` +(`PyTorch/wl-fraud-detection/verify_integration.py`). + +--- + +## 4. Configuration (environment variables) + +Authoritative reference: `weightslab/docs/configuration.rst`. High-signal ones: + +**Backend:** + +| Variable | Default | Why | +|---|---|---| +| `WEIGHTSLAB_LOG_LEVEL` | `INFO` | `DEBUG` for detail (`WATCHDOG` level sits between WARNING/ERROR). | +| `GRPC_BACKEND_HOST`/`PORT` | `0.0.0.0`/`50051` | Backend gRPC bind address. | +| `GRPC_TLS_ENABLED` | `0` | TLS on the gRPC socket; set with `weightslab start --certs`. | +| `GRPC_TLS_REQUIRE_CLIENT_AUTH` | `0` | mTLS; must match `--certs`. | +| `WEIGHTSLAB_CERTS_DIR` | `~/.weightslab-certs` | Cert lookup — single source of truth. | +| `GRPC_AUTH_TOKEN` | unset | Optional token auth on top of mTLS. | +| `GRPC_MAX_MESSAGE_BYTES` | `268435456` | Raise if large tensors/images fail to transfer. | +| `WEIGHTSLAB_DISABLE_WATCHDOGS` | `0` | Set `1` when breakpoint-debugging (§5). | +| `GRPC_WATCHDOG_STUCK_SECONDS` | `60` | Lock/RPC stuck threshold + lock-acquire timeout. | + +**Frontend — runtime `window.*` globals (injected at `weightslab start` time; restart+reload to apply):** + +| Variable | Default | Why | +|---|---|---| +| `WS_SERVER_HOST`/`PORT`/`PROTOCOL` | `localhost`/`8080`/`http` | How the browser reaches the server — #1 connection knob. | +| `WS_HISTOGRAM_MAX_BINS` | `512` | Metadata histogram bar cap. | +| `BB_THUMB_RENDER` | `10` | Max boxes per thumbnail, per overlay (GT/PRED independent). | +| `BB_MODAL_RENDER` | `100` | Max boxes per modal image, per overlay. | +| `ENABLE_PLOTS` | `1` | `0` removes plots board + Signals card. | +| `ENABLE_DATA_EXPLORATION` | `1` | `0` removes data grid + metadata panel. | +| `ENABLE_HYPERPARAMETERS_OPTIMIZATION` | `1` | `0` makes HP inputs read-only, stops HP poll. | +| `ENABLE_AGENT` | `1` | `0` removes agent chat bar. | +| `ENABLE_NOTEBOOK` | `1` | `0` removes the notebook (shared in-process kernel against the live experiment; persisted as `notebook.ipynb`; `>`-prefixed cells ask the agent for code). | + +`VITE_*` vars are build-time (need a frontend rebuild); `WS_*`/`BB_*`/`ENABLE_*` +are runtime (need only restart + reload). `ENABLE_*` default on; `0`/`false`/`no`/`off` disables. + +--- + +## 5. Troubleshooting + +**Sample grid empty / "failed to fetch" / gRPC errors.** Check in order: (1) +backend serving on `0.0.0.0:50051`; (2) `weightslab start` running, browser +reaches `:8080`; (3) TLS mismatch if using `--certs` — run `weightslab se` +first, export `WEIGHTSLAB_CERTS_DIR` (or drop TLS: omit `--certs`, `GRPC_TLS_ENABLED=0`). + +**Env var change not taking effect.** `VITE_*` → rebuild frontend. +`WS_*`/`BB_*`/`ENABLE_*` → restart `weightslab start` + reload tab. + +**Grid flashes empty on auto-refresh.** Refreshes now skip while a +`GetDataSamples` fetch is in flight (`isFetchInProgress()` in +`weights_studio/src/grid_data/gridDataManager.ts`) — confirm your build has this guard. + +**Detection overlays slow/cluttered.** Cap with `BB_THUMB_RENDER` / +`BB_MODAL_RENDER` (GT and PRED capped independently; render-only, no data dropped). + +**Training hangs; `RESOURCE_EXHAUSTED`; server "restarts".** A watchdog flags +locks/RPCs held past `GRPC_WATCHDOG_STUCK_SECONDS` (60s) and restarts the gRPC +server after repeated unhealthy polls. Debugging with breakpoints that +intentionally exceed this? Set `WEIGHTSLAB_DISABLE_WATCHDOGS=1`. +`RESOURCE_EXHAUSTED` = a handler couldn't get the lock in time — find what's holding it. + +**Pause/resume broken, or train/test stats mixed up.** Train/eval step isn't +wrapped in `guard_training_context`/`guard_testing_context` — see §3.5. + +**Large weights/images fail to transfer.** Raise `GRPC_MAX_MESSAGE_BYTES`. + +**Agent bar says unconfigured.** Backed by a local OpenCode server +(`OPENCODE_URL`, default `http://127.0.0.1:4096`), auto-started on first use. +`/init` from the UI (then `/model`, `/reset`). See `docs/agent.rst`, `docs/weights_studio.rst`. + +**Agent says it "cannot run code."** It has bash/read/write/edit tools rooted +at this workspace directory (the frontend sends no `tools` restriction) — a +model claiming otherwise is declining to call a tool it actually has, not +reporting a real limitation. Ask it directly and concretely: "use your bash +tool to run `python @@ -560,6 +794,7 @@ def _dataframe_section_html(stats: dict) -> str: WeightsLab Experiment Report +
@@ -577,6 +812,8 @@ def _dataframe_section_html(stats: dict) -> str: {signals_html}
+ {distributions_section_html} +

Loss-Shape Classification

{loss_shape_html} @@ -589,6 +826,31 @@ def _dataframe_section_html(stats: dict) -> str:
WeightsLab · generated automatically, review before sharing
+ + """ @@ -629,6 +891,7 @@ def render_report(context: dict, output_path, narrative: Optional[str] = None) - root_log_dir=html.escape(context.get("root_log_dir", "")), narrative=narrative_html, signals_html=signals_html, + distributions_section_html=_distributions_section_html(context.get("distributions") or []), loss_shape_html=_loss_shape_section_html(context.get("loss_shape_tags") or []), dataframe_html=_dataframe_section_html(context.get("dataframe") or {}), ) @@ -644,6 +907,29 @@ def default_report_path(root_log_dir) -> Path: return Path(root_log_dir) / "reports" / f"experiment_report_{stamp}.html" +def list_reports(root_log_dir) -> list: + """Every ``*.html`` report already written for this experiment, newest + first (by mtime) -- the same directory and ordering the Studio report + button's right-click dropdown uses (``weightslab/ui/server.py``'s + ``_list_experiment_reports``), duplicated here rather than imported since + that lives in the UI server module, not this LLM/agent-agnostic one.""" + reports_dir = Path(root_log_dir) / "reports" + if not reports_dir.is_dir(): + return [] + paths = [p for p in reports_dir.iterdir() if p.is_file() and p.suffix == ".html"] + paths.sort(key=lambda p: p.stat().st_mtime, reverse=True) + return paths + + +def latest_report_path(root_log_dir) -> Optional[Path]: + """The most recently written report for this experiment, or ``None`` if + none exists yet -- used to resolve "update the report"/"add X to the + report" requests to the file they mean, see ``generate_report``'s + ``update_existing``.""" + reports = list_reports(root_log_dir) + return reports[0] if reports else None + + def summarize_context_for_llm(context: dict) -> str: """The JSON summary of ``context`` that gets handed to an LLM for the narrative: everything except the base64 plot images (bytes an LLM can do @@ -658,6 +944,10 @@ def summarize_context_for_llm(context: dict) -> str: {k: v for k, v in entry.items() if k != "plot_b64"} for entry in context.get("signals", []) ], + "distributions": [ + {k: v for k, v in entry.items() if k != "plot_b64"} + for entry in context.get("distributions", []) + ], "loss_shape_tags": context.get("loss_shape_tags", []), "dataframe": context.get("dataframe", {}), }, indent=2, default=str) @@ -670,6 +960,8 @@ def generate_report( signals: Optional[list] = None, output_path=None, narrative_fn: Optional[Callable[[str], str]] = None, + distributions: Optional[list] = None, + update_existing: bool = False, ) -> dict: """Collect → narrate → render, in one call. The single implementation behind every user-facing entry point (the Studio report button and the @@ -684,11 +976,29 @@ def generate_report( configured, request timed out, ...) degrades the same way — a report with no written analysis rather than no report at all. - Returns ``{"path", "n_signals", "narrative", "narrative_error"}``. Failures - to gather data or write the file are NOT swallowed: they raise, because - there is no report to hand back. + ``distributions`` names columns to render as value-distribution + histograms in an extra Distributions section — opt-in, e.g. a follow-up + "add a histogram of train_loss to the report" — see + ``compute_distribution_entries``. + + ``update_existing`` (default ``False``, ignored if ``output_path`` is + given explicitly) — "update the report"/"add X to the report" should + overwrite the most recently written report for this experiment rather + than create a new timestamped one; "generate a report" with no reference + to an existing one should always create a fresh file. Resolved via + ``latest_report_path``; when there is nothing to update yet (a fresh + experiment, or its ``reports/`` directory was cleared), this falls back + to creating a new report exactly as if ``update_existing`` were ``False`` + — there is nothing wrong with "update" on a first-ever report. + + Returns ``{"path", "n_signals", "narrative", "narrative_error", + "updated_existing"}`` — the last key tells the caller whether an existing + file was overwritten (``True``) or a new one was created (``False``), so + it can phrase its reply accordingly. Failures to gather data or write the + file are NOT swallowed: they raise, because there is no report to hand + back. """ - context = collect_report_context(root_log_dir, logger_q, df, signals=signals) + context = collect_report_context(root_log_dir, logger_q, df, signals=signals, distributions=distributions) narrative = None narrative_error = None @@ -699,14 +1009,21 @@ def generate_report( narrative_error = str(exc) logger.warning("experiment report: narrative generation failed: %s", exc) - path = render_report( - context, - output_path if output_path is not None else default_report_path(root_log_dir), - narrative=narrative, - ) + resolved_output_path = output_path + updated_existing = False + if resolved_output_path is None and update_existing: + existing = latest_report_path(root_log_dir) + if existing is not None: + resolved_output_path = existing + updated_existing = True + if resolved_output_path is None: + resolved_output_path = default_report_path(root_log_dir) + + path = render_report(context, resolved_output_path, narrative=narrative) return { "path": path, "n_signals": len(context.get("signals") or []), "narrative": narrative, "narrative_error": narrative_error, + "updated_existing": updated_existing, } diff --git a/weightslab/src.py b/weightslab/src.py index 23fa802f..575f707a 100644 --- a/weightslab/src.py +++ b/weightslab/src.py @@ -1525,6 +1525,57 @@ def start_training(timeout: int = None) -> None: pause_ctrl.resume() # Ensure we're not paused if start_training is called after serve +def _register_pid_with_ui_server() -> None: + """Tell the ``weightslab start`` UI server that this training process + exists, so stopping that workspace (Ctrl+C, closing the terminal) stops + this process too. + + The UI server already kills PIDs registered with it on every termination + path (``weightslab/ui/server.py``'s ``_TrackedProcesses``) -- but until + now the ONLY thing that ever registered one was the agent remembering to + POST it by hand right after launching something detached. That second + step is skipped often enough to matter: the turn gets interrupted between + launch and registration, the model forgets on a relaunch, or the run was + started by hand in a terminal that never involved an agent at all. Each + miss leaves an orphaned python process holding the GPU and the gRPC port + after the UI is long gone. Registering from inside ``serve()`` makes it + unskippable for anything that serves. + + Entirely best-effort, and off the calling thread: a stale/unreachable + origin must never delay or break a training run. Only ever posts to + loopback (the endpoint refuses anything else anyway). + """ + origin = os.environ.get("WEIGHTSLAB_UI_ORIGIN") + if not origin: + return + + def _post() -> None: + try: + import json as _json + import ssl as _ssl + import urllib.request as _urlreq + + payload = _json.dumps({"pid": os.getpid()}).encode("utf-8") + request = _urlreq.Request( + f"{origin.rstrip('/')}/agent-server/track-process", + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + # `weightslab start --certs` serves the UI over HTTPS with its own + # self-signed material, which nothing in this process trusts. The + # payload is one integer going to loopback, so verification buys + # nothing here and would just turn tracking off whenever certs are on. + context = _ssl._create_unverified_context() if origin.startswith("https") else None + with _urlreq.urlopen(request, timeout=3.0, context=context): + pass + logger.debug(f"Registered PID {os.getpid()} with the WeightsLab UI at {origin}.") + except Exception as exc: + logger.debug(f"Could not register this process with the WeightsLab UI at {origin}: {exc}") + + threading.Thread(target=_post, name="wl-track-process", daemon=True).start() + + def serve(serving_cli: bool = True, serving_grpc: bool = True, spawn_cli_client: bool = False, serving_bore: bool = False, bore_port: int = None, allow_unconfigured: bool = True, **kwargs): @@ -1607,6 +1658,10 @@ def serve(serving_cli: bool = True, serving_grpc: bool = True, ) logger.warning(base_msg) + # Before anything that can block: whatever happens next, this process + # should not outlive the workspace that owns it. + _register_pid_with_ui_server() + # Embed a real Jupyter kernel (shares this process's live objects) so the # studio notebook panel — and any external Jupyter client — can attach to # it, unless we're already inside a notebook/Colab kernel ourselves (mode @@ -4895,6 +4950,7 @@ def ai_report_generation( output_path: str | None = None, root_log_dir: str | None = None, use_agent: bool = True, + distributions: list | None = None, ) -> str: """Generate the experiment health report and return the path written. @@ -4926,6 +4982,11 @@ def ai_report_generation( the same report minus the prose. With ``True`` but no agent available (none configured, or no experiment being served in this process), the report is still written — just without the analysis, never an error. + distributions : list of str, optional + Add a "Distributions" section with a value-distribution histogram for + each named column/signal (e.g. ``["train_loss"]``), computed over the + CURRENT per-sample dataframe rather than the aggregated training + curve. Omitted (default): no Distributions section at all. Raises ------ @@ -4940,10 +5001,12 @@ def ai_report_generation( wl.ai_report_generation() # everything, with analysis wl.ai_report_generation(signals=["train_loss"]) # one signal wl.ai_report_generation(use_agent=False) # no LLM call + wl.ai_report_generation(distributions=["train_loss"]) # + a histogram section """ return _ai_report_generation_result( signals=signals, output_path=output_path, root_log_dir=root_log_dir, use_agent=use_agent, + distributions=distributions, )["path"] @@ -4952,6 +5015,7 @@ def _ai_report_generation_result( output_path: str | None = None, root_log_dir: str | None = None, use_agent: bool = True, + distributions: list | None = None, ) -> dict: """:func:`ai_report_generation`'s implementation, keeping the full ``reporting.generate_report`` result dict (path + signal count + narrative @@ -5013,6 +5077,7 @@ def _ai_report_generation_result( result = reporting.generate_report( root_log_dir, logger_q, df, signals=signals, output_path=output_path, narrative_fn=narrative_fn, + distributions=distributions, ) if result["narrative_error"]: logger.warning( diff --git a/weightslab/trainer/services/agent/intent_prompt.py b/weightslab/trainer/services/agent/intent_prompt.py index b9e61154..8b07ac6a 100644 --- a/weightslab/trainer/services/agent/intent_prompt.py +++ b/weightslab/trainer/services/agent/intent_prompt.py @@ -44,7 +44,8 @@ | **Load weights** | "Load the model weights from step 500", "Roll back weights to step 500", "Load weights at step 500 from hash " | `action` (`action_name="load_weights"`, `action_params={{"step": 500}}`) | | **Tune hyperparameter** | "Set the batch size to 32", "Increase the learning rate by 10%", "Change the dumping model ratio to 15", "Change the evaluation ratio to 20" | `action` (`action_name="set_hyperparam"`; see rule 11) | | **Config question (READ-ONLY)** | "Show me the root log dir", "What is the batch size?", "Display the whole configuration", "Show the config" | `action` (`action_name="show_config"`; `action_params={{"param": ""}}` for one value, omit for the whole config) | -| **Experiment report** | "Generate a report", "How is this experiment going?", "Create a report on training progress", "Summarize the experiment" | `action` (`action_name="generate_experiment_report"`; optional `action_params={{"signals": ["", ...]}}`) | +| **Experiment report (new)** | "Generate a report", "How is this experiment going?", "Create a report on training progress", "Summarize the experiment" | `action` (`action_name="generate_experiment_report"`; optional `action_params={{"signals": ["", ...], "distributions": ["", ...]}}`) | +| **Experiment report (update)** | "Update the report", "Regenerate the report with...", "Add a histogram of X to the report", "Also include Y in it" | `action` (`action_name="generate_experiment_report"`; `action_params={{"update_existing": true, ...}}`) | | **History query** | "...that never had train loss below 0.5", "...whose loss was ever above 5", "min/max/mean loss OVER TRAINING" | `transform`/`keep` using `signal_history(...)` (see rule 10) | --- @@ -100,7 +101,12 @@ - `"load_weights"` — load ONLY model weights, optionally at a specific step: `action_params={{"step": }}` (and optionally `"hash": ""`; defaults to the current experiment). - `"set_hyperparam"` — change a training hyperparameter: `action_params={{"param": "", "op": "", "value": }}` (see rule 11). - `"show_config"` — READ-ONLY: display the experiment configuration. Omit `action_params` to dump the whole config, or pass `action_params={{"param": ""}}` to show a single value (e.g. `"root_log_dir"`, `"batch_size"`). Never modifies anything — use it for any "show/what is/display the config/setting" question. - - `"generate_experiment_report"` — READ-ONLY: build an HTML report (signal trajectory plots + health classification + dataset stats + a written analysis) summarizing how the current experiment is going, saved under the experiment's `reports/` directory. Optionally pass `action_params={{"signals": ["", ...]}}` to report on specific signals instead of the automatically-selected most-important ones. Never modifies anything. + - `"generate_experiment_report"` — READ-ONLY: build an HTML report (signal trajectory plots + health classification + dataset stats + a written analysis) summarizing how the current experiment is going, saved under the experiment's `reports/` directory. Optionally pass `action_params={{"signals": ["", ...]}}` to report on specific signals instead of the automatically-selected most-important ones, and/or `action_params={{"distributions": ["", ...]}}` to ADD a "Distributions" section with a value-distribution histogram for each named column/signal (e.g. "add a histogram of train_loss to the report" -> `distributions=["train_loss"]`). Both may be combined in one call. Never modifies anything. This is the ONLY correct way to satisfy "generate a report"/"how is this experiment going"/"summarize training progress" — never decompose a report request into several separate `analysis`/`data_analysis` steps to compute numbers yourself; the report already computes and plots everything in one deterministic pass. + + **`"update_existing"` (bool, default omitted = `false`)** — decides which FILE the report is written to: + - `false`/omitted → always write a brand-new, separately timestamped report. Use this for a request that doesn't reference an existing report: "Generate a report", "How is this experiment going?", "Create a report on training progress". + - `true` → overwrite the MOST RECENTLY generated report for this experiment instead of creating another one. Use this whenever the user's wording refers to a report that (as far as they're concerned) already exists: "Update the report", "Regenerate the report with...", "Add a histogram of X to the report", "Also include Y in it", "Redo the report but...". If nothing has been generated yet, the backend transparently falls back to creating a fresh one — you never need to check first. + - When you set `update_existing=true` for a request that only adds ONE new thing (e.g. "also add a histogram of val_loss"), check the conversation History for a prior `generate_experiment_report` call in this session and CARRY FORWARD its `signals`/`distributions` alongside the new one (e.g. previous `distributions=["train_loss"]` + this request -> `distributions=["train_loss", "val_loss"]`) — updating overwrites the whole file, so anything not re-listed would otherwise silently disappear from it. If you can't find a prior call to carry forward from, just pass what was asked for this turn. - `action_params`: Optional dict of parameters for the action (e.g. `{{"architecture": true}}`, `{{"hash": "abc123..."}}`, `{{"step": 500}}`, `{{"param": "batch_size", "op": "set", "value": 32}}`). --- @@ -860,7 +866,7 @@ **Ex51: Generate An Experiment Report (READ-ONLY)** User: "How is this experiment going? Generate a report." {{ - "reasoning": "Read-only request for a summary of experiment health. Use generate_experiment_report with no params so it auto-selects the most important signals.", + "reasoning": "Read-only request for a summary of experiment health. Use generate_experiment_report with no params so it auto-selects the most important signals. This one action already builds the whole report (plots, stats, analysis) -- it must NOT be split into separate analysis/data_analysis steps.", "primary_goal": "action", "steps": [ {{ @@ -871,6 +877,52 @@ }} +**Ex52: Add A Histogram Distribution To The Report (READ-ONLY, Update-In-Place)** +User: "Add a histogram distribution of train_loss to the report" +{{ + "reasoning": "The wording 'to THE report' refers to a report that (as far as the user is concerned) already exists, so this overwrites the most recent one (update_existing=true) rather than creating yet another timestamped file, in addition to requesting the Distributions section for this column.", + "primary_goal": "action", + "steps": [ + {{ + "kind": "action", + "action_name": "generate_experiment_report", + "action_params": {{ "update_existing": true, "distributions": ["train_loss"] }} + }} + ] +}} + + +**Ex52b: Second Follow-Up Add Carries The First One Forward (READ-ONLY, Update-In-Place)** +History (last turn): user asked "Add a histogram distribution of train_loss to the report" -> generate_experiment_report with action_params={{"update_existing": true, "distributions": ["train_loss"]}} +User: "Now also add one for val_loss" +{{ + "reasoning": "Another update to the SAME report. Overwriting replaces the whole file, so train_loss's distribution (added last turn, per History) must be re-listed alongside val_loss or it would silently disappear from the report.", + "primary_goal": "action", + "steps": [ + {{ + "kind": "action", + "action_name": "generate_experiment_report", + "action_params": {{ "update_existing": true, "distributions": ["train_loss", "val_loss"] }} + }} + ] +}} + + +**Ex53: Report On Specific Signals Plus A Distribution (READ-ONLY)** +User: "Generate a report on train_loss and val_loss, and include a distribution of val_loss" +{{ + "reasoning": "Combine both optional params on the same generate_experiment_report action: signals restricts the trajectory plots, distributions adds the histogram section.", + "primary_goal": "action", + "steps": [ + {{ + "kind": "action", + "action_name": "generate_experiment_report", + "action_params": {{ "signals": ["train_loss", "val_loss"], "distributions": ["val_loss"] }} + }} + ] +}} + + --- diff --git a/weightslab/trainer/services/agent/opencode_chat.py b/weightslab/trainer/services/agent/opencode_chat.py index de03bd13..e67581a4 100644 --- a/weightslab/trainer/services/agent/opencode_chat.py +++ b/weightslab/trainer/services/agent/opencode_chat.py @@ -317,20 +317,26 @@ def _ensure_model_resolved(self) -> None: this class's structured-JSON intent-parsing, since it isn't a text-reasoning model at all). - Resolution order mirrors weightslab/ui/server.py's own - _opencode_resolve_model (a /loop check-in hits the identical - problem, with no chat attached to inherit a model from): + Resolution order: 1. `GET /config`'s own `model` field -- the one the model picker writes back to opencode.json on every pick (opencodeClient.ts's setDefaultModel), so it's "whatever the user last actually chose", and survives across the browser, the CLI, and other - machines. - 2. `/config/providers`'s own `default` mapping -- what OpenCode - itself would otherwise have fallen back to for whichever - provider is configured. - 3. `_DEFAULT_MODEL` -- a free-tier model, used only when both of the - above come back empty (a fresh OpenCode install with no - provider credentials configured at all). + machines. The ONLY thing that counts as an actual choice here. + 2. `_DEFAULT_MODEL` -- a free-tier, always-available text/tool + model, used whenever step 1 comes back empty. + + `/config/providers`'s own `default` mapping used to be tried in + between (what OpenCode itself would otherwise fall back to for + whichever provider happens to be configured) -- dropped after this + was confirmed live to still resolve to an arbitrary, sometimes + non-text-reasoning model (an image-generation preview, the same + failure mode this method exists to avoid) whenever ANY provider had + credentials configured, even with no real model chosen -- silently + pre-empting _DEFAULT_MODEL every time. "At UI init, with nothing + explicitly chosen, land on the known-good free model" now means + exactly that, with no provider-reported default able to override it. + Resolved once and cached on self.model. An explicit model (model_is_explicit=True) is left alone -- deliberately chosen, not a placeholder to override. @@ -344,17 +350,6 @@ def _ensure_model_resolved(self) -> None: if isinstance(model_id, str) and "/" in model_id: self.model = model_id return - except Exception: # noqa: BLE001 - fall through to the next source - pass - try: - with self._request("/config/providers") as resp: - providers_data = json.loads(resp.read().decode("utf-8")) - defaults = (providers_data or {}).get("default") or {} - for provider in (providers_data or {}).get("providers") or []: - provider_id = provider.get("id") - if provider_id and defaults.get(provider_id): - self.model = f"{provider_id}/{defaults[provider_id]}" - return except Exception: # noqa: BLE001 - fall through to the hardcoded default pass self.model = _DEFAULT_MODEL diff --git a/weightslab/trainer/services/agent/report_prompt.py b/weightslab/trainer/services/agent/report_prompt.py index 588f99e2..d764ae2a 100644 --- a/weightslab/trainer/services/agent/report_prompt.py +++ b/weightslab/trainer/services/agent/report_prompt.py @@ -19,9 +19,9 @@ You are given already-computed statistics below and must not recompute, \ extrapolate, or second-guess them -- write an analysis grounded ONLY in what's \ -listed. The statistics have three parts, each bounded in size on purpose (so this \ -prompt stays small even for a dataset with millions of samples -- neither of us \ -ever sees a full per-sample dump): +listed. The statistics have up to four parts, each bounded in size on purpose (so \ +this prompt stays small even for a dataset with millions of samples -- neither of \ +us ever sees a full per-sample dump): - `signals`: per-signal AGGREGATE trajectory (the mean training curve) with its own shape label, plus (under `outliers`) a SMALL handful of specific sample_ids that @@ -33,6 +33,12 @@ any concerning label. Absent/empty means it hasn't been computed for this experiment -- say so briefly rather than assuming every sample is fine. - `dataframe`: sample counts, discard rate, splits, other tags. +- `distributions`: OPTIONAL -- only present when the user explicitly asked for a + histogram of a specific column (e.g. "add a histogram of train_loss"). Each entry + has `name`, `n`, `mean`, `std`, `min`, `max` for that column's CURRENT per-sample + values -- a snapshot spread, not a trajectory. Usually absent/empty; when present + you may mention it briefly (e.g. a wide spread, or min/max worth flagging), but + never invent one that isn't listed. Rules: - 3 to 6 sentences of plain prose. No JSON, no markdown headers/bullets, no code. diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py index f233e3ac..c7dc4c1b 100755 --- a/weightslab/trainer/services/data_service.py +++ b/weightslab/trainer/services/data_service.py @@ -2814,7 +2814,8 @@ def _agent_show_config(self, param=None) -> str: text = text[:max_len] + "\n... (truncated; ask for a specific key, e.g. 'show the root log dir')" return f"Configuration ({hp_name}):\n{text}" - def _agent_generate_experiment_report(self, signals=None) -> str: + def _agent_generate_experiment_report(self, signals=None, distributions=None, + update_existing=False) -> str: """Agent action (READ-ONLY): build an HTML experiment report -- signal trajectory plots + health classification + dataset stats + a written analysis -- saved under the experiment's ``reports/`` directory. @@ -2827,6 +2828,17 @@ def _agent_generate_experiment_report(self, signals=None) -> str: the collected numbers, never raw history, see report_prompt.py. A failed/unavailable LLM degrades to a report with no narrative rather than no report at all. + + ``distributions`` (optional column/signal names) adds a Distributions + section of value-distribution histograms -- e.g. "generate a report + with a histogram of train_loss" -> ``distributions=["train_loss"]``. + + ``update_existing`` (default ``False``) means "update the report" / + "add X to the report" -- overwrite the most recently generated report + for this experiment instead of writing a new timestamped one. A plain + "generate a report" (no reference to an existing one) should leave + this ``False`` and always get a fresh file; see intent_prompt.py's + `generate_experiment_report` guidance for when the LLM should set it. """ from weightslab.backend import ledgers from weightslab import reporting @@ -2852,12 +2864,14 @@ def _agent_generate_experiment_report(self, signals=None) -> str: result = reporting.generate_report( root_log_dir, logger_q, df, signals=signals, narrative_fn=self._agent.generate_report_narrative, + distributions=distributions, update_existing=bool(update_existing), ) except Exception as e: return f"Action: failed to generate report: {e}" suffix = "" if result["narrative"] else " (no written analysis -- agent LLM unavailable)" - return (f"Action: generated experiment report ({result['n_signals']} signal(s)) " + verb = "updated" if result["updated_existing"] else "generated" + return (f"Action: {verb} experiment report ({result['n_signals']} signal(s)) " f"at {result['path']}{suffix}") @staticmethod @@ -2974,7 +2988,10 @@ def _apply_agent_operation(self, df, func: str, params: dict) -> str: # classification + dataset stats + a written analysis). elif action_name in ("generate_experiment_report", "experiment_report", "create_report", "generate_report", "report"): - return self._agent_generate_experiment_report(signals=params.get("signals")) + return self._agent_generate_experiment_report( + signals=params.get("signals"), distributions=params.get("distributions"), + update_existing=bool(params.get("update_existing")), + ) return f"Action triggered: {action_name} (Not implemented)" diff --git a/weightslab/ui/server.py b/weightslab/ui/server.py index 763feb21..11d01f60 100644 --- a/weightslab/ui/server.py +++ b/weightslab/ui/server.py @@ -576,7 +576,7 @@ def ensure(self, workspace_dir: str, origin: Optional[str]) -> dict: OPENCODE_URL set by anyone -- still adopts that same server instead of spawning a second one for the identical experiment directory. """ - _ensure_workspace_agents_md(workspace_dir) + _ensure_workspace_agent_files(workspace_dir) with self._lock: if self._running_locked(): return {"ok": True, "url": self._url_locked(), "reused": True, @@ -722,24 +722,11 @@ def _read_repo_doc(filename: str) -> Optional[str]: return None -# Best-effort: drop AGENTS.md directly into a freshly-used OpenCode workspace -# so the agent can just `read AGENTS.md` itself -- it has file tools rooted -# right there -- rather than depending entirely on the landing chat happening -# to attach it via /agent-server/docs (which the experiment-bar-driven -# `/loop` agent never calls at all). Never overwrites an existing AGENTS.md -# already in the workspace: that may be the user's OWN project instructions, -# not ours to replace. -def _ensure_workspace_agents_md(workspace_dir: str) -> None: - target = Path(workspace_dir) / "AGENTS.md" - if target.exists(): - return - source = _repo_doc_path("AGENTS.md") - if source is None: - return - try: - shutil.copyfile(source, target) - except OSError: - pass +# Seeding AGENTS.md + opencode.json into the workspace lives in +# opencode_process, not here: the backend SDK agent reaches that module +# directly without ever going through this server, and both spawn paths have +# to leave a workspace in the same state. See its WORKSPACE_SEED_FILES. +_ensure_workspace_agent_files = opencode_process.ensure_workspace_agent_files # The MNIST preset's counterpart to _read_repo_doc: a complete, working @@ -827,6 +814,14 @@ def _read_example_main(usecase: str) -> Optional[str]: " status -- component NAMES + model age only, NO metric/hyperparam values\n" " agent query \"\" -- plain-English read OR edit (e.g. \"what is the last train loss\", " "\"discard samples where loss > 5\") -- the right tool for any metric/signal question, status never has that\n" + " agent query \"Generate an experiment report.\" -- the SAME command also builds a complete, " + "styled HTML report (signal plots, health classification, dataset stats, written analysis) in " + "ONE call when the task asks for a report/summary of how the run is going -- never assemble a " + "report yourself from several separate 'agent query' answers plus your own write tool, that " + "skips the plots/styling the backend already produces. On a LATER tick, \"agent query \\\"Update " + "the report with a histogram of val_loss.\\\"\" overwrites that SAME report file instead of " + "creating another one each check-in -- use that phrasing (\"update\"/\"add X to it\") once a " + "report already exists for this run, plain \"generate\" only for the first one\n" " pause / resume -- freeze/resume weight updates\n" " discard -- discard one sample by id\n\n" "Be fast: for a simple question, pipe one command in and answer from its " @@ -838,15 +833,42 @@ def _read_example_main(usecase: str) -> Optional[str]: "-WindowStyle Hidden -PassThru`, never a bare `python train.py` in the " "foreground. A foreground launch blocks THIS tool call -- and therefore " "this whole check-in, and the next one after it -- until the entire run " - "finishes, since this session processes one turn at a time.\n\n" + "finishes, since this session processes one turn at a time. This applies " + "just as much once you've decided the run genuinely needs to happen (e.g. " + "relaunching a crashed one) as it does anywhere else -- that is exactly " + "the case this rule is for, not an exception to it. If a bash call you " + "already made hasn't finished, the fix is to switch that launch to " + "Start-Process, never to raise its timeout and wait longer -- a bigger " + "timeout still blocks the same way, only for longer, and leaves you with " + "no PID to act on if you're later asked to stop it.\n\n" "A DETACHED process like that is NOT automatically stopped when this " "workspace is (Ctrl+C on `weightslab start`, or the process otherwise " "exiting) -- it has no OS-level relationship to this workspace's own " - "process tree once launched this way. Register its PID right after " - "launching it, so it is: `Invoke-RestMethod -Method Post -Uri " + "process tree once launched this way. A script that calls `wl.serve()` " + "registers itself and needs nothing from you; for anything else, register " + "its PID right after launching it, so it is stopped too: " + "`Invoke-RestMethod -Method Post -Uri " "'{origin}/agent-server/track-process' -ContentType 'application/json' " "-Body (@{{pid=$p.Id}} | ConvertTo-Json)`. Do this for anything you " - "relaunch, not just the first run.\n\n" + "relaunch, not just the first run -- and keep that PID in mind for the " + "rest of this job's life: it's what makes a later stop/kill instant " + "(`Stop-Process -Id -Force`) instead of having to go rediscover " + "which process is the right one.\n\n" + "NEVER stop/kill a process you did not yourself start in this session (or " + "the one crashed run you were explicitly asked to relaunch) without " + "asking the user first and naming the exact PID/command. Finding an " + "unfamiliar, duplicate-looking, or seemingly-stale process via ps/" + "tasklist/Get-CimInstance is not authorization to end it -- describe what " + "you found and wait for a go-ahead instead of cleaning it up yourself.\n\n" + "An exec/bash call blocks THIS check-in until it returns, with no " + "per-call timeout -- a command that can hang (a broad WMI/" + "Get-CimInstance scan, Stop-Process on something with stuck threads, " + "anything that can wait on a lock or a confirmation prompt) stalls the " + "whole check-in indefinitely, not just that one command. Prefer the " + "narrowest command that answers the question (a specific known PID over " + "a broad process scan); if something doesn't return in a few seconds, " + "don't just retry the same broad query -- report what you know and ask " + "rather than keep guessing.\n\n" "Your reply renders Markdown and LaTeX ($...$/$$...$$ or \\(...\\)/\\[...\\]) " "-- use real formulas for anything math-shaped instead of describing them in prose.\n\n" # weights_studio/src/agent/loopChatPane.ts looks for this EXACT trailing @@ -1716,6 +1738,11 @@ def _start_agent_server(self): origin = f"{scheme}://{host}" result = _opencode_session.ensure(workspace, origin) + if result.get("ok"): + # Told outright rather than left for the model to discover by + # trial and error -- see opencode_process.shell_platform_note() + # and _default_shell()'s docstring for why guessing was the bug. + result.update(opencode_process.shell_platform_note()) status = HTTPStatus.OK if result.get("ok") else HTTPStatus.INTERNAL_SERVER_ERROR self._send_json(status, result) @@ -2338,6 +2365,16 @@ def serve_ui( display_host = "localhost" if ui_host in ("0.0.0.0", "::", "") else ui_host url = f"{scheme}://{display_host}:{ui_port}" + # Where /agent-server/track-process lives, for anything launched from an + # environment descending from this process -- most importantly the + # OpenCode server (spawned by this server, so it inherits this) and + # therefore the agent's own shell, and therefore any training process the + # agent starts from it. weightslab.src.serve() reads this and registers + # ITSELF, so a run is tracked even when nobody POSTs its PID by hand; see + # _register_pid_with_ui_server there for why that mattered. Loopback + # rather than display_host because the endpoint only answers loopback. + os.environ["WEIGHTSLAB_UI_ORIGIN"] = f"{scheme}://127.0.0.1:{ui_port}" + sys.stdout.write( "\n" " WeightsLab UI is running:\n"