LLM Proxy Features Added

canyonos_core/llm_proxy/README.md

@@ -21,7 +21,8 @@ your app (unchanged)         localhost:8080                 real upstream
 - **OpenAI / Anthropic** — straight HTTP reverse-proxy: rewrite host, swap in the
   real key, forward with `requests`, return the response.
 - **Bedrock** — re-issued through the proxy's own `boto3` client (handles SigV4
-  signing + URL-encoding correctly). Only `invoke` is wired up.
+  signing + URL-encoding correctly). `invoke`, `converse`, `converse-stream`,
+  and `invoke-with-response-stream` are all wired up.
 
 ## Run
 
@@ -78,33 +79,47 @@ boto3.client("bedrock-runtime").invoke_model(
 
 ## Telemetry & Metrics
 
-**Automatic telemetry is currently Bedrock-only.** The proxy captures:
+The proxy captures, per response:
 - Model ID
 - Input/output/total token counts
-- Cache tokens (read & write)
+- Cache tokens (read & write, where the provider reports them)
 - Error status
 
-Telemetry is automatically written to Redis under `future:<future_id>` keys.
+Telemetry is written to Redis under `future:<future_id>` keys, keyed off an
+`X-Canyonos-Future-ID` request header.
 
-### How it works (Bedrock only)
+### Usage extraction coverage (`hooks.py::Hooks._extract_usage`)
 
-1. **Auto-injection:** boto3 hook (`proxy.py`) injects `X-Canyonos-Future-ID` header from thread-local context
-2. **Token extraction:** `hooks.py` parses response `usage` field
-3. **Redis write:** All metrics written to `future:<future_id>` hash
+| Provider / op | Usage schema used | Status |
+| --- | --- | --- |
+| Bedrock `converse` / `converse-stream` | Bedrock's own camelCase (`inputTokens`, ...) | works for any model |
+| Bedrock `invoke` / `invoke-with-response-stream`, `anthropic.*` model | Anthropic's native (`input_tokens`, ...) | works |
+| Bedrock `invoke` / `invoke-with-response-stream`, other model families | model-specific, unknown | no usage (schema not implemented yet) |
+| Direct Anthropic API (`/anthropic/...`) | Anthropic's native | works |
+| Direct OpenAI API (`/openai/...`) | OpenAI's native (`prompt_tokens`, ...) | works |
 
-### Why Bedrock-only?
+### How it works
+
+1. **Auto-injection:** boto3 hook (`proxy.py`) injects `X-Canyonos-Future-ID` header from thread-local context -- **Bedrock (boto3) only**; the OpenAI/Anthropic SDKs don't fire this hook, so callers using those SDKs directly won't get the header auto-attached.
+2. **Usage extraction:** `hooks.py`'s `_extract_usage` parses the response `usage` field with the schema matching that provider/op/model (table above) -- this part works for all three providers whenever the header is present.
+3. **Redis write:** All metrics written to `future:<future_id>` hash.
+
+So in practice, automatic end-to-end telemetry (header injection + extraction) is Bedrock-only for now; OpenAI/Anthropic usage extraction works, but nothing auto-attaches `X-Canyonos-Future-ID` for those SDKs yet.
+
+### Why is header auto-injection Bedrock-only?
 
 OpenAI and Anthropic use their own Python SDKs (`openai`, `anthropic`), not boto3.
-The boto3 event hook doesn't fire for non-AWS SDKs. To add telemetry for those:
+The boto3 event hook doesn't fire for non-AWS SDKs. To add auto-injection for those:
 - Would need separate hooks in each SDK's HTTP client
-- Or callers would need to use proxy directly (not through SDKs)
-
-The proxy *forwards* OpenAI/Anthropic requests and *can* extract tokens, but doesn't
-automatically inject headers or write telemetry.
+- Or callers would need to attach `X-Canyonos-Future-ID` themselves
 
 ## Limitations
 
-- **No streaming.** `stream=True` / `invoke-with-response-stream` are not handled.
+- **Bedrock `converse-stream` and `invoke-with-response-stream` only.**
+  OpenAI/Anthropic `stream=True` is still not handled and remains fully
+  buffered. `invoke-with-response-stream` also has no usage/token telemetry
+  regardless of model (unlike `converse-stream`, it has no metadata event to
+  read usage from -- see the usage extraction coverage table above).
 - **Bedrock error bodies are reconstructed**, not passed through byte-for-byte
   (boto3 raises on 4xx/5xx; we rebuild a JSON body with the real status +
   message). OpenAI/Anthropic errors pass through unchanged.

canyonos_core/llm_proxy/core.py

@@ -6,7 +6,7 @@
 import time
 from typing import Optional
 
-from flask import Response
+from flask import Response, stream_with_context
 
 from canyonos_core.llm_proxy.hooks import Ctx
 
@@ -53,5 +53,20 @@ def proxy_request(provider, subpath, flask_request):
     else:
         pr = provider.forward(flask_request, subpath, body)
 
+    if pr.stream is not None:
+        # Streamed responses: relay chunks as they arrive, and only take telemetry when the whole response is done.
+        def relay():
+            try:
+                yield from pr.stream
+            finally:
+                hooks.on_response(ctx, pr)
+
+        return Response(
+            stream_with_context(relay()),
+            status=pr.status,
+            headers=pr.headers,
+            direct_passthrough=True,
+        )
+
     hooks.on_response(ctx, pr)
     return Response(pr.content, status=pr.status, headers=pr.headers)

canyonos_core/llm_proxy/hooks.py

@@ -1,9 +1,10 @@
 """The metrics seam.
 
-Every proxied call passes through ``on_request`` / ``on_response``. Today these
-only log. Token accounting lands here later: because the whole response is
-buffered, usage extraction is a one-liner, e.g. ``resp.json().get("usage")`` for
-OpenAI/Anthropic (Bedrock's usage lives in its per-model response body).
+Every proxied call passes through ``on_request`` / ``on_response``, which logs
+and (when Redis is configured) extracts token usage per provider/op/model --
+see ``Hooks._extract_usage``. Bedrock invoke's usage schema is only known for
+anthropic.* models today; other model families and OpenAI/Anthropic streaming
+remain unhandled.
 """
 
 from __future__ import annotations
@@ -76,15 +77,17 @@ def on_request(self, ctx: Ctx) -> None:
         )
 
     def on_response(self, ctx: Ctx, resp: Any) -> None:
