diff --git a/README.md b/README.md index 8fc4084..4cb7ed0 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ Edit `.car/config/global_controller.yaml` in your project directory to list the Agents that need API keys read them from environment variables. Point `env_file` at a `.env` file to have Ventis inject it into every agent container: ```yaml -# config/global_controller.yaml +# .car/config/global_controller.yaml env_file: .env ``` diff --git a/cli/README.md b/cli/README.md index 4313946..4acfdb7 100644 --- a/cli/README.md +++ b/cli/README.md @@ -23,4 +23,4 @@ leaving every other line unchanged. If you have a workflow running, and want to make a config change, canyonos config automatically would reload the project with your config. If you change the workflow files itself though and want the changes to take effect, you need to redeploy from scratch, running canyonos build for good measure -# Use: canyonos -h \ No newline at end of file +# Use: canyonos -h diff --git a/examples/portfolio/agents/advisor_agent.py b/examples/portfolio/agents/advisor_agent.py index 82d8936..33ac8fc 100644 --- a/examples/portfolio/agents/advisor_agent.py +++ b/examples/portfolio/agents/advisor_agent.py @@ -2,8 +2,9 @@ # # 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.controller.bedrock so token/cost telemetry gets -# recorded onto this execution's future: hash. Configure +# (Converse API), called directly via boto3. Token/cost telemetry is recorded +# onto this execution's future: hash transparently by the Ventis LLM +# proxy each agent container's boto3 calls are routed through. Configure # with env vars: # BEDROCK_MODEL_ID (default: meta.llama3-8b-instruct-v1:0) # AWS_REGION (default: us-east-1) @@ -15,10 +16,7 @@ import os -try: - from ventis.controller.bedrock import call_bedrock -except ImportError: - from bedrock import call_bedrock +import boto3 class AdvisorAgent(object): @@ -28,16 +26,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 04124cf..4bb915c 100644 --- a/examples/portfolio/agents/intent_agent.py +++ b/examples/portfolio/agents/intent_agent.py @@ -7,9 +7,11 @@ # -> {"holdings": {"AAPL": 0.4, "MSFT": 0.35, "NVDA": 0.25}, # "lookback_days": 180} # -# Calls AWS Bedrock (Converse API) via ventis.controller.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) directly via boto3 -- same pattern as +# AdvisorAgent. Token/cost telemetry is recorded onto this execution's +# future: hash transparently by the Ventis LLM proxy, which each +# agent container's boto3 calls are routed through (AWS_ENDPOINT_URL_BEDROCK_RUNTIME). +# Configure with env vars: # BEDROCK_MODEL_ID (default: meta.llama3-8b-instruct-v1:0) # AWS_REGION (default: us-east-1) # @@ -24,10 +26,7 @@ import re import json -try: - from ventis.controller.bedrock import call_bedrock -except ImportError: - from bedrock import call_bedrock +import boto3 DEFAULT_LOOKBACK_DAYS = 365 @@ -39,14 +38,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 205a672..60abb9e 100644 --- a/examples/portfolio/config/global_controller.yaml +++ b/examples/portfolio/config/global_controller.yaml @@ -7,7 +7,7 @@ agents: # Stage 0: parse the free-text request into structured holdings + lookback - # window (calls Bedrock directly via ventis.llm.bedrock). Cheap CPU, one + # window (calls Bedrock via boto3, routed through the Ventis LLM proxy). Cheap CPU, one # call per request, on the critical path before the fan-out. - name: IntentAgent redis_port: 6379 diff --git a/examples/text2sql/agents/vllm_agent.py b/examples/text2sql/agents/vllm_agent.py index ea23616..1d8141f 100644 --- a/examples/text2sql/agents/vllm_agent.py +++ b/examples/text2sql/agents/vllm_agent.py @@ -1,9 +1,10 @@ # VLLM Agent # # LLM backend for SQL candidate generation, called remotely by -# SQLGeneratorAgent. Calls AWS Bedrock (Converse API) via ventis.controller.bedrock -# so token/cost telemetry gets recorded onto this execution's -# future: hash — same pattern as +# SQLGeneratorAgent. Calls AWS Bedrock (Converse API) directly via boto3. +# Token/cost telemetry is recorded onto this execution's future: +# hash transparently by the Ventis LLM proxy each agent container's boto3 calls +# are routed through — same pattern as # examples/portfolio/agents/advisor_agent.py. # Configure with env vars: # BEDROCK_MODEL_ID (default: meta.llama3-8b-instruct-v1:0) @@ -13,10 +14,7 @@ import os -try: - from ventis.controller.bedrock import call_bedrock -except ImportError: - from bedrock import call_bedrock +import boto3 class VllmAgent(object): @@ -24,15 +22,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/pyproject.toml b/pyproject.toml index 41169a9..a94a9b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "psycopg[binary]", "pyyaml", "flask", + "requests", "psutil", "opentelemetry-api>=1.44.0", "opentelemetry-sdk>=1.44.0", diff --git a/requirements.txt b/requirements.txt index f06e0b0..20abfff 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ grpcio-tools redis pyyaml flask +requests sqlalchemy psycopg[binary] psutil diff --git a/tests/test_instance_manager_runtime.py b/tests/test_instance_manager_runtime.py index 13f9876..2c481c3 100644 --- a/tests/test_instance_manager_runtime.py +++ b/tests/test_instance_manager_runtime.py @@ -170,6 +170,8 @@ def test_local_instances_keep_default_host_and_increment_host_ports(self): "VENTIS_REDIS_PORT=6379", "-e", "VENTIS_POLL_INTERVAL=5", + "-e", + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://127.0.0.1:8081/bedrock", "ventis-alpha", ], "localhost", @@ -226,6 +228,8 @@ def test_local_workflow_and_resource_flags_stay_the_same(self): "VENTIS_REDIS_PORT=6379", "-e", "VENTIS_POLL_INTERVAL=5", + "-e", + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://127.0.0.1:8081/bedrock", "-p", "8080:8080", "--cpus", diff --git a/uv.lock b/uv.lock index b7e8f96..017bf9e 100644 --- a/uv.lock +++ b/uv.lock @@ -1229,6 +1229,7 @@ dependencies = [ { name = "psycopg", extra = ["binary"] }, { name = "pyyaml" }, { name = "redis" }, + { name = "requests" }, { name = "sqlalchemy" }, ] @@ -1252,6 +1253,7 @@ requires-dist = [ { name = "psycopg", extras = ["binary"] }, { name = "pyyaml" }, { name = "redis" }, + { name = "requests" }, { name = "sqlalchemy" }, ] diff --git a/ventis/FUTURE_SCHEMA.md b/ventis/FUTURE_SCHEMA.md index 360b1af..0d122b4 100644 --- a/ventis/FUTURE_SCHEMA.md +++ b/ventis/FUTURE_SCHEMA.md @@ -19,16 +19,16 @@ Fields currently written into `future:{future_id}`, and where: | `created_at` | `future.py` only (origin submission time) | | `result` | `future.py`, `local_controller.py` | | `failed` | `future.py`, `local_controller.py` | -| `error` | `future.py` (`_submit_request`), `local_controller.py` (`_mark_future_failed`) -- the sole failure-message field; `bedrock.py` deliberately never writes it | +| `error` | `future.py` (`_submit_request`), `local_controller.py` (`_mark_future_failed`) -- the sole failure-message field; the LLM proxy deliberately never writes it | | `finished_at` | `local_controller.py` (`_execute_locally` finally block) | | `cpu_resource` | `local_controller.py` | | `gpu_resource` | `local_controller.py` | | `agent` | `local_controller.py` (agent_id that executed this step) | | `queue_time` | `local_controller.py` (only when `submitted_at` is known) | -| `model` | `llm/bedrock.py` (`call_bedrock`) | -| `input_token_count` | `llm/bedrock.py` | -| `output_token_count` | `llm/bedrock.py` | -| `token_count` | `llm/bedrock.py` | -| `errors` | `llm/bedrock.py` (Bedrock call error count) | -| `input_cache_tokens` | `llm/bedrock.py` | -| `input_cache_write_tokens` | `llm/bedrock.py` | +| `model` | `llm_proxy/hooks.py` (on_response) | +| `input_token_count` | `llm_proxy/hooks.py` | +| `output_token_count` | `llm_proxy/hooks.py` | +| `token_count` | `llm_proxy/hooks.py` | +| `errors` | `llm_proxy/hooks.py` (Bedrock call error flag) | +| `input_cache_tokens` | `llm_proxy/hooks.py` | +| `input_cache_write_tokens` | `llm_proxy/hooks.py` | diff --git a/ventis/controller/bedrock.py b/ventis/controller/bedrock.py deleted file mode 100644 index 97c6a3f..0000000 --- a/ventis/controller/bedrock.py +++ /dev/null @@ -1,52 +0,0 @@ -import os - -try: - from ventis.controller.utils.redis_client import RedisClient - import ventis.controller.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/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py index be3bb18..ce68744 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -288,6 +288,10 @@ 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)}", + # Route the agent's boto3 Bedrock calls through the in-container LLM + # proxy (started by LocalController) so token/cost telemetry is captured. + "-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 88e4ec3..6329311 100644 --- a/ventis/controller/cloud_provider_logic/Local/_runtime.py +++ b/ventis/controller/cloud_provider_logic/Local/_runtime.py @@ -93,6 +93,10 @@ 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)}", + # Route the agent's boto3 Bedrock calls through the in-container LLM + # proxy (started by LocalController) so token/cost telemetry is captured. + "-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 8d9b525..1cc879f 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -35,6 +35,18 @@ import ventis.controller.ventis_context as ventis_context except ImportError: import ventis_context + +# Auto-inject X-Ventis-Future-ID into all boto3 Bedrock calls so the LLM proxy +# can attribute token/cost telemetry to the executing future. Import for its +# global boto3 event-hook side effect; safe no-op if the proxy isn't present. +try: + from ventis.llm_proxy import proxy as _llm_proxy_autoinject # noqa: F401 +except ImportError: + try: + from llm_proxy import proxy as _llm_proxy_autoinject # noqa: F401 + except ImportError: + pass # No proxy available; agents call Bedrock directly. + import local_controler_pb2 import local_controler_pb2_grpc @@ -102,6 +114,11 @@ def __init__(self, port=50051): max_instances = int(os.environ.get("VENTIS_MAX_AGENT_INSTANCES", 8)) self._executor = ThreadPoolExecutor(max_workers=max_instances) + # Start the LLM proxy alongside the agent in this container. Bedrock + # calls are routed to it via AWS_ENDPOINT_URL_BEDROCK_RUNTIME (injected + # by the runtime), and it writes token/cost telemetry to Redis. + 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 +128,33 @@ 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 the LLM proxy as a subprocess in this container (127.0.0.1:8081). + + Best-effort: a failure here must never stop the controller from coming up. + """ + import subprocess + + try: + 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), + }) + proxy_process = subprocess.Popen( + [sys.executable, "-m", "ventis.llm_proxy"], + env=proxy_env, + ) + 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/llm_proxy/README.md b/ventis/llm_proxy/README.md new file mode 100644 index 0000000..873144d --- /dev/null +++ b/ventis/llm_proxy/README.md @@ -0,0 +1,112 @@ +# llm_proxy + +A local, single-machine pass-through proxy for **OpenAI**, **Anthropic**, and +**Bedrock**. Callers keep their exact SDK calling convention — the only change is +one base-URL env var per provider. Every call flows through one function +(`core.proxy_request`) where token/metrics hooks fire. + +**Scope:** request/response ("call and return") only. Streaming is intentionally +not implemented yet. + +## How it works + +``` +your app (unchanged) localhost:8080 real upstream + openai SDK ─/openai/... ─┐ + anthropic SDK ─/anthropic/ ─┼─▶ proxy_request(ctx) ─▶ provider ─▶ api.openai.com + boto3 bedrock ─/bedrock/... ┘ (metrics hooks) adapter api.anthropic.com + bedrock-runtime..amazonaws.com +``` + +- **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. + +## Run + +```bash +pip install -r llm_proxy/requirements.txt + +# real upstream credentials live here; callers can use dummy keys +export OPENAI_API_KEY=sk-... +export ANTHROPIC_API_KEY=sk-ant-... +export AWS_REGION=us-east-1 # + normal AWS creds (env / ~/.aws / role) + +python -m llm_proxy # listens on 127.0.0.1:8080 +``` + +## Point your SDKs at it + +No code changes — just env vars: + +```bash +export OPENAI_BASE_URL=http://localhost:8080/openai/v1 +export ANTHROPIC_BASE_URL=http://localhost:8080/anthropic +export AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://localhost:8080/bedrock +``` + +Then your existing code works unchanged: + +```python +from openai import OpenAI +OpenAI().chat.completions.create(model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}]) + +from anthropic import Anthropic +Anthropic().messages.create(model="claude-3-5-sonnet-20241022", max_tokens=64, + messages=[{"role": "user", "content": "hi"}]) + +import boto3, json +boto3.client("bedrock-runtime").invoke_model( + modelId="anthropic.claude-3-5-sonnet-20240620-v1:0", + body=json.dumps({"anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}]})) +``` + +## Configuration (env vars) + +| Var | Default | Purpose | +|---|---|---| +| `PROXY_HOST` / `PROXY_PORT` | `127.0.0.1` / `8080` | where the proxy listens | +| `PROXY_CONNECT_TIMEOUT` / `PROXY_READ_TIMEOUT` | `10` / `600` | upstream timeouts (s) | +| `OPENAI_API_KEY` / `ANTHROPIC_API_KEY` | — | real upstream keys the proxy injects | +| `OPENAI_UPSTREAM_BASE` / `ANTHROPIC_UPSTREAM_BASE` | official APIs | override upstream (e.g. Azure/gateway) | +| `BEDROCK_REGION` (or `AWS_REGION`) | `us-east-1` | Bedrock region | +| `BEDROCK_UPSTREAM_HOST` | `bedrock-runtime..amazonaws.com` | override Bedrock host | + +## Telemetry & Metrics + +**Automatic telemetry is currently Bedrock-only.** The proxy captures: +- Model ID +- Input/output/total token counts +- Cache tokens (read & write) +- Error status + +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 + (boto3 raises on 4xx/5xx; we rebuild a JSON body with the real status + + message). OpenAI/Anthropic errors pass through unchanged. +- **Dev server.** Runs on Flask's built-in server — fine for a local proxy, not + meant for production traffic. diff --git a/ventis/llm_proxy/__init__.py b/ventis/llm_proxy/__init__.py new file mode 100644 index 0000000..827377e --- /dev/null +++ b/ventis/llm_proxy/__init__.py @@ -0,0 +1,14 @@ +"""Local LLM proxy. + +A transparent, single-machine pass-through for OpenAI, Anthropic, and Bedrock. +Point each provider's SDK at this service via its base-URL env var and calls flow +through one choke point (``llm_proxy.core.proxy_request``) where request/response +metrics hooks fire. + +Scope: request/response ("call and return") only. Streaming is intentionally +not implemented yet. +""" + +__all__ = ["__version__"] + +__version__ = "0.1.0" diff --git a/ventis/llm_proxy/__main__.py b/ventis/llm_proxy/__main__.py new file mode 100644 index 0000000..c0af2e9 --- /dev/null +++ b/ventis/llm_proxy/__main__.py @@ -0,0 +1,29 @@ +"""Entry point: ``python -m llm_proxy``.""" + +from __future__ import annotations + +import logging + +from ventis.llm_proxy.app import create_app +from ventis.llm_proxy.config import Config + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + cfg = Config.from_env() + app = create_app(cfg) + logging.getLogger("llm_proxy").info( + "llm_proxy on http://%s:%d (openai=%s, anthropic=%s, bedrock=%s [%s])", + cfg.host, cfg.port, cfg.openai.upstream_base, cfg.anthropic.upstream_base, + cfg.bedrock_upstream_host, cfg.bedrock_region, + ) + # threaded so concurrent callers don't serialize; dev server is fine for a + # local proxy. + app.run(host=cfg.host, port=cfg.port, threaded=True) + + +if __name__ == "__main__": + main() diff --git a/ventis/llm_proxy/app.py b/ventis/llm_proxy/app.py new file mode 100644 index 0000000..b2e3b09 --- /dev/null +++ b/ventis/llm_proxy/app.py @@ -0,0 +1,46 @@ +"""Flask app: one catch-all route per provider prefix, all funneled through +``proxy_request``.""" + +from __future__ import annotations + +import logging + +from flask import Flask, jsonify, request + +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") + +ALL_METHODS = ["GET", "POST", "PUT", "PATCH", "DELETE"] + + +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(): + return jsonify(status="ok", providers=sorted(registry.keys())) + + @app.route("//", methods=ALL_METHODS) + def dispatch(provider, subpath): + prov = registry.get(provider) + if prov is None: + return ( + jsonify(error=f"unknown provider '{provider}'", known=sorted(registry.keys())), + 404, + ) + try: + return proxy_request(prov, subpath, request) + except Exception as exc: # surface upstream/adapter errors as 502 + log.exception("proxy error for %s/%s", provider, subpath) + return jsonify(error="proxy_error", detail=str(exc)), 502 + + return app diff --git a/ventis/llm_proxy/config.py b/ventis/llm_proxy/config.py new file mode 100644 index 0000000..9e85cfa --- /dev/null +++ b/ventis/llm_proxy/config.py @@ -0,0 +1,66 @@ +"""Configuration, read once from the environment at startup. + +The proxy holds the *real* upstream credentials; callers can send dummy keys. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class ProviderConfig: + upstream_base: str + api_key: Optional[str] = None + + +@dataclass +class Config: + host: str + port: int + connect_timeout: float + read_timeout: float + + openai: ProviderConfig + anthropic: ProviderConfig + + bedrock_region: str + bedrock_upstream_host: str + + redis_host: str + redis_port: int + + @classmethod + def from_env(cls) -> "Config": + region = ( + os.getenv("BEDROCK_REGION") + or os.getenv("AWS_REGION") + or os.getenv("AWS_DEFAULT_REGION") + or "us-east-1" + ) + return cls( + host=os.getenv("PROXY_HOST", "127.0.0.1"), + port=int(os.getenv("PROXY_PORT", "8080")), + connect_timeout=float(os.getenv("PROXY_CONNECT_TIMEOUT", "10")), + read_timeout=float(os.getenv("PROXY_READ_TIMEOUT", "600")), + openai=ProviderConfig( + upstream_base=os.getenv( + "OPENAI_UPSTREAM_BASE", "https://api.openai.com" + ).rstrip("/"), + api_key=os.getenv("OPENAI_API_KEY"), + ), + anthropic=ProviderConfig( + upstream_base=os.getenv( + "ANTHROPIC_UPSTREAM_BASE", "https://api.anthropic.com" + ).rstrip("/"), + api_key=os.getenv("ANTHROPIC_API_KEY"), + ), + bedrock_region=region, + 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/ventis/llm_proxy/core.py b/ventis/llm_proxy/core.py new file mode 100644 index 0000000..6284e00 --- /dev/null +++ b/ventis/llm_proxy/core.py @@ -0,0 +1,46 @@ +"""The single choke point every proxied call flows through.""" + +from __future__ import annotations + +import json +import time +from typing import Optional + +from flask import Response + +from ventis.llm_proxy.hooks import Ctx + + +def _guess_model(body: bytes) -> Optional[str]: + """Best-effort model name from the JSON body, for logging/metrics. + + Never raises. Returns None for requests whose model isn't in the body + (e.g. Bedrock, where it's in the path and already shown via the subpath). + """ + try: + model = json.loads(body).get("model") + return model if isinstance(model, str) else None + except Exception: + return None + + +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, + method=flask_request.method, + subpath=subpath, + body=body, + headers=dict(flask_request.headers), + t0=time.monotonic(), + model=_guess_model(body), + ) + hooks.on_request(ctx) + + pr = provider.forward(flask_request, subpath, body) + + hooks.on_response(ctx, pr) + return Response(pr.content, status=pr.status, headers=pr.headers) diff --git a/ventis/llm_proxy/hooks.py b/ventis/llm_proxy/hooks.py new file mode 100644 index 0000000..7cc9957 --- /dev/null +++ b/ventis/llm_proxy/hooks.py @@ -0,0 +1,159 @@ +"""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: + try: + from ventis.controller.utils.redis_client import RedisClient + except ImportError: + # In-container the framework files are copied flat to /app. + from 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/ventis/llm_proxy/providers/__init__.py b/ventis/llm_proxy/providers/__init__.py new file mode 100644 index 0000000..d9ff021 --- /dev/null +++ b/ventis/llm_proxy/providers/__init__.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +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): + """Map the URL prefix -> provider instance.""" + return { + "openai": OpenAIProvider(cfg), + "anthropic": AnthropicProvider(cfg), + "bedrock": BedrockProvider(cfg), + } diff --git a/ventis/llm_proxy/providers/anthropic.py b/ventis/llm_proxy/providers/anthropic.py new file mode 100644 index 0000000..33e14aa --- /dev/null +++ b/ventis/llm_proxy/providers/anthropic.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from ventis.llm_proxy.providers.base import HttpProvider, UpstreamRequest, client_headers + + +class AnthropicProvider(HttpProvider): + name = "anthropic" + + def target(self, req, subpath, body): + headers = client_headers(req, drop=["x-api-key", "authorization"]) + if self.cfg.anthropic.api_key: + headers["x-api-key"] = self.cfg.anthropic.api_key + # `anthropic-version` is supplied by the SDK and passes through untouched. + return UpstreamRequest( + method=req.method, + url=f"{self.cfg.anthropic.upstream_base}/{subpath}", + headers=headers, + params=req.args.to_dict(flat=True), + ) diff --git a/ventis/llm_proxy/providers/base.py b/ventis/llm_proxy/providers/base.py new file mode 100644 index 0000000..ef5db41 --- /dev/null +++ b/ventis/llm_proxy/providers/base.py @@ -0,0 +1,96 @@ +"""Provider abstraction and shared HTTP plumbing. + +A provider's only job is to take the incoming request and produce a +``ProxyResponse``. Straight HTTP reverse-proxy providers (OpenAI, Anthropic) +subclass ``HttpProvider`` and just describe the upstream target; Bedrock owns +its own ``forward`` because it re-issues through boto3. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Dict, Iterable, List, Tuple + +import requests + +# Request headers we never forward: hop-by-hop (RFC 7230), ones we rewrite, and +# accept-encoding (we let the HTTP client negotiate + decode, then re-frame the +# response ourselves). +DROP_REQUEST_HEADERS = { + "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", + "te", "trailers", "transfer-encoding", "upgrade", + "host", "content-length", "accept-encoding", +} + +# Response headers we drop: we return already-decoded content and let the WSGI +# layer recompute framing headers. +DROP_RESPONSE_HEADERS = { + "content-encoding", "content-length", "transfer-encoding", + "connection", "keep-alive", +} + + +@dataclass +class UpstreamRequest: + method: str + url: str + headers: Dict[str, str] + params: Dict[str, str] = field(default_factory=dict) + + +@dataclass +class ProxyResponse: + status: int + headers: List[Tuple[str, str]] + content: bytes + + def json(self): + return json.loads(self.content.decode("utf-8")) + + +def client_headers(incoming, drop: Iterable[str] = ()) -> Dict[str, str]: + """Copy the caller's headers minus the ones we must not forward.""" + extra = {d.lower() for d in drop} + return { + k: v + for k, v in incoming.headers.items() + if k.lower() not in DROP_REQUEST_HEADERS and k.lower() not in extra + } + + +def filter_response_headers(headers) -> List[Tuple[str, str]]: + return [(k, v) for k, v in headers.items() if k.lower() not in DROP_RESPONSE_HEADERS] + + +class Provider: + name = "base" + + def __init__(self, cfg): + self.cfg = cfg + + def forward(self, req, subpath: str, body: bytes) -> ProxyResponse: + raise NotImplementedError + + +class HttpProvider(Provider): + """Providers that are a straight HTTP reverse-proxy (OpenAI, Anthropic).""" + + def target(self, req, subpath: str, body: bytes) -> UpstreamRequest: + raise NotImplementedError + + def forward(self, req, subpath, body): + up = self.target(req, subpath, body) + resp = requests.request( + up.method, + up.url, + headers=up.headers, + params=up.params, + data=body, + timeout=(self.cfg.connect_timeout, self.cfg.read_timeout), + ) + return ProxyResponse( + status=resp.status_code, + headers=filter_response_headers(resp.headers), + content=resp.content, + ) 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/ventis/llm_proxy/providers/openai.py b/ventis/llm_proxy/providers/openai.py new file mode 100644 index 0000000..67457eb --- /dev/null +++ b/ventis/llm_proxy/providers/openai.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from ventis.llm_proxy.providers.base import HttpProvider, UpstreamRequest, client_headers + + +class OpenAIProvider(HttpProvider): + name = "openai" + + def target(self, req, subpath, body): + headers = client_headers(req, drop=["authorization"]) + if self.cfg.openai.api_key: + headers["Authorization"] = f"Bearer {self.cfg.openai.api_key}" + return UpstreamRequest( + method=req.method, + url=f"{self.cfg.openai.upstream_base}/{subpath}", + headers=headers, + params=req.args.to_dict(flat=True), + ) diff --git a/ventis/llm_proxy/proxy.py b/ventis/llm_proxy/proxy.py new file mode 100644 index 0000000..ec0956d --- /dev/null +++ b/ventis/llm_proxy/proxy.py @@ -0,0 +1,58 @@ +"""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.controller.ventis_context as ventis_context +except ImportError: + # In-container the framework files are copied flat to /app. + try: + import ventis_context + except ImportError: + ventis_context = None + +log = logging.getLogger(__name__) + + +def _inject_ventis_headers(params=None, **kwargs): + """Inject X-Ventis-Future-ID into the outgoing Bedrock HTTP request. + + Registered on boto3's ``before-call.bedrock-runtime`` event, whose handlers + receive the prepared-request ``params`` dict (with a mutable ``headers``). + The ``request`` object only exists on the later ``before-send`` event, so + reading it here would always be None and silently drop the header. + """ + if not ventis_context or params is None: + return + + # Get current future_id from thread-local context + try: + future_id = ventis_context.get_current_future_id() + if future_id: + params.setdefault("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/ventis/llm_proxy/requirements.txt b/ventis/llm_proxy/requirements.txt new file mode 100644 index 0000000..2f7091c --- /dev/null +++ b/ventis/llm_proxy/requirements.txt @@ -0,0 +1,3 @@ +flask>=2.0 +requests>=2.28 +boto3>=1.28 diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index 3c13032..3edc970 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -17,10 +17,10 @@ import yaml # Packages every agent container needs regardless of its specific business logic. -BASE_AGENT_REQUIREMENTS = ["grpcio", "grpcio-tools", "redis", "pyyaml", "psutil", "boto3"] +BASE_AGENT_REQUIREMENTS = ["grpcio", "grpcio-tools", "redis", "pyyaml", "psutil", "boto3", "flask", "requests"] # Workflow will always require these -BASE_WORKFLOW_REQUIREMENTS = BASE_AGENT_REQUIREMENTS + ["flask", "sqlalchemy", "psycopg[binary]"] +BASE_WORKFLOW_REQUIREMENTS = BASE_AGENT_REQUIREMENTS + ["sqlalchemy", "psycopg[binary]"] def _build_import_nodes(): @@ -305,6 +305,23 @@ def _stub_destination(stub_file, stub_entrypoints): return basename +def _copy_llm_proxy(output_dir, script_dir): + """Copy the ventis.llm_proxy package into the build context as an importable + `ventis` package so the in-container proxy can run via `python -m ventis.llm_proxy`. + Its cross-package imports (redis_client, ventis_context) fall back to the flat + copies already placed at the context root.""" + 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.copy2( + os.path.join(script_dir, "__init__.py"), + os.path.join(output_dir, "ventis", "__init__.py"), + ) + + def _copy_files(output_dir, files_to_copy): """Copy each (src, dst) pair into output_dir, refusing to write outside it (e.g. via a symlinked destination parent).""" real_output_dir = os.path.realpath(output_dir) @@ -392,7 +409,6 @@ def generate_docker( os.path.join(script_dir, "controller", "utils", "gpu_metrics.py"), "gpu_metrics.py", ), - (os.path.join(script_dir, "controller", "bedrock.py"), "bedrock.py"), ] # Copy provided agent stubs, overwriting the swept real file at the same path @@ -408,6 +424,13 @@ def generate_docker( files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) _copy_files(output_dir, files_to_copy) + _copy_llm_proxy(output_dir, script_dir) + + # Copy the real agent entrypoint to the context root. + shutil.copy2( + os.path.abspath(agent_file), + os.path.join(output_dir, os.path.basename(agent_file)), + ) # Copy the real agent entrypoint to the context root. shutil.copy2( @@ -529,6 +552,13 @@ def generate_workflow_docker( files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) _copy_files(output_dir, files_to_copy) + _copy_llm_proxy(output_dir, script_dir) + + # Copy the real workflow entrypoint to the context root. + shutil.copy2( + os.path.abspath(workflow_file), + os.path.join(output_dir, workflow_basename), + ) # Copy the real workflow entrypoint to the context root. shutil.copy2(