[sdk] support streaming

.speakeasy/speakeasy-modifications-overlay.yaml

@@ -8,8 +8,6 @@ info:
     before: ""
     type: speakeasy-modifications
 actions:
-  - target: $["paths"]["/textql.rpc.public.chat.ChatService/StreamChat"]
-    remove: true
   - target: $["paths"]["/textql.rpc.public.chat.ChatService/WatchChat"]
     remove: true
   - target: $["paths"]["/textql.rpc.public.agent.AgentService/StreamAgentStatus"]

STREAMING.md

@@ -0,0 +1,114 @@
+# Streaming (Connect-RPC bridge)
+
+The TextQL API exposes several **server-streaming** RPCs that have no HTTP/JSON
+shape in the OpenAPI spec, so they are not part of the Speakeasy-generated SDK
+surface. This package bridges them with [Connect-RPC](https://connectrpc.com)
+via `textql_sdk.streaming` — a hand-written module that talks the Connect
+protocol directly to the same gateway, authenticated with the same
+`tql_api_key`.
+
+> Prefer `watch_chat` for anything long-lived — it carries run lifecycle events
+> (`run_started`, `run_complete`, `run_error`) and heartbeats; `stream_chat` is
+> the run-scoped cell firehose for one-shot scripts.
+
+## Usage
+
+Configure the server and API key once on the `Textql` SDK; streaming inherits
+both. You never pass a server URL or deal with the `/rpc/public` mount:
+
+```python
+import asyncio, os
+from textql_sdk import Textql
+from textql_sdk.streaming import create_streaming_client
+from textql_sdk._connect.public.chat_pb2 import WatchChatRequest
+
+
+async def main():
+    sdk = Textql(api_key=os.environ["TEXTQL_API_KEY"])  # server_url optional
+    streaming = create_streaming_client(sdk)
+
+    async for event in streaming.chats.watch_chat(WatchChatRequest(chat_id=chat_id)):
+        payload = event.WhichOneof("payload")
+        if payload == "cell":
+            ...  # event.cell is a Cell
+        elif payload == "run_complete":
+            ...  # run finished
+        elif payload == "heartbeat":
+            pass  # keepalive, safe to ignore
+
+
+asyncio.run(main())
+```
+
+An on-prem/dev host set on the SDK is picked up automatically:
+
+```python
+sdk = Textql(api_key=..., server_url="https://your-host")  # e.g. from TEXTQL_SERVER_URL
+streaming = create_streaming_client(sdk)  # streams to https://your-host/rpc/public
+```
+
+Without an SDK instance, pass `api_key=` directly (`server_url` then defaults to
+the same server list the generated SDK uses, from the Speakeasy config):
+
+```python
+streaming = create_streaming_client(api_key=os.environ["TEXTQL_API_KEY"])
+```
+
+### Sync
+
+Use `create_streaming_client_sync` for a blocking iterator instead of an async
+one:
+
+```python
+from textql_sdk.streaming import create_streaming_client_sync
+
+streaming = create_streaming_client_sync(sdk)
+for update in streaming.agents.stream_agent_status(StreamAgentStatusRequest()):
+    print(update.agent_id, update.status)
+```
+
+## Streaming methods
+
+| Method | Emits |
+| --- | --- |
+| `chats.watch_chat(WatchChatRequest(chat_id=...))` | `WatchChatEvent` (opened, cell, run lifecycle, handoff, heartbeat) |
+| `chats.stream_chat(RunChatRequest(...))` | `Cell` per update while a run executes |
+| `agents.stream_agent_status(StreamAgentStatusRequest())` | `AgentStatusUpdate` for every visible agent run transition |
+| `apps.stream_app_activity(StreamAppActivityRequest(app_id=...))` | `AppActivityStreamEvent` (activity batches, presence, heartbeat) |
+| `dashboards.watch_dashboard_health(WatchDashboardHealthRequest(dashboard_id=...))` | `DashboardHealthEvent` on health transitions |
+| `playbooks.stream_template_data_status(StreamTemplateDataStatusRequest(...))` | `TemplateDataStatusUpdate` per template-data row |
+
+Request/response types live under `textql_sdk._connect.public.<service>_pb2`.
+
+For any other service in `textql_sdk._connect`, use the escape hatch:
+
+```python
+from textql_sdk.streaming import create_connect_client
+from textql_sdk._connect.public.feed_connect import FeedServiceClient
+
+feed = create_connect_client(FeedServiceClient, sdk)
+```
+
+## Notes
+
+- **Transport**: server-streaming rides on the `connect-python` runtime
+  (`pyqwest`). Client-streaming RPCs are not exercised by this bridge.
+- **Types**: streaming methods return protobuf message types (vendored under
+  `textql_sdk._connect`), which differ in shape from the Speakeasy-generated
+  Pydantic models for the same protos.
+- **Long-lived idle streams** may be closed by intermediary proxies if nothing
+  is sent for a while; `watch_chat` and `stream_app_activity` send periodic
+  heartbeats, the others emit only on activity. Wrap consumption in a reconnect
+  loop for anything long-running.
+
+## Regenerating the vendored types
+
+`src/textql_sdk/_connect` is generated from the platform protos (not by
+Speakeasy — it survives `speakeasy run`):
+
+```bash
+DEMO2_DIR=/path/to/demo2 ./scripts/generate-connect.sh
+```
+
+The buf plugin versions in that script are pinned to match the `connect-python`
+runtime in `pyproject.toml`; bump them together.

examples/watch_chat.py

@@ -0,0 +1,56 @@
+"""Watch a chat over the Connect-RPC streaming bridge.
+
+Server-streaming RPCs aren't part of the generated REST SDK, so they live in
+`textql_sdk.streaming`. `watch_chat` emits the full run lifecycle plus every
+cell as it's produced. See STREAMING.md for the other streaming methods.
+
+    uv run python examples/watch_chat.py <chat_id>
+"""
+
+import asyncio
+import os
+import sys
+
+from dotenv import load_dotenv
+from connectrpc.errors import ConnectError
+
+from textql_sdk import Textql
+from textql_sdk.streaming import create_streaming_client
+from textql_sdk._connect.public.chat_pb2 import WatchChatRequest
+
+
+async def main() -> None:
+    load_dotenv()
+    if len(sys.argv) < 2:
+        raise SystemExit("usage: python examples/watch_chat.py <chat_id>")
+    chat_id = sys.argv[1]
+
+    # Configure the SDK once; streaming inherits its server + API key.
+    # Set TEXTQL_SERVER_URL for on-prem/dev; it defaults to the cloud server.
+    sdk = Textql(
+        api_key=os.environ["TEXTQL_API_KEY"],
+        server_url=os.environ.get("TEXTQL_SERVER_URL"),
+    )
+    streaming = create_streaming_client(sdk)
+
+    try:
+        async for event in streaming.chats.watch_chat(WatchChatRequest(chat_id=chat_id)):
+            kind = event.WhichOneof("payload")
+            if kind == "cell":
+                print(f"cell: {event.cell.id}")
+            elif kind == "run_started":
+                print("run started")
+            elif kind == "run_complete":
+                print("run complete")
+                break
+            elif kind == "run_error":
+                print(f"run error: {event.run_error}")
+                break
+            elif kind == "heartbeat":
+                pass  # keepalive
+    except ConnectError as e:
+        raise RuntimeError(f"watch_chat failed: {e.code}") from e
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

pyproject.toml

@@ -9,6 +9,8 @@ dependencies = [
     "httpcore >=1.0.9",
     "httpx >=0.28.1",
     "pydantic >=2.11.2,<2.13",
+    "connect-python >=0.9.0",
+    "protobuf >=6.31,<8",
 ]
 urls.repository = "https://github.com/TextQLLabs/textql-python-v3.git"
 license = { text = "Apache-2.0" }
@@ -24,7 +26,7 @@ dev = [
 where = ["src"]
 
 [tool.setuptools.package-data]
-"*" = ["py.typed"]
+"*" = ["py.typed", "*.pyi"]
 
 [build-system]
 requires = ["setuptools>=80", "wheel"]

scripts/generate-connect.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Regenerates the vendored Connect-RPC client under src/textql_sdk/_connect from
+# the platform protos. This is NOT produced by Speakeasy — it survives
+# `speakeasy run` (Speakeasy only owns the files it generates). Mirrors the TS
+# SDK's scripts/generate-connect.sh.
+#
+#   DEMO2_DIR=/path/to/demo2 ./scripts/generate-connect.sh
+#
+# Plugin versions are pinned to match the `connect-python` runtime declared in
+# pyproject.toml. Bump them together.
+
+SDK_DIR="$(cd "$(dirname "$0")/.." && pwd)"
+DEMO2_DIR="${DEMO2_DIR:-$SDK_DIR/../demo2}"
+PROTO_DIR="$DEMO2_DIR/proto/api"
+OUT_DIR="$SDK_DIR/src/textql_sdk/_connect"
+
+CONNECT_PLUGIN_VERSION="v0.9.0"   # keep in sync with connect-python in pyproject.toml
+PROTOBUF_PLUGIN_VERSION="v31.1"   # keep in sync with the protobuf runtime range
+
+TEMPLATE_DIR="$(mktemp -d)"
+trap 'rm -rf "$TEMPLATE_DIR"' EXIT
+TEMPLATE="$TEMPLATE_DIR/buf.gen.yaml"
+cat > "$TEMPLATE" <<EOF
+version: v2
+clean: true
+plugins:
+  - remote: buf.build/protocolbuffers/python:$PROTOBUF_PLUGIN_VERSION
+    out: $OUT_DIR
+  - remote: buf.build/protocolbuffers/pyi:$PROTOBUF_PLUGIN_VERSION
+    out: $OUT_DIR
+  - remote: buf.build/connectrpc/python:$CONNECT_PLUGIN_VERSION
+    out: $OUT_DIR
+EOF
+
+cd "$PROTO_DIR"
+buf generate --include-imports --template "$TEMPLATE"
+
+# protoc/connect emit absolute imports (`import public.chat_pb2`), which only
+# resolve if the output dir is on sys.path. Rewrite them to package-relative so
+# _connect works as a plain subpackage of textql_sdk. google.protobuf is left
+# absolute so the well-known types resolve from the installed protobuf runtime.
+FDS="$TEMPLATE_DIR/fds.binpb"
+buf build --as-file-descriptor-set -o "$FDS"
+# Run isolated (uvx, not `uv run`): protoletariat pins protobuf<6, which clashes
+# with the SDK's runtime protobuf. It's a build-time text rewriter, never imported.
+uvx --from protoletariat protol \
+  --python-out "$OUT_DIR" --in-place --create-package --exclude-google-imports \
+  -s _pb2.py -s _pb2.pyi -s _connect.py \
+  raw "$FDS"
+
+# protoletariat doesn't rewrite the `import x.y as z` form that connect-python
+# emits in *_connect.py. Every service lives under public/, so its sibling
+# message imports become relative; google.protobuf stays absolute (runtime).
+find "$OUT_DIR/public" -name '*_connect.py' -exec \
+  sed -i.bak -E 's/^import public\.([A-Za-z0-9_]+) as /from . import \1 as /' {} +
+find "$OUT_DIR" -name '*.bak' -delete
+
+# Trim internal-only proto trees (not part of the public SDK surface) and the
+# vendored google/protobuf well-known types (we use the runtime's — a second
+# copy would double-register in the descriptor pool).
+rm -rf "$OUT_DIR/platform" "$OUT_DIR/demo" "$OUT_DIR/google/protobuf"
+
+echo "regenerated $OUT_DIR"

src/textql_sdk/_connect/__init__.pyi

@@ -0,0 +1,10 @@
+from . import auth_pb2
+from . import buf
+from . import compute_pb2
+from . import demo
+from . import google
+from . import paradigm_params_pb2
+from . import platform
+from . import powerbi_selection_pb2
+from . import public
+from . import textable_pb2