-        # Extract tokens for Bedrock
-        usage = None
-        if ctx.provider == "bedrock":
-            usage = self._extract_bedrock_tokens(resp)
-
+        usage = self._extract_usage(ctx, resp)
+        is_stream = getattr(resp, "stream", None) is not None
+
+        status = getattr(resp, "status", "?")
+        if is_stream and getattr(resp, "stream_error", False):
+            status = f"{status} (stream-error)"
+
         log.info(
             "← %s %s /%s -> %s in %.0fms | %s",
             ctx.provider, ctx.method, ctx.subpath,
-            getattr(resp, "status", "?"), ctx.elapsed_ms(),
+            status, ctx.elapsed_ms(),
             usage or "no usage"
         )
 
@@ -98,8 +101,10 @@ def on_response(self, ctx: Ctx, resp: Any) -> None:
                     # Extract model ID
                     model_id = self._extract_model_id(ctx)
 
-                    is_error = resp.status >= 400
-
+                    is_error = resp.status >= 400 or (
+                        is_stream and getattr(resp, "stream_error", False)
+                    )
+
                     # Build telemetry data
                     data = {
                         "model": model_id,
@@ -125,35 +130,90 @@ def _extract_model_id(self, ctx: Ctx) -> str:
         """Extract model ID from context or subpath."""
         if ctx.model:
             return ctx.model
-
-        # For Bedrock: subpath is "model/<modelId>/operation"
-        # Use rpartition to peel operation off the right (same as provider logic)
-        if ctx.provider == "bedrock" and ctx.subpath.startswith("model/"):
-            model_id, sep, op = ctx.subpath[len("model/"):].rpartition("/")
-            if sep:  # Found a separator
+        if ctx.provider == "bedrock":
+            model_id, _op = self._bedrock_model_and_op(ctx.subpath)
+            if model_id:
                 return model_id
-
         return "unknown"
-
-    def _extract_bedrock_tokens(self, resp: Any) -> Optional[TokenUsage]:
-        """Extract token usage from Bedrock response. It requires diff logic from OpenAI/Anthropic"""
-        if resp.status != 200:
+
+    @staticmethod
+    def _bedrock_model_and_op(subpath: str):
+        """Split a Bedrock subpath ("model/<modelId>/<op>") into (model_id, op); (None, None) if unrecognized."""
+        if not subpath.startswith("model/"):
+            return None, None
+        model_id, sep, op = subpath[len("model/"):].rpartition("/")
+        return (model_id, op) if sep else (None, None)
+
+    def _extract_usage(self, ctx: Ctx, resp: Any) -> Optional[TokenUsage]:
+        """Dispatch to the right usage schema for this provider/op/model."""
+        is_stream = getattr(resp, "stream", None) is not None
+
+        if ctx.provider == "bedrock":
+            if is_stream:
+                return self._usage_from_dict(getattr(resp, "stream_usage", None))
+            model_id, op = self._bedrock_model_and_op(ctx.subpath)
+            # invoke's body is model-native; only anthropic.*'s schema is known so far.
+            if op == "invoke" and (model_id or "").startswith("anthropic."):
+                return self._extract_json_usage(resp, self._usage_from_anthropic_dict)
+            return self._extract_json_usage(resp, self._usage_from_dict)
+
+        if ctx.provider == "anthropic":
+            return self._extract_json_usage(resp, self._usage_from_anthropic_dict)
+
+        if ctx.provider == "openai":
+            return self._extract_json_usage(resp, self._usage_from_openai_dict)
+
+        return None
+
+    @staticmethod
+    def _extract_json_usage(resp: Any, parser) -> Optional[TokenUsage]:
+        """Parse resp.content as JSON and hand its "usage" key to `parser`."""
+        if getattr(resp, "status", None) != 200:
             return None
-
         try:
             data = json.loads(resp.content.decode("utf-8"))
-            usage = data.get("usage", {})
-            if usage:
-                return TokenUsage(
-                    input_tokens=usage.get("inputTokens", 0),
-                    output_tokens=usage.get("outputTokens", 0),
-                    total_tokens=usage.get("totalTokens", 0),
-                    input_cache_tokens=usage.get("cacheReadInputTokens", 0),
-                    input_cache_write_tokens=usage.get("cacheCreationInputTokens", 0),
-                )
-        except:
-            pass
-        return None
+        except Exception:
+            return None
+        return parser(data.get("usage"))
+
+    @staticmethod
+    def _usage_from_dict(usage: Optional[Dict[str, Any]]) -> Optional[TokenUsage]:
+        """Bedrock Converse's usage schema (camelCase), used for both converse and converse-stream."""
+        if not usage:
+            return None
+        return TokenUsage(
+            input_tokens=usage.get("inputTokens", 0),
+            output_tokens=usage.get("outputTokens", 0),
+            total_tokens=usage.get("totalTokens", 0),
+            input_cache_tokens=usage.get("cacheReadInputTokens", 0),
+            input_cache_write_tokens=usage.get("cacheCreationInputTokens", 0),
+        )
+
+    @staticmethod
+    def _usage_from_anthropic_dict(usage: Optional[Dict[str, Any]]) -> Optional[TokenUsage]:
+        """Anthropic's native usage schema (snake_case, no total field); shared by direct Anthropic API calls and Bedrock invoke for anthropic.* models, since Bedrock returns Anthropic's own response body unchanged for that op."""
+        if not usage:
+            return None
+        input_tokens = usage.get("input_tokens", 0)
+        output_tokens = usage.get("output_tokens", 0)
+        return TokenUsage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=input_tokens + output_tokens,
+            input_cache_tokens=usage.get("cache_read_input_tokens", 0),
+            input_cache_write_tokens=usage.get("cache_creation_input_tokens", 0),
+        )
+
+    @staticmethod
+    def _usage_from_openai_dict(usage: Optional[Dict[str, Any]]) -> Optional[TokenUsage]:
+        """OpenAI's native usage schema."""
+        if not usage:
+            return None
+        return TokenUsage(
+            input_tokens=usage.get("prompt_tokens", 0),
+            output_tokens=usage.get("completion_tokens", 0),
+            total_tokens=usage.get("total_tokens", 0),
+        )
 
 
 hooks = Hooks()

canyonos_core/llm_proxy/providers/base.py

@@ -10,7 +10,7 @@
 
 import json
 from dataclasses import dataclass, field
-from typing import Dict, Iterable, List, Tuple
+from typing import Dict, Iterable, Iterator, List, Optional, Tuple
 
 import requests
 
@@ -43,7 +43,12 @@ class UpstreamRequest:
 class ProxyResponse:
     status: int
     headers: List[Tuple[str, str]]
-    content: bytes
+    content: bytes = b""
+    # Set instead of ``content`` for streamed responses; core.proxy_request streams these chunks through directly.
+    stream: Optional[Iterator[bytes]] = None
+    # Filled in by the stream generator once usage is known (only available after the trailing "metadata" event).
+    stream_usage: Optional[dict] = None
+    stream_error: bool = False
 
     def json(self):
         return json.loads(self.content.decode("utf-8"))