diff --git a/cli/ARCHITECTURE.md b/cli/ARCHITECTURE.md new file mode 100644 index 0000000..62f869a --- /dev/null +++ b/cli/ARCHITECTURE.md @@ -0,0 +1,200 @@ +# CanyonOS CLI — Architecture + +## The one idea to keep in your head + +**The CLI does almost nothing. The container does everything.** + +`canyonos` is a thin client. It never builds, compiles, or runs your workflow +itself — it manages a **Global Controller (GC) container**, ships your project +into it, and drives it over a small HTTP API. Everything you see in your +terminal is the CLI *narrating* what the container is doing. + +If you remember only one picture, remember this: + +``` + YOU CLI (host) GLOBAL CONTROLLER (container) + │ │ │ + │ canyonos deploy │ │ + ├──────────────────────▶│ pull + run container │ + │ ├───────────────────────────────▶│ + │ │ copy project in (docker cp) │ + │ ├───────────────────────────────▶│ /workspace + │ │ POST /deploy │ + │ ├───────────────────────────────▶│ ventis build + launch + │ │◀── log stream (docker logs) ───┤ │ + │◀── readable progress ─┤ │ ▼ + │ │ spawns Redis + agents + │ │ (sibling containers) +``` + +--- + +## How the pieces connect + +``` +┌───────────────────────────── your machine ─────────────────────────────┐ +│ │ +│ ┌───────────┐ HTTP :8000 ┌──────────────────────────┐ │ +│ │ canyonos │ ───── /deploy /clean ────▶│ Global Controller │ │ +│ │ CLI │ /status /endpoints │ container │ │ +│ │ │ ───── docker cp ─────────▶│ ├─ /workspace (a copy │ │ +│ │ │ ───── docker logs -f ────▶│ │ of your project) │ │ +│ └─────┬─────┘ │ └─ runs `ventis` │ │ +│ │ └───────────┬──────────────┘ │ +│ │ docker compose │ docker.sock │ +│ ▼ ▼ (spawns siblings)│ +│ ┌───────────────────────┐ ┌───────────────────────────┐ │ +│ │ Dashboard stack │◀── traces ───│ Redis + your agent / │ │ +│ │ web · api · postgres │ (OTLP) │ workflow containers │ │ +│ └───────────────────────┘ └───────────────────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────────────────┘ +``` + +Three things worth internalizing about this diagram: + +1. **The container talks to the host Docker daemon.** The GC mounts the host's + `docker.sock`, so the Redis and agent/workflow containers it launches are + **siblings on your machine**, not nested inside it. (This is why teardown has + to be explicit — see `stop` vs `quit` below.) +2. **Your project is a *copy*, not a live mount.** Files are `docker cp`'d into + a named volume at `/workspace`. Editing files on the host after a deploy does + **not** reach the running build. +3. **The dashboard is separate.** It's its own compose stack that just *renders* + the OTLP traces your workflow emits — it isn't in the deploy critical path. + +State connecting the CLI to its container is a single file: +`~/.canyonos/state.json` (container id + port). Every command that needs the +container reads it. + +--- + +## The three commands that matter + +### `build` — get your code into CanyonOS shape + +``` + you ──▶ canyonos build ──▶ pick agent (Claude/Codex) + scope + └─▶ fetch the porting skill from GitHub + └─▶ launch your coding agent with it + │ + ▼ + generates .car/ ◀── canyonos-formatted project + (originals untouched) +``` + +A **host-side, agent-driven** step. The CLI installs the CanyonOS porting skill +onto your coding agent and hands it a prompt; the agent produces a `.car/` +folder — the canyonos-ready version of your project plus its config. **No +container is involved yet.** + +### `deploy` — the main path + +``` + canyonos deploy + │ + ├─ 1. start fresh → ensure Docker up, tear down any old controller, + │ pull + run the GC container, save state (previous canyonos init) + │ + ├─ 2. ship code → docker cp your project into /workspace + │ + ├─ 3. trigger → POST /deploy (container runs `ventis`: + │ build stubs/images + launch the workflow) + │ + └─ 4. narrate → tail container logs, boil them down to phases, + and when the workflow reports "up": + • auto-start the dashboard (canyonos serve) + • print where everything lives +``` + +Everything after step 3 happens *inside* the container. The CLI's real job in +step 4 is turning a very noisy log stream (a full image-build transcript, etc.) +into a short, readable progression — and, on failure, revealing the part it had +been hiding so you can see the actual cause. + +When it finishes you get a summary panel: the **dashboard URL** and each +**workflow endpoint** (`POST /main`), using the real address the container +placed the workflow at. + +``` + ┌─ Deploy is live ─────────────────────────────┐ + │ Dashboard http://127.0.0.1:8080 │ + │ POST http://127.0.0.1:8000/main │ + │ body {"query": "..."} │ + └──────────────────────────────────────────────┘ +``` + +### `config` — view or edit settings + +``` + canyonos config ──▶ View → pretty tables of agents / otel / general + └─▶ Change → interactive editor (comments & order preserved) +``` + +The important mental model isn't the editor — it's **what a change costs you**: +canyonos config only allows you to change the config file, changing the source code requires a redeploy. + +``` + change type takes effect by... + ─────────────────────── ─────────────────────────────────── + config value only reloads in place (no rebuild) + workflow *code* changes full redeploy (container holds a copy) +``` + +--- + +## Lifecycle: what stays and what goes + +Because the deploy spawns real sibling containers, "make it stop" has two levels of "stop": + +``` + deploy sibling GC project files + stops? containers? container? (volume)? + ─────────────── ─────── ─────────── ───────── ───────────── + canyonos stop ✅ ✅ keep keep + canyonos quit ✅ ✅ remove remove +``` + +- **`stop`** — pause the show, keep the stage set. Redeploy without re-pulling. +- **`quit`** — full teardown. Removes the container *and* the `/workspace` + volume (your copied files). Every `deploy` quietly does this to any previous + controller, so each deploy starts clean. + +And to observe without changing anything: + +- **`logs`** — re-attach to the same live log stream `deploy` shows. Useful + after you Ctrl+C out of a deploy: the deploy keeps running; you just stopped + *watching*. (Ctrl+C on `logs` likewise only detaches.) + +``` + deploy ──▶ (Ctrl+C) ──▶ still running in the container + │ ▲ + └── logs ─────────────────┘ re-attach anytime +``` + +--- + +## The whole loop, one screen + +``` + cd your-project + │ + ▼ + build port your code → .car/ (opens your coding agent) + │ + ▼ + deploy build + launch in the container (dashboard opens itself) + │ + ├─ status where does the workflow answer? + ├─ config tweak settings (live reload); redeploy for code changes + ├─ logs re-attach to the stream + │ + ▼ + stop halt the deploy, keep container + files + or + quit full teardown, remove everything +``` + +That's the entire system: a thin CLI, one container that does the heavy +lifting, a pile of sibling containers it spawns, and a dashboard watching the +whole thing. diff --git a/cli/README.md b/cli/README.md index 76cd4e1..f708d62 100644 --- a/cli/README.md +++ b/cli/README.md @@ -2,6 +2,12 @@ Lightweight CLI for CanyonOS Serves as a thin API layer, connecting to the global controller container. +## Architecture + +For a full walkthrough of the `build`, `deploy`, and `config` flows — plus how +`logs`, `stop`, and `quit` fit into the container lifecycle — see +[ARCHITECTURE.md](ARCHITECTURE.md). + ## Serve `canyonos serve` starts the local CanyonOS dashboard against the Postgres bundled in its compose @@ -14,15 +20,3 @@ Need uv or pip Need docker and docker compose 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 -### To Republish to PyPi - -```Terminal -cd cli -# Go into, pyproject.toml, and increment version number -rm -rf dist/ # Removes the old distro, causes conflicts - -uv build -uv publish # Needs PyPi Auth Token, ask Saaketh -``` \ No newline at end of file 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 dbff7bb..25100a6 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 8278401..dd3ff5f 100644 --- a/uv.lock +++ b/uv.lock @@ -1238,6 +1238,7 @@ dependencies = [ { name = "psycopg", extra = ["binary"] }, { name = "pyyaml" }, { name = "redis" }, + { name = "requests" }, { name = "sqlalchemy" }, ] @@ -1261,6 +1262,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 1108be1..6f8fc14 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -19,10 +19,10 @@ # 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", "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(): @@ -307,6 +307,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) @@ -394,7 +411,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 @@ -410,6 +426,7 @@ 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( @@ -531,6 +548,7 @@ 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(