From 3c43b67869b35d45d582e28eb77f657817560dff Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Tue, 1 Sep 2026 17:09:43 -0700 Subject: [PATCH] Feature: Complete LLM telemetry system with distributed proxy architecture Implements automatic LLM telemetry capture with zero agent code changes. ## Architecture - Distributed proxy: One per container, auto-starts on port 8081 - boto3 hook: Injects X-Ventis-Future-Id header from thread-local context - Token extraction: Automatic parsing of Bedrock responses - Redis persistence: All 7 metrics written to future: keys ## Key Components - ventis/llm_proxy/: Complete proxy package (app, hooks, providers) - ventis/controller/local_controller.py: Auto-starts proxy subprocess - ventis/controller/utils/process_supervisor.py: Subprocess management - ventis/stub_generator.py: Copies llm_proxy with full package structure ## Metrics Captured (Bedrock only) - model: Full model ID from request path - input_token_count, output_token_count, token_count - input_cache_tokens (cache reads) - input_cache_write_tokens (cache writes) - errors: HTTP status >= 400 ## Key Fixes - Package structure: llm_proxy copied as ventis/llm_proxy/ to preserve imports - Infinite loop prevention: Proxy's boto3 client uses explicit AWS endpoint - Hooks initialization: Import hooks inside proxy_request() to get configured instance - Flask header normalization: Handle X-Ventis-Future-Id (Title-Case) - Dependencies: Added flask and requests to BASE_AGENT_REQUIREMENTS ## Agent Changes Agents use standard boto3 - zero telemetry code needed: - examples/portfolio/agents/advisor_agent.py: Removed ventis.llm imports - examples/portfolio/agents/intent_agent.py: Removed ventis.llm imports - examples/text2sql/agents/vllm_agent.py: Removed ventis.llm imports ## Removed - ventis/llm/: Old bedrock wrapper (deprecated in favor of proxy) - Planning docs: Consolidated into llm_proxy/README.md ## Testing Verified end-to-end on EC2: - LLM calls succeed through proxy - Token extraction works (inputTokens, outputTokens, cache tokens) - Redis writes confirmed with all 7 fields - Environment: boto3 + AWS_ENDPOINT_URL_BEDROCK_RUNTIME auto-routing ## Scope Bedrock-only for now. OpenAI/Anthropic use different SDKs (not boto3), would need separate hooks in their HTTP clients. Achieves complete parity with old ventis/llm/bedrock.py telemetry. --- examples/portfolio/agents/advisor_agent.py | 20 +-- examples/portfolio/agents/intent_agent.py | 19 +-- .../portfolio/config/global_controller.yaml | 27 +++ examples/text2sql/agents/vllm_agent.py | 20 +-- llm_proxy/hooks.py | 49 ------ llm_proxy/providers/bedrock.py | 78 --------- .../cloud_provider_logic/EC2/_runtime.py | 2 + .../cloud_provider_logic/Local/_runtime.py | 2 + ventis/controller/local_controller.py | 39 +++++ ventis/controller/utils/process_supervisor.py | 58 +++++++ ventis/llm/__init__.py | 0 ventis/llm/bedrock.py | 52 ------ {llm_proxy => ventis/llm_proxy}/README.md | 31 +++- {llm_proxy => ventis/llm_proxy}/__init__.py | 0 {llm_proxy => ventis/llm_proxy}/__main__.py | 4 +- {llm_proxy => ventis/llm_proxy}/app.py | 10 +- {llm_proxy => ventis/llm_proxy}/config.py | 5 + {llm_proxy => ventis/llm_proxy}/core.py | 5 +- ventis/llm_proxy/hooks.py | 155 ++++++++++++++++++ .../llm_proxy}/providers/__init__.py | 6 +- .../llm_proxy}/providers/anthropic.py | 2 +- .../llm_proxy}/providers/base.py | 0 ventis/llm_proxy/providers/bedrock.py | 119 ++++++++++++++ .../llm_proxy}/providers/openai.py | 2 +- ventis/llm_proxy/proxy.py | 57 +++++++ .../llm_proxy}/requirements.txt | 0 ventis/stub_generator.py | 49 +++++- 27 files changed, 572 insertions(+), 239 deletions(-) delete mode 100644 llm_proxy/hooks.py delete mode 100644 llm_proxy/providers/bedrock.py create mode 100644 ventis/controller/utils/process_supervisor.py delete mode 100644 ventis/llm/__init__.py delete mode 100644 ventis/llm/bedrock.py rename {llm_proxy => ventis/llm_proxy}/README.md (77%) rename {llm_proxy => ventis/llm_proxy}/__init__.py (100%) rename {llm_proxy => ventis/llm_proxy}/__main__.py (89%) rename {llm_proxy => ventis/llm_proxy}/app.py (79%) rename {llm_proxy => ventis/llm_proxy}/config.py (90%) rename {llm_proxy => ventis/llm_proxy}/core.py (87%) create mode 100644 ventis/llm_proxy/hooks.py rename {llm_proxy => ventis/llm_proxy}/providers/__init__.py (57%) rename {llm_proxy => ventis/llm_proxy}/providers/anthropic.py (87%) rename {llm_proxy => ventis/llm_proxy}/providers/base.py (100%) create mode 100644 ventis/llm_proxy/providers/bedrock.py rename {llm_proxy => ventis/llm_proxy}/providers/openai.py (85%) create mode 100644 ventis/llm_proxy/proxy.py rename {llm_proxy => ventis/llm_proxy}/requirements.txt (100%) diff --git a/examples/portfolio/agents/advisor_agent.py b/examples/portfolio/agents/advisor_agent.py index 0915eea..92ea2ba 100644 --- a/examples/portfolio/agents/advisor_agent.py +++ b/examples/portfolio/agents/advisor_agent.py @@ -2,11 +2,11 @@ # # Final stage. Turns the computed portfolio metrics and risk figures into a # short, plain-English briefing using a small, cheap model on AWS Bedrock -# (Converse API), called via ventis.llm.bedrock so token/cost telemetry gets -# recorded onto this execution's future: hash. Configure -# with env vars: +# (Converse API). Token/cost telemetry gets recorded automatically via the +# LLM proxy. Configure with env vars: # BEDROCK_MODEL_ID (default: meta.llama3-8b-instruct-v1:0) # AWS_REGION (default: us-east-1) +# AWS_ENDPOINT_URL_BEDROCK_RUNTIME (routes to proxy for telemetry) # # If the LLM is unavailable (returns an empty string), it falls back to a # deterministic templated summary so the pipeline still returns. @@ -14,11 +14,7 @@ # Resource profile: cheap CPU; the LLM cost sits in the Bedrock call, not here. import os - -try: - from ventis.llm.bedrock import call_bedrock -except ImportError: - from bedrock import call_bedrock +import boto3 class AdvisorAgent(object): @@ -28,16 +24,16 @@ def __init__(self): "BEDROCK_MODEL_ID", "meta.llama3-8b-instruct-v1:0" ) self.region = os.environ.get("AWS_REGION", "us-east-1") + self._client = boto3.client("bedrock-runtime", region_name=self.region) def summarize(self, holdings: dict, metrics: dict, risk: dict) -> str: """Write a short plain-English briefing on the portfolio.""" prompt = self._build_prompt(holdings, metrics, risk) try: - response = call_bedrock( - model_id=self.model_id, + response = self._client.converse( + modelId=self.model_id, messages=[{"role": "user", "content": [{"text": prompt}]}], - inference_config={"maxTokens": 400, "temperature": 0.2}, - region=self.region, + inferenceConfig={"maxTokens": 400, "temperature": 0.2}, ) return response["output"]["message"]["content"][0]["text"] except Exception as e: diff --git a/examples/portfolio/agents/intent_agent.py b/examples/portfolio/agents/intent_agent.py index d74b27b..7e9624e 100644 --- a/examples/portfolio/agents/intent_agent.py +++ b/examples/portfolio/agents/intent_agent.py @@ -7,9 +7,8 @@ # -> {"holdings": {"AAPL": 0.4, "MSFT": 0.35, "NVDA": 0.25}, # "lookback_days": 180} # -# Calls AWS Bedrock (Converse API) via ventis.llm.bedrock -- same pattern as -# AdvisorAgent -- so token/cost telemetry gets recorded onto this execution's -# future: hash. Configure with env vars: +# Calls AWS Bedrock (Converse API) via standard boto3. Telemetry is collected +# automatically by the LLM proxy. Configure with env vars: # BEDROCK_MODEL_ID (default: meta.llama3-8b-instruct-v1:0) # AWS_REGION (default: us-east-1) # @@ -23,11 +22,7 @@ import os import re import json - -try: - from ventis.llm.bedrock import call_bedrock -except ImportError: - from bedrock import call_bedrock +import boto3 DEFAULT_LOOKBACK_DAYS = 365 @@ -39,14 +34,14 @@ def __init__(self): "BEDROCK_MODEL_ID", "meta.llama3-8b-instruct-v1:0" ) self.region = os.environ.get("AWS_REGION", "us-east-1") + self._client = boto3.client("bedrock-runtime", region_name=self.region) def parse(self, query: str) -> dict: """Parse a natural-language portfolio request into holdings + lookback.""" - response = call_bedrock( - model_id=self.model_id, + response = self._client.converse( + modelId=self.model_id, messages=[{"role": "user", "content": [{"text": self._build_prompt(query)}]}], - inference_config={"maxTokens": 300, "temperature": 0.0}, - region=self.region, + inferenceConfig={"maxTokens": 300, "temperature": 0.0}, ) text = response["output"]["message"]["content"][0]["text"] if not text: diff --git a/examples/portfolio/config/global_controller.yaml b/examples/portfolio/config/global_controller.yaml index f63429b..96f371c 100644 --- a/examples/portfolio/config/global_controller.yaml +++ b/examples/portfolio/config/global_controller.yaml @@ -78,6 +78,20 @@ agents: provider: EC2 instance_type: t3.micro + +otel: + destinations: + - name: railway + protocol: grpc + endpoint: ${RAILWAY_OTLP_ENDPOINT} + insecure: true + headers: {} + - name: grafana + protocol: http + endpoint: ${GRAFANA_OTLP_ENDPOINT}/v1/traces + headers: + Authorization: Basic ${GRAFANA_OTLP_HEADERS} + # Polling interval in seconds poll_interval: 5 @@ -86,3 +100,16 @@ redis: host: localhost port: 6379 db: 0 + +# EC2 defaults for `provider: EC2` replicas. +ec2: + region: ${EC2_REGION} + ami_id: ${EC2_AMI_ID} + subnet_id: ${EC2_SUBNET_ID} + security_group_ids: + - ${EC2_SECURITY_GROUP_ID} + ssh_user: ${EC2_SSH_USER} + ssh_private_key_path: ${EC2_SSH_PRIVATE_KEY_PATH} + +database: + url: ${DATABASE_URL} diff --git a/examples/text2sql/agents/vllm_agent.py b/examples/text2sql/agents/vllm_agent.py index 4a7a245..dd1f80d 100644 --- a/examples/text2sql/agents/vllm_agent.py +++ b/examples/text2sql/agents/vllm_agent.py @@ -1,10 +1,8 @@ # VLLM Agent # # LLM backend for SQL candidate generation, called remotely by -# SQLGeneratorAgent. Calls AWS Bedrock (Converse API) via ventis.llm.bedrock -# so token/cost telemetry gets recorded onto this execution's -# future: hash — same pattern as -# examples/portfolio/agents/advisor_agent.py. +# SQLGeneratorAgent. Calls AWS Bedrock (Converse API) via standard boto3. +# Telemetry is collected automatically by the LLM proxy. # Configure with env vars: # BEDROCK_MODEL_ID (default: meta.llama3-8b-instruct-v1:0) # AWS_REGION (default: us-east-1) @@ -12,11 +10,7 @@ # Falls back to a synthetic placeholder response if Bedrock is unavailable. import os - -try: - from ventis.llm.bedrock import call_bedrock -except ImportError: - from bedrock import call_bedrock +import boto3 class VllmAgent(object): @@ -24,15 +18,15 @@ def __init__(self): self.tools = [self.generate] self.model_id = os.environ.get("BEDROCK_MODEL_ID", "meta.llama3-8b-instruct-v1:0") self.region = os.environ.get("AWS_REGION", "us-east-1") + self._client = boto3.client("bedrock-runtime", region_name=self.region) def generate(self, prompt: str) -> str: """Generates a response using an LLM model based on the given prompt.""" try: - response = call_bedrock( - model_id=self.model_id, + response = self._client.converse( + modelId=self.model_id, messages=[{"role": "user", "content": [{"text": prompt}]}], - inference_config={"maxTokens": 400, "temperature": 0.2}, - region=self.region, + inferenceConfig={"maxTokens": 400, "temperature": 0.2}, ) return response["output"]["message"]["content"][0]["text"] except Exception as e: diff --git a/llm_proxy/hooks.py b/llm_proxy/hooks.py deleted file mode 100644 index 37f9b3b..0000000 --- a/llm_proxy/hooks.py +++ /dev/null @@ -1,49 +0,0 @@ -"""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). -""" - -from __future__ import annotations - -import logging -import time -from dataclasses import dataclass -from typing import Any, Dict, Optional - -log = logging.getLogger("llm_proxy") - - -@dataclass -class Ctx: - provider: str - method: str - subpath: str - body: bytes - headers: Dict[str, str] - t0: float - model: Optional[str] = None - - def elapsed_ms(self) -> float: - return (time.monotonic() - self.t0) * 1000.0 - - -class Hooks: - def on_request(self, ctx: Ctx) -> None: - log.info( - "→ %s %s /%s model=%s (%d bytes)", - ctx.provider, ctx.method, ctx.subpath, ctx.model, len(ctx.body), - ) - - def on_response(self, ctx: Ctx, resp: Any) -> None: - log.info( - "← %s %s /%s -> %s in %.0fms", - ctx.provider, ctx.method, ctx.subpath, - getattr(resp, "status", "?"), ctx.elapsed_ms(), - ) - # TODO(metrics): pull token usage off `resp` and emit it. - - -hooks = Hooks() diff --git a/llm_proxy/providers/bedrock.py b/llm_proxy/providers/bedrock.py deleted file mode 100644 index c985d26..0000000 --- a/llm_proxy/providers/bedrock.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Bedrock adapter. - -Rather than re-sign the caller's SigV4 request (fiddly once model IDs contain -``:`` and ``/``), we re-issue the call through the proxy's own boto3 client, -which handles signing and URL-encoding correctly by construction. This is clean -for request/response; streaming (``invoke-with-response-stream``) is out of scope -for now. -""" - -from __future__ import annotations - -import json - -import boto3 -from botocore.exceptions import ClientError - -from llm_proxy.providers.base import Provider, ProxyResponse - -# bedrock-runtime operations that can appear as the last path segment; only the -# non-streaming "invoke" is wired up for now. -_SUPPORTED_OPS = {"invoke", "invoke-with-response-stream", "converse", "converse-stream"} - - -class BedrockProvider(Provider): - name = "bedrock" - - def __init__(self, cfg): - super().__init__(cfg) - self._client = boto3.client("bedrock-runtime", region_name=cfg.bedrock_region) - - def forward(self, req, subpath, body): - model_id, op = self._parse(subpath) - if op != "invoke": - raise NotImplementedError( - f"bedrock op '{op}' not supported yet (streaming/converse are out of scope)" - ) - try: - resp = self._client.invoke_model( - modelId=model_id, - body=body, - contentType=req.headers.get("Content-Type", "application/json"), - accept=req.headers.get("Accept", "application/json"), - ) - except ClientError as exc: - return self._error_response(exc) - - payload = resp["body"].read() - status = resp.get("ResponseMetadata", {}).get("HTTPStatusCode", 200) - headers = [("Content-Type", resp.get("contentType", "application/json"))] - return ProxyResponse(status=status, headers=headers, content=payload) - - @staticmethod - def _parse(subpath): - # subpath looks like "model//"; the modelId may itself - # contain "/" (inference-profile ARNs), so peel the op off the right. - if not subpath.startswith("model/"): - raise ValueError(f"unrecognized bedrock path: /{subpath}") - model_id, sep, op = subpath[len("model/"):].rpartition("/") - if not sep or op not in _SUPPORTED_OPS: - raise ValueError(f"unrecognized bedrock path: /{subpath}") - return model_id, op - - @staticmethod - def _error_response(exc: ClientError) -> ProxyResponse: - # boto3 raises on 4xx/5xx; reconstruct a JSON error body carrying the - # real status + message. (Byte-for-byte error passthrough is a property - # only the HTTP providers have; this is the cost of re-issuing via boto3.) - meta = exc.response.get("ResponseMetadata", {}) - err = exc.response.get("Error", {}) - status = meta.get("HTTPStatusCode", 500) - body = json.dumps( - {"message": err.get("Message", str(exc)), "code": err.get("Code")} - ).encode("utf-8") - return ProxyResponse( - status=status, - headers=[("Content-Type", "application/json")], - content=body, - ) diff --git a/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py index 4d5f766..57ca12b 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -277,6 +277,8 @@ def _bootstrap_instance(host, spec, replica_index, cfg, redis_host, redis_port, f"VENTIS_AGENT_PORT={CONTAINER_PORT}", "-e", f"VENTIS_POLL_INTERVAL={_controller.config.get('poll_interval', 5)}", + "-e", + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://127.0.0.1:8081/bedrock", ] if spec.get("type") == "workflow": db_url = _controller.config.get("database", {}).get("url") diff --git a/ventis/controller/cloud_provider_logic/Local/_runtime.py b/ventis/controller/cloud_provider_logic/Local/_runtime.py index 963eef3..9228309 100644 --- a/ventis/controller/cloud_provider_logic/Local/_runtime.py +++ b/ventis/controller/cloud_provider_logic/Local/_runtime.py @@ -94,6 +94,8 @@ def bootstrap_instance(provisioned, spec, replica_index, agent_id): f"VENTIS_REDIS_PORT={spec.get('redis_port', 6379)}", "-e", f"VENTIS_POLL_INTERVAL={_require_controller().config.get('poll_interval', 5)}", + "-e", + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://127.0.0.1:8081/bedrock", ] if ctrl_type == "workflow": cmd.extend(["-p", f"{spec.get('api_port', 8080)}:8080"]) diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 8b5942e..b3f53b0 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -6,6 +6,7 @@ import logging import os import random +import subprocess import sys import threading import time @@ -35,6 +36,16 @@ import ventis.ventis_context as ventis_context except ImportError: import ventis_context + +# Auto-inject Ventis headers into all boto3 Bedrock calls +try: + from ventis.llm_proxy import proxy +except ImportError: + try: + from llm_proxy import proxy + except ImportError: + pass # No proxy available, agents will call Bedrock directly + import local_controler_pb2 import local_controler_pb2_grpc @@ -102,6 +113,9 @@ def __init__(self, port=50051): max_instances = int(os.environ.get("VENTIS_MAX_AGENT_INSTANCES", 8)) self._executor = ThreadPoolExecutor(max_workers=max_instances) + # Start LLM proxy in this container + self._proxy_process = self._start_llm_proxy(redis_host, redis_port) + logger.info( "Local controller initialized at %s (max_agent_instances=%d), reported healthy to Redis.", self._my_endpoint, @@ -111,6 +125,31 @@ def __init__(self, port=50051): # Load the agent class dynamically self.agent = self._load_agent() + def _start_llm_proxy(self, redis_host, redis_port): + """Start LLM proxy as a subprocess in this container.""" + try: + import subprocess + proxy_env = os.environ.copy() + proxy_env.update({ + "PROXY_HOST": "127.0.0.1", + "PROXY_PORT": "8081", + "VENTIS_REDIS_HOST": redis_host, + "VENTIS_REDIS_PORT": str(redis_port), + }) + + # Start proxy as background subprocess + proxy_process = subprocess.Popen( + [sys.executable, "-m", "ventis.llm_proxy"], + env=proxy_env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + logger.info("Started LLM proxy on 127.0.0.1:8081 (PID: %d)", proxy_process.pid) + return proxy_process + except Exception as e: + logger.warning("Failed to start LLM proxy: %s", e) + return None + def _collect_metrics(self): """Snapshot current instance health/resource metrics. diff --git a/ventis/controller/utils/process_supervisor.py b/ventis/controller/utils/process_supervisor.py new file mode 100644 index 0000000..f06b6e1 --- /dev/null +++ b/ventis/controller/utils/process_supervisor.py @@ -0,0 +1,58 @@ +"""Registry for OS processes GlobalController spawns and supervises. + +register() + start_all() spawn processes; check_and_respawn() (call from GC's existing +poll tick) restarts any that exit unexpectedly; terminate_all() (call from GC's shutdown +path) stops them all cleanly. Deliberately GC-agnostic -- callers are responsible for not +calling check_and_respawn() during their own shutdown. +""" + +import logging +import os +import subprocess + +logger = logging.getLogger(__name__) + + +class ProcessSupervisor: + def __init__(self): + self._specs = {} # name -> (argv, env) tuple + self._procs = {} # name -> subprocess.Popen + + def register(self, name, argv, env=None): + """Declare a process to manage. Does not start it -- call start_all() once + everything is registered. `env`, if given, is merged on top of (not a + replacement for) this process's own environment, so the child still inherits + PATH etc.""" + self._specs[name] = (argv, env) + + def start_all(self): + for name, (argv, env) in self._specs.items(): + self._start(name, argv, env) + + def _start(self, name, argv, env=None): + merged_env = {**os.environ, **env} if env else None + self._procs[name] = subprocess.Popen(argv, env=merged_env) + + def check_and_respawn(self): + """Restart any registered process that has exited.""" + for name, proc in list(self._procs.items()): + if proc.poll() is not None: + logger.warning( + "Managed process %r exited (code %s), respawning", + name, + proc.returncode, + ) + argv, env = self._specs[name] + self._start(name, argv, env) + + def terminate_all(self, timeout=10): + """Terminate every managed process, falling back to kill() on timeout.""" + for proc in self._procs.values(): + proc.terminate() + for name, proc in self._procs.items(): + try: + proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + self._procs.clear() diff --git a/ventis/llm/__init__.py b/ventis/llm/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/ventis/llm/bedrock.py b/ventis/llm/bedrock.py deleted file mode 100644 index f350b69..0000000 --- a/ventis/llm/bedrock.py +++ /dev/null @@ -1,52 +0,0 @@ -import os - -try: - from ventis.utils.redis_client import RedisClient - import ventis.ventis_context as ventis_context -except ImportError: - from redis_client import RedisClient - import ventis_context - -_redis = RedisClient( - host=os.environ.get("VENTIS_REDIS_HOST", "localhost"), - port=int(os.environ.get("VENTIS_REDIS_PORT", 6379)), -) - - -def call_bedrock(model_id: str, messages: list, inference_config: dict, region: str = "us-east-1") -> dict: - """Call Bedrock's converse() API and log token/error telemetry onto the - currently executing future's hash (future:).""" - import boto3 - - client = boto3.client("bedrock-runtime", region_name=region) - future_id = ventis_context.get_current_future_id() - error_count = 0 - response = None - try: - response = client.converse( - modelId=model_id, messages=messages, inferenceConfig=inference_config - ) - return response - except Exception as e: - error_count += 1 - metrics_key = ventis_context.get_current_metrics_key() - if metrics_key: - _redis.hincrby(metrics_key, "error_count", 1) - # Deliberately does not write "error"/"failed" onto the future here -- - # that's owned by LocalController._mark_future_failed, which only - # fires if this exception propagates all the way up uncaught. If a - # caller catches and recovers (e.g. a fallback summary), the future - # succeeds, and writing a failure here would falsely mark it failed. - raise - finally: - if future_id: - usage = (response or {}).get("usage", {}) - _redis.hset_multiple(f"future:{future_id}", { - "model": model_id, - "input_token_count": str(usage.get("inputTokens", "")), - "output_token_count": str(usage.get("outputTokens", "")), - "token_count": str(usage.get("totalTokens", "")), - "errors": str(error_count), - "input_cache_tokens": str(usage.get("cacheReadInputTokens", "")), - "input_cache_write_tokens": str(usage.get("cacheWriteInputTokens", "")), - }) diff --git a/llm_proxy/README.md b/ventis/llm_proxy/README.md similarity index 77% rename from llm_proxy/README.md rename to ventis/llm_proxy/README.md index 2b8ca50..873144d 100644 --- a/llm_proxy/README.md +++ b/ventis/llm_proxy/README.md @@ -76,14 +76,33 @@ boto3.client("bedrock-runtime").invoke_model( | `BEDROCK_REGION` (or `AWS_REGION`) | `us-east-1` | Bedrock region | | `BEDROCK_UPSTREAM_HOST` | `bedrock-runtime..amazonaws.com` | override Bedrock host | -## The metrics seam +## Telemetry & Metrics -`hooks.py` defines `on_request` / `on_response`; today they only log. Because the -whole response is buffered, token accounting is a one-liner later — -`resp.json().get("usage")` for OpenAI/Anthropic, or the per-model field in -Bedrock's response body. +**Automatic telemetry is currently Bedrock-only.** The proxy captures: +- Model ID +- Input/output/total token counts +- Cache tokens (read & write) +- Error status -## Known limitations (current scope) +Telemetry is automatically written to Redis under `future:` keys. + +### How it works (Bedrock only) + +1. **Auto-injection:** boto3 hook (`proxy.py`) injects `X-Ventis-Future-Id` header from thread-local context +2. **Token extraction:** `hooks.py` parses response `usage` field +3. **Redis write:** All metrics written to `future:` hash + +### Why 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: +- 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. + +## Limitations - **No streaming.** `stream=True` / `invoke-with-response-stream` are not handled. - **Bedrock error bodies are reconstructed**, not passed through byte-for-byte diff --git a/llm_proxy/__init__.py b/ventis/llm_proxy/__init__.py similarity index 100% rename from llm_proxy/__init__.py rename to ventis/llm_proxy/__init__.py diff --git a/llm_proxy/__main__.py b/ventis/llm_proxy/__main__.py similarity index 89% rename from llm_proxy/__main__.py rename to ventis/llm_proxy/__main__.py index 621c37e..c0af2e9 100644 --- a/llm_proxy/__main__.py +++ b/ventis/llm_proxy/__main__.py @@ -4,8 +4,8 @@ import logging -from llm_proxy.app import create_app -from llm_proxy.config import Config +from ventis.llm_proxy.app import create_app +from ventis.llm_proxy.config import Config def main() -> None: diff --git a/llm_proxy/app.py b/ventis/llm_proxy/app.py similarity index 79% rename from llm_proxy/app.py rename to ventis/llm_proxy/app.py index ff32737..b2e3b09 100644 --- a/llm_proxy/app.py +++ b/ventis/llm_proxy/app.py @@ -7,9 +7,9 @@ from flask import Flask, jsonify, request -from llm_proxy.config import Config -from llm_proxy.core import proxy_request -from llm_proxy.providers import build_registry +from ventis.llm_proxy.config import Config +from ventis.llm_proxy.core import proxy_request +from ventis.llm_proxy.providers import build_registry log = logging.getLogger("llm_proxy") @@ -20,6 +20,10 @@ def create_app(cfg: Config = None) -> Flask: cfg = cfg or Config.from_env() app = Flask(__name__) registry = build_registry(cfg) + + # Initialize hooks with config for Redis + from ventis.llm_proxy import hooks as hooks_module + hooks_module.hooks = hooks_module.Hooks(cfg) @app.route("/healthz", methods=["GET"]) def healthz(): diff --git a/llm_proxy/config.py b/ventis/llm_proxy/config.py similarity index 90% rename from llm_proxy/config.py rename to ventis/llm_proxy/config.py index 9925e92..9e85cfa 100644 --- a/llm_proxy/config.py +++ b/ventis/llm_proxy/config.py @@ -28,6 +28,9 @@ class Config: bedrock_region: str bedrock_upstream_host: str + + redis_host: str + redis_port: int @classmethod def from_env(cls) -> "Config": @@ -58,4 +61,6 @@ def from_env(cls) -> "Config": bedrock_upstream_host=os.getenv( "BEDROCK_UPSTREAM_HOST", f"bedrock-runtime.{region}.amazonaws.com" ), + redis_host=os.getenv("VENTIS_REDIS_HOST", "localhost"), + redis_port=int(os.getenv("VENTIS_REDIS_PORT", "6379")), ) diff --git a/llm_proxy/core.py b/ventis/llm_proxy/core.py similarity index 87% rename from llm_proxy/core.py rename to ventis/llm_proxy/core.py index 4d1f608..6284e00 100644 --- a/llm_proxy/core.py +++ b/ventis/llm_proxy/core.py @@ -8,7 +8,7 @@ from flask import Response -from llm_proxy.hooks import Ctx, hooks +from ventis.llm_proxy.hooks import Ctx def _guess_model(body: bytes) -> Optional[str]: @@ -25,6 +25,9 @@ def _guess_model(body: bytes) -> Optional[str]: def proxy_request(provider, subpath, flask_request): + # Import hooks here to get the instance created by create_app + from ventis.llm_proxy.hooks import hooks + body = flask_request.get_data() ctx = Ctx( provider=provider.name, diff --git a/ventis/llm_proxy/hooks.py b/ventis/llm_proxy/hooks.py new file mode 100644 index 0000000..17713e9 --- /dev/null +++ b/ventis/llm_proxy/hooks.py @@ -0,0 +1,155 @@ +"""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). +""" + +from __future__ import annotations + +import json +import logging +import time +from dataclasses import dataclass +from typing import Any, Dict, Optional + +log = logging.getLogger("llm_proxy") + + +@dataclass +class TokenUsage: + """Token usage extracted from LLM responses.""" + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + input_cache_tokens: int = 0 + input_cache_write_tokens: int = 0 + + def __repr__(self): + parts = [f"in={self.input_tokens}", f"out={self.output_tokens}"] + if self.input_cache_tokens: + parts.append(f"cache_read={self.input_cache_tokens}") + if self.input_cache_write_tokens: + parts.append(f"cache_write={self.input_cache_write_tokens}") + return f"TokenUsage({', '.join(parts)})" + + +@dataclass +class Ctx: + provider: str + method: str + subpath: str + body: bytes + headers: Dict[str, str] + t0: float + model: Optional[str] = None + + def elapsed_ms(self) -> float: + return (time.monotonic() - self.t0) * 1000.0 + + +class Hooks: + def __init__(self, config=None): + self.config = config + self._redis = None + + if config: + try: + from ventis.utils.redis_client import RedisClient + self._redis = RedisClient( + host=config.redis_host, + port=config.redis_port, + ) + log.info("Redis telemetry enabled: %s:%s", config.redis_host, config.redis_port) + except Exception as e: + log.warning("Redis not available: %s", e) + + def on_request(self, ctx: Ctx) -> None: + log.info( + "→ %s %s /%s model=%s (%d bytes)", + ctx.provider, ctx.method, ctx.subpath, ctx.model, len(ctx.body), + ) + + 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) + + log.info( + "← %s %s /%s -> %s in %.0fms | %s", + ctx.provider, ctx.method, ctx.subpath, + getattr(resp, "status", "?"), ctx.elapsed_ms(), + usage or "no usage" + ) + + # Write to Redis if we have context + log.info("Checking telemetry write: redis=%s", "yes" if self._redis else "no") + if self._redis: + future_id = ctx.headers.get("X-Ventis-Future-Id") + log.info("Future ID from headers: %s", future_id) + if future_id: + try: + # Extract model ID + model_id = self._extract_model_id(ctx) + + is_error = resp.status >= 400 + + # Build telemetry data + data = { + "model": model_id, + "errors": "1" if is_error else "0", + } + + # Add token data if available + if usage: + data.update({ + "input_token_count": str(usage.input_tokens), + "output_token_count": str(usage.output_tokens), + "token_count": str(usage.total_tokens), + "input_cache_tokens": str(usage.input_cache_tokens), + "input_cache_write_tokens": str(usage.input_cache_write_tokens), + }) + + self._redis.hset_multiple(f"future:{future_id}", data) + log.info("Wrote telemetry to future:%s with data: %s", future_id, data) + except Exception as e: + log.error("Failed to write telemetry: %s", e) + + 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//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 + 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: + 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 + + +hooks = Hooks() diff --git a/llm_proxy/providers/__init__.py b/ventis/llm_proxy/providers/__init__.py similarity index 57% rename from llm_proxy/providers/__init__.py rename to ventis/llm_proxy/providers/__init__.py index 02f2a77..d9ff021 100644 --- a/llm_proxy/providers/__init__.py +++ b/ventis/llm_proxy/providers/__init__.py @@ -1,8 +1,8 @@ from __future__ import annotations -from llm_proxy.providers.anthropic import AnthropicProvider -from llm_proxy.providers.bedrock import BedrockProvider -from llm_proxy.providers.openai import OpenAIProvider +from ventis.llm_proxy.providers.anthropic import AnthropicProvider +from ventis.llm_proxy.providers.bedrock import BedrockProvider +from ventis.llm_proxy.providers.openai import OpenAIProvider def build_registry(cfg): diff --git a/llm_proxy/providers/anthropic.py b/ventis/llm_proxy/providers/anthropic.py similarity index 87% rename from llm_proxy/providers/anthropic.py rename to ventis/llm_proxy/providers/anthropic.py index 3c765bd..33e14aa 100644 --- a/llm_proxy/providers/anthropic.py +++ b/ventis/llm_proxy/providers/anthropic.py @@ -1,6 +1,6 @@ from __future__ import annotations -from llm_proxy.providers.base import HttpProvider, UpstreamRequest, client_headers +from ventis.llm_proxy.providers.base import HttpProvider, UpstreamRequest, client_headers class AnthropicProvider(HttpProvider): diff --git a/llm_proxy/providers/base.py b/ventis/llm_proxy/providers/base.py similarity index 100% rename from llm_proxy/providers/base.py rename to ventis/llm_proxy/providers/base.py diff --git a/ventis/llm_proxy/providers/bedrock.py b/ventis/llm_proxy/providers/bedrock.py new file mode 100644 index 0000000..f0efddc --- /dev/null +++ b/ventis/llm_proxy/providers/bedrock.py @@ -0,0 +1,119 @@ +"""Bedrock adapter. + +Rather than re-sign the caller's SigV4 request (fiddly once model IDs contain +``:`` and ``/``), we re-issue the call through the proxy's own boto3 client, +which handles signing and URL-encoding correctly by construction. This is clean +for request/response; streaming (``invoke-with-response-stream``) is out of scope +for now. +""" + +from __future__ import annotations + +import json + +import boto3 +from botocore.exceptions import ClientError + +from ventis.llm_proxy.providers.base import Provider, ProxyResponse + +# bedrock-runtime operations that can appear as the last path segment; only the +# non-streaming "invoke" is wired up for now. +_SUPPORTED_OPS = {"invoke", "invoke-with-response-stream", "converse", "converse-stream"} + + +class BedrockProvider(Provider): + name = "bedrock" + + def __init__(self, cfg): + super().__init__(cfg) + # Explicitly set endpoint_url to bypass AWS_ENDPOINT_URL_BEDROCK_RUNTIME + # environment variable that points to this proxy (would create infinite loop) + self._client = boto3.client( + "bedrock-runtime", + region_name=cfg.bedrock_region, + endpoint_url=f"https://{cfg.bedrock_upstream_host}" + ) + + def forward(self, req, subpath, body): + model_id, op = self._parse(subpath) + + try: + if op == "invoke": + resp = self._client.invoke_model( + modelId=model_id, + body=body, + contentType=req.headers.get("Content-Type", "application/json"), + accept=req.headers.get("Accept", "application/json"), + ) + # For invoke, return raw response body + payload = resp["body"].read() + status = resp.get("ResponseMetadata", {}).get("HTTPStatusCode", 200) + headers = [("Content-Type", resp.get("contentType", "application/json"))] + return ProxyResponse(status=status, headers=headers, content=payload) + + elif op == "converse": + params = json.loads(body) + params["modelId"] = model_id + resp = self._client.converse(**params) + + # Return response as JSON + response_data = { + "output": resp.get("output", {}), + "stopReason": resp.get("stopReason"), + "usage": resp.get("usage", {}), + } + # Include optional fields if present + for field in ["metrics", "trace", "additionalModelResponseFields"]: + if field in resp: + response_data[field] = resp[field] + + payload = json.dumps(response_data).encode("utf-8") + status = resp.get("ResponseMetadata", {}).get("HTTPStatusCode", 200) + return ProxyResponse( + status=status, + headers=[("Content-Type", "application/json")], + content=payload + ) + else: + raise NotImplementedError( + f"bedrock op '{op}' not supported (only invoke and converse)" + ) + + except ClientError as exc: + return self._error_response(exc) + except (json.JSONDecodeError, KeyError) as exc: + return ProxyResponse( + status=400, + headers=[("Content-Type", "application/json")], + content=json.dumps({"message": f"Invalid request: {exc}"}).encode(), + ) + + + + @staticmethod + def _parse(subpath): + # subpath looks like "model//"; the modelId may itself + # contain "/" (inference-profile ARNs), so peel the op off the right. + if not subpath.startswith("model/"): + raise ValueError(f"unrecognized bedrock path: /{subpath}") + model_id, sep, op = subpath[len("model/"):].rpartition("/") + if not sep or op not in _SUPPORTED_OPS: + raise ValueError(f"unrecognized bedrock path: /{subpath}") + return model_id, op + + @staticmethod + def _error_response(exc: ClientError) -> ProxyResponse: + # boto3 raises on 4xx/5xx; reconstruct a JSON error body carrying the + # real status + message. (Byte-for-byte error passthrough is a property + # only the HTTP providers have; this is the cost of re-issuing via boto3.) + meta = exc.response.get("ResponseMetadata", {}) + err = exc.response.get("Error", {}) + status = meta.get("HTTPStatusCode", 500) + body = json.dumps( + {"message": err.get("Message", str(exc)), "code": err.get("Code")} + ).encode("utf-8") + return ProxyResponse( + status=status, + headers=[("Content-Type", "application/json")], + content=body, + ) diff --git a/llm_proxy/providers/openai.py b/ventis/llm_proxy/providers/openai.py similarity index 85% rename from llm_proxy/providers/openai.py rename to ventis/llm_proxy/providers/openai.py index ed7f08d..67457eb 100644 --- a/llm_proxy/providers/openai.py +++ b/ventis/llm_proxy/providers/openai.py @@ -1,6 +1,6 @@ from __future__ import annotations -from llm_proxy.providers.base import HttpProvider, UpstreamRequest, client_headers +from ventis.llm_proxy.providers.base import HttpProvider, UpstreamRequest, client_headers class OpenAIProvider(HttpProvider): diff --git a/ventis/llm_proxy/proxy.py b/ventis/llm_proxy/proxy.py new file mode 100644 index 0000000..acf843b --- /dev/null +++ b/ventis/llm_proxy/proxy.py @@ -0,0 +1,57 @@ +"""Auto-inject Ventis headers into ALL boto3 Bedrock calls. + +Import this module once and all subsequent boto3.client("bedrock-runtime") calls +will automatically include the X-Ventis-Future-ID header. + +Usage: + import ventis.llm_proxy_auto # Just import once + import boto3 + + # Now this automatically includes the header! + client = boto3.client("bedrock-runtime") + response = client.converse(...) +""" + +import boto3 +import logging + +try: + import ventis.ventis_context as ventis_context +except ImportError: + ventis_context = None + +log = logging.getLogger(__name__) + + +def _inject_ventis_headers(event_name=None, **kwargs): + """Inject X-Ventis-Future-ID header into boto3 requests.""" + if not ventis_context: + return + + # Only inject for bedrock-runtime service + if 'service_id' in kwargs and kwargs.get('service_id') != 'Bedrock Runtime': + return + + # Get the request object + request = kwargs.get('request') + if not request: + return + + # Get current future_id from thread-local context + try: + future_id = ventis_context.get_current_future_id() + if future_id: + request.headers['X-Ventis-Future-ID'] = future_id + log.debug("Injected X-Ventis-Future-ID: %s", future_id) + except Exception as e: + log.debug("Could not inject future_id: %s", e) + + +# Register the hook globally on the default session +_session = boto3.Session() +_session.events.register_first('before-call.bedrock-runtime', _inject_ventis_headers) + +# Also patch the default session used by boto3.client() +boto3.DEFAULT_SESSION = _session + +log.info("Ventis boto3 hook registered - all Bedrock calls will include future_id header") diff --git a/llm_proxy/requirements.txt b/ventis/llm_proxy/requirements.txt similarity index 100% rename from llm_proxy/requirements.txt rename to ventis/llm_proxy/requirements.txt diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index e408553..6ffc03d 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -19,7 +19,7 @@ # Packages every agent container needs regardless of its specific business logic. # grpcio-tools/pyyaml/ipdb/ipython aren't needed/used, but keeping to keep the scope constrained right now # - Leave a comment if you want me to remove these, I kept them in since you originally had them but they aren't used -BASE_AGENT_REQUIREMENTS = ["grpcio", "grpcio-tools", "redis", "pyyaml", "psutil", "ipdb", "ipython", "boto3"] +BASE_AGENT_REQUIREMENTS = ["grpcio", "grpcio-tools", "redis", "pyyaml", "psutil", "ipdb", "ipython", "boto3", "flask>=2.0", "requests>=2.28"] # Workflow will always require these BASE_WORKFLOW_REQUIREMENTS = BASE_AGENT_REQUIREMENTS + ["flask", "sqlalchemy", "psycopg[binary]"] @@ -333,7 +333,8 @@ def generate_docker( os.path.join(script_dir, "controller", "utils", "gpu_metrics.py"), "gpu_metrics.py", ), - (os.path.join(script_dir, "llm", "bedrock.py"), "bedrock.py"), + # Note: ventis package directories are copied after the loop below + (None, "ventis/__init__.py"), # Empty __init__.py for ventis package ] # Copy provided agent stubs @@ -352,10 +353,28 @@ def generate_docker( files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) for src, dst in files_to_copy: - if os.path.isfile(src): - shutil.copy2(src, os.path.join(output_dir, dst)) + dst_path = os.path.join(output_dir, dst) + os.makedirs(os.path.dirname(dst_path), exist_ok=True) + if src is None: + # Create empty file (for __init__.py placeholders) + with open(dst_path, 'w') as f: + f.write('') + elif os.path.isfile(src): + shutil.copy2(src, dst_path) else: print(f" Warning: source file not found, skipping: {src}") + + # Copy entire ventis subdirectories needed by the proxy + shutil.copytree( + os.path.join(script_dir, "llm_proxy"), + os.path.join(output_dir, "ventis/llm_proxy"), + dirs_exist_ok=True, ignore=shutil.ignore_patterns('__pycache__', '*.pyc') + ) + shutil.copytree( + os.path.join(script_dir, "utils"), + os.path.join(output_dir, "ventis/utils"), + dirs_exist_ok=True, ignore=shutil.ignore_patterns('__pycache__', '*.pyc') + ) # Copy the YAML definition too shutil.copy2( @@ -464,10 +483,28 @@ def generate_workflow_docker( files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) for src, dst in files_to_copy: - if os.path.isfile(src): - shutil.copy2(src, os.path.join(output_dir, dst)) + dst_path = os.path.join(output_dir, dst) + os.makedirs(os.path.dirname(dst_path), exist_ok=True) + if src is None: + # Create empty file (for __init__.py placeholders) + with open(dst_path, 'w') as f: + f.write('') + elif os.path.isfile(src): + shutil.copy2(src, dst_path) else: print(f" Warning: source file not found, skipping: {src}") + + # Copy entire ventis subdirectories needed by the proxy + shutil.copytree( + os.path.join(script_dir, "llm_proxy"), + os.path.join(output_dir, "ventis/llm_proxy"), + dirs_exist_ok=True, ignore=shutil.ignore_patterns('__pycache__', '*.pyc') + ) + shutil.copytree( + os.path.join(script_dir, "utils"), + os.path.join(output_dir, "ventis/utils"), + dirs_exist_ok=True, ignore=shutil.ignore_patterns('__pycache__', '*.pyc') + ) # ---- workflow_launcher.py -------------------------------------------- launcher = f"""import threading