mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-26 06:52:57 +08:00
docs(examples): add opencode privacy-redaction demo
This commit is contained in:
@@ -0,0 +1,102 @@
|
||||
# opencode + privacy-preserving image redaction
|
||||
|
||||
Run [opencode](https://opencode.ai) inside an OpenShell sandbox where image payloads are kept away from the public inference provider and recovered, on demand, through a private local model.
|
||||
|
||||
opencode talks to `inference.local`, which is routed to a host-side OpenAI-compatible **redaction proxy**. The proxy decides per request, from a content policy, whether to let an image through to the public model or to strip it and tell opencode to recover the data via a **local-model MCP server** running on the same host.
|
||||
|
||||
```
|
||||
opencode (sandbox)
|
||||
├── chat ── inference.local ──▶ redaction proxy (host :8000) ──▶ public model (gpt-5.5)
|
||||
│ │ if content_policy image: redact
|
||||
│ └─ strips image, injects "use the local_model MCP"
|
||||
└── tool ── local_model MCP (host :8001) ──▶ private/local model (e.g. LM Studio :1234)
|
||||
```
|
||||
|
||||
The "demo magic": the proxy reads an OpenShell-style policy file (`policies/content-policy.yaml`) on **every** request. OpenShell itself does not interpret the `content_policy` block (that would need code changes) - only this proxy does. Flip one line and the next image request changes routing, no restart.
|
||||
|
||||
## Layout
|
||||
|
||||
| Path | What |
|
||||
|------|------|
|
||||
| `proxy/` | The redaction proxy + `call_inference` MCP server (a `uv` project). |
|
||||
| `proxy/Makefile` | Start/stop the proxy and MCP; flip the demo policy. |
|
||||
| `policies/opencode-policy.yaml` | Sandbox policy you give to OpenShell at create time. |
|
||||
| `policies/content-policy.yaml` | Copy of the policy **plus** the `content_policy` block the proxy reads. |
|
||||
| `sandbox/opencode.json` | opencode config: `inference.local` provider + `local_model` MCP. |
|
||||
| `sandbox/opencode-provision.sh` | Writes the config + workspace inside the sandbox. |
|
||||
| `sandbox/setup-sandbox.sh` | One command: create sandbox, provision, run web UI, expose it. |
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- An OpenShell gateway registered and selected: `openshell gateway list` (the active one must be reachable as `host.openshell.internal` from sandboxes - this is the default for the local podman/VM gateway).
|
||||
- A private/local OpenAI-compatible model for the MCP to call (the example defaults to LM Studio at `http://localhost:1234`).
|
||||
- An `NVIDIA_API_KEY` for the public upstream (`https://inference-api.nvidia.com`).
|
||||
- `uv` installed.
|
||||
|
||||
## 1. Start the host services
|
||||
|
||||
```shell
|
||||
cd proxy
|
||||
cp .env.example .env # then put your NVIDIA_API_KEY in .env
|
||||
uv sync
|
||||
make start # starts the MCP (:8001) and proxy (:8000)
|
||||
```
|
||||
|
||||
The proxy and MCP bind to `0.0.0.0` so the sandbox can reach them via `host.openshell.internal`.
|
||||
|
||||
## 2. One-time gateway wiring (inference.local -> the proxy)
|
||||
|
||||
```shell
|
||||
openshell provider create --name local-proxy --type openai \
|
||||
--credential OPENAI_API_KEY=dummy \
|
||||
--config OPENAI_BASE_URL=http://host.openshell.internal:8000/v1
|
||||
|
||||
openshell inference set --provider local-proxy --model openai/openai/gpt-5.5 --no-verify
|
||||
```
|
||||
|
||||
`--no-verify` because host-side validation cannot resolve `host.openshell.internal`; runtime egress originates inside the sandbox and works.
|
||||
|
||||
## 3. Create the opencode sandbox
|
||||
|
||||
```shell
|
||||
cd sandbox
|
||||
./setup-sandbox.sh
|
||||
```
|
||||
|
||||
This creates the `opencode` sandbox with `policies/opencode-policy.yaml`, writes `opencode.json`, starts `opencode web` on loopback inside the sandbox, and exposes it. Open the printed URL (e.g. `http://opencode--web.openshell.localhost:<gateway-port>/`).
|
||||
|
||||
Sanity check from the CLI:
|
||||
|
||||
```shell
|
||||
ssh -F /tmp/opencode-ssh-config openshell-opencode \
|
||||
'export HOME=/sandbox; cd /sandbox/workspace; opencode run --model inference-local/openai/openai/gpt-5.5 "say hi"'
|
||||
```
|
||||
|
||||
## 4. Demo: flip the content policy
|
||||
|
||||
Default state ships as `redact`. Watch the proxy log while you drive opencode from the web UI:
|
||||
|
||||
```shell
|
||||
cd proxy && make tail-proxy-log
|
||||
```
|
||||
|
||||
- **Allow** (`make demo-allow`): attach an image in opencode and ask about it. The proxy logs `image_redaction_skipped reason=allowed_by_policy` and the image reaches the public model.
|
||||
- **Redact** (`make demo-redact`): re-send. The proxy logs `image_redaction_applied`, the image is replaced with the `redact_description`, and opencode calls the `local_model` MCP tool to recover the data from the private model. The image never reaches the public model.
|
||||
|
||||
`make demo-status` shows the current rule. The proxy re-reads the file per request, so no restart is needed between toggles.
|
||||
|
||||
Request/response captures land in `proxy/request-logs/` (gitignored) for inspection.
|
||||
|
||||
## Tear down
|
||||
|
||||
```shell
|
||||
cd proxy && make stop
|
||||
openshell service delete opencode web
|
||||
openshell sandbox delete opencode
|
||||
```
|
||||
|
||||
## Tests
|
||||
|
||||
```shell
|
||||
cd proxy && make test
|
||||
```
|
||||
@@ -0,0 +1,64 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Demo "content policy" file.
|
||||
#
|
||||
# This is a COPY of opencode-policy.yaml with one extra top-level block:
|
||||
# `content_policy`. The OpenShell policy engine does not interpret
|
||||
# `content_policy` today (it would need code changes); only the redaction proxy
|
||||
# in this example reads it. The proxy re-reads this file on EVERY request, so it
|
||||
# is a live toggle:
|
||||
#
|
||||
# content_type.image: allow -> images pass through to the public model
|
||||
# content_type.image: redact -> images are stripped before the public model,
|
||||
# and `redact_description` is injected at the
|
||||
# redaction site telling the model to recover
|
||||
# the data via the local-model MCP tool.
|
||||
#
|
||||
# Flip the one line below and re-send an image from opencode to see the routing
|
||||
# change with no restart.
|
||||
|
||||
version: 1
|
||||
|
||||
filesystem_policy:
|
||||
include_workdir: true
|
||||
read_only:
|
||||
- /usr
|
||||
- /lib
|
||||
- /proc
|
||||
- /dev/urandom
|
||||
- /app
|
||||
- /etc
|
||||
- /var/log
|
||||
read_write:
|
||||
- /sandbox
|
||||
- /tmp
|
||||
- /dev/null
|
||||
- /home
|
||||
- /root
|
||||
- /dev/pts
|
||||
|
||||
landlock:
|
||||
compatibility: best_effort
|
||||
|
||||
process:
|
||||
run_as_user: sandbox
|
||||
run_as_group: sandbox
|
||||
|
||||
network_policies:
|
||||
opencode_runtime:
|
||||
name: opencode_runtime
|
||||
endpoints:
|
||||
- { host: host.openshell.internal, port: 8001 }
|
||||
- { host: registry.npmjs.org, port: 443 }
|
||||
- { host: "*.npmjs.org", port: 443 }
|
||||
- { host: models.dev, port: 443 }
|
||||
- { host: app.opencode.ai, port: 443 }
|
||||
binaries:
|
||||
- { path: "**" }
|
||||
|
||||
# --- demo magic: interpreted by the redaction proxy, ignored by OpenShell ---
|
||||
content_policy:
|
||||
content_type:
|
||||
image: redact
|
||||
redact_description: Image has been redacted, use a local model through the `local_model` MCP server.
|
||||
@@ -0,0 +1,51 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
version: 1
|
||||
|
||||
filesystem_policy:
|
||||
include_workdir: true
|
||||
read_only:
|
||||
- /usr
|
||||
- /lib
|
||||
- /proc
|
||||
- /dev/urandom
|
||||
- /app
|
||||
- /etc
|
||||
- /var/log
|
||||
read_write:
|
||||
- /sandbox
|
||||
- /tmp
|
||||
- /dev/null
|
||||
- /home
|
||||
- /root
|
||||
# PTY allocation for the opencode TUI / web terminal and bash tool.
|
||||
# /dev/ptmx is a symlink into pts/ptmx, so granting the devpts mount covers it.
|
||||
- /dev/pts
|
||||
|
||||
landlock:
|
||||
compatibility: best_effort
|
||||
|
||||
process:
|
||||
run_as_user: sandbox
|
||||
run_as_group: sandbox
|
||||
|
||||
network_policies:
|
||||
# opencode runtime egress. inference.local is ALWAYS reachable regardless of
|
||||
# network policy (the supervisor routes it), so it is intentionally absent here.
|
||||
opencode_runtime:
|
||||
name: opencode_runtime
|
||||
endpoints:
|
||||
# Redact MCP server, running on the host (same host as the inference proxy),
|
||||
# reached from inside the sandbox via the host alias on its HTTP port.
|
||||
- { host: host.openshell.internal, port: 8001 }
|
||||
# opencode fetches the openai-compatible provider SDK (@ai-sdk/openai-compatible)
|
||||
# and the models.dev metadata snapshot on first run.
|
||||
- { host: registry.npmjs.org, port: 443 }
|
||||
- { host: "*.npmjs.org", port: 443 }
|
||||
- { host: models.dev, port: 443 }
|
||||
# `opencode web` serves its SPA front-end from app.opencode.ai. Without this
|
||||
# the local web server returns 403 (egress proxy denies the asset fetch).
|
||||
- { host: app.opencode.ai, port: 443 }
|
||||
binaries:
|
||||
- { path: "**" }
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copy to .env and fill in NVIDIA_API_KEY. Loaded automatically by the proxy/MCP.
|
||||
NVIDIA_API_KEY=your-nvapikey
|
||||
ENV_FILE=.env
|
||||
|
||||
# Public upstream the proxy forwards to when an image is NOT redacted.
|
||||
# This is what inference.local resolves to from the sandbox.
|
||||
UPSTREAM_BASE_URL=https://inference-api.nvidia.com
|
||||
UPSTREAM_MODEL=openai/openai/gpt-5.5
|
||||
|
||||
# Private/local model the MCP `call_inference` tool calls (e.g. LM Studio).
|
||||
PRIVATE_UPSTREAM_BASE_URL=http://localhost:1234
|
||||
PRIVATE_UPSTREAM_MODEL=nvidia/nemotron-3-nano-omni
|
||||
PRIVATE_UPSTREAM_API_KEY=
|
||||
MCP_FORCE_UPSTREAM_MODEL=true
|
||||
|
||||
# Bind to 0.0.0.0 so the podman/VM-hosted sandbox can reach the host via
|
||||
# host.openshell.internal.
|
||||
PROXY_HOST=0.0.0.0
|
||||
PROXY_PORT=8000
|
||||
MCP_HOST=0.0.0.0
|
||||
MCP_PORT=8001
|
||||
MCP_TRANSPORT=http
|
||||
|
||||
PROXY_LOG_LEVEL=INFO
|
||||
REQUEST_LOG_DIR=request-logs
|
||||
|
||||
# Demo magic: the proxy reads this OpenShell-style policy file on every request
|
||||
# and redacts images only when content_policy.content_type.image == redact.
|
||||
# (REDACTION_ENABLED is the fallback when CONTENT_POLICY_FILE is unset.)
|
||||
CONTENT_POLICY_FILE=../policies/content-policy.yaml
|
||||
REDACTION_ENABLED=false
|
||||
|
||||
# Store redacted images so the MCP can re-load them via image_reference.
|
||||
STORE_REDACTED_IMAGES=true
|
||||
REDACTED_IMAGE_DIR=/tmp/redact-spike-images
|
||||
@@ -0,0 +1,14 @@
|
||||
.venv/
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
.pytest_cache/
|
||||
.env
|
||||
.env.*
|
||||
!.env.example
|
||||
uv.lock
|
||||
redact-spike-images/
|
||||
request-logs/
|
||||
proxy.log
|
||||
proxy.pid
|
||||
mcp.log
|
||||
mcp.pid
|
||||
@@ -0,0 +1,67 @@
|
||||
ENV_FILE ?= $(if $(wildcard .env),.env,.env.example)
|
||||
CONTENT_POLICY ?= ../policies/content-policy.yaml
|
||||
|
||||
.PHONY: start-proxy stop-proxy restart-proxy tail-proxy-log \
|
||||
start-mcp stop-mcp restart-mcp tail-mcp-log \
|
||||
start stop demo-allow demo-redact demo-status test
|
||||
|
||||
## --- proxy (OpenAI-compatible, inference.local upstream + image redaction) ---
|
||||
start-proxy:
|
||||
@if [ -f proxy.pid ] && kill -0 "$$(cat proxy.pid)" 2>/dev/null; then \
|
||||
echo "proxy already running: pid=$$(cat proxy.pid)"; \
|
||||
else \
|
||||
rm -f proxy.log proxy.pid; \
|
||||
ENV_FILE="$(ENV_FILE)" uv run redact-proxy > proxy.log 2>&1 & \
|
||||
echo $$! > proxy.pid; \
|
||||
echo "proxy started: pid=$$(cat proxy.pid), log=proxy.log, env=$(ENV_FILE)"; \
|
||||
fi
|
||||
|
||||
stop-proxy:
|
||||
@if [ -f proxy.pid ] && kill -0 "$$(cat proxy.pid)" 2>/dev/null; then \
|
||||
kill "$$(cat proxy.pid)"; echo "proxy stopped: pid=$$(cat proxy.pid)"; rm -f proxy.pid; \
|
||||
else echo "proxy not running"; rm -f proxy.pid; fi
|
||||
|
||||
restart-proxy: stop-proxy start-proxy
|
||||
|
||||
tail-proxy-log:
|
||||
tail -f proxy.log
|
||||
|
||||
## --- MCP (call_inference -> private/local model) ---
|
||||
start-mcp:
|
||||
@if [ -f mcp.pid ] && kill -0 "$$(cat mcp.pid)" 2>/dev/null; then \
|
||||
echo "mcp already running: pid=$$(cat mcp.pid)"; \
|
||||
else \
|
||||
rm -f mcp.log mcp.pid; \
|
||||
ENV_FILE="$(ENV_FILE)" uv run redact-mcp > mcp.log 2>&1 & \
|
||||
echo $$! > mcp.pid; \
|
||||
echo "mcp started: pid=$$(cat mcp.pid), log=mcp.log, env=$(ENV_FILE)"; \
|
||||
fi
|
||||
|
||||
stop-mcp:
|
||||
@if [ -f mcp.pid ] && kill -0 "$$(cat mcp.pid)" 2>/dev/null; then \
|
||||
kill "$$(cat mcp.pid)"; echo "mcp stopped: pid=$$(cat mcp.pid)"; rm -f mcp.pid; \
|
||||
else echo "mcp not running"; rm -f mcp.pid; fi
|
||||
|
||||
restart-mcp: stop-mcp start-mcp
|
||||
|
||||
tail-mcp-log:
|
||||
tail -f mcp.log
|
||||
|
||||
## --- convenience ---
|
||||
start: start-mcp start-proxy
|
||||
stop: stop-proxy stop-mcp
|
||||
|
||||
## --- demo magic: flip the content policy live (proxy re-reads per request) ---
|
||||
demo-allow:
|
||||
@sed -i.bak -E 's/^( *image: *).*/\1allow/' "$(CONTENT_POLICY)" && rm -f "$(CONTENT_POLICY).bak"
|
||||
@echo "content policy -> image: allow (images pass through to the public model)"
|
||||
|
||||
demo-redact:
|
||||
@sed -i.bak -E 's/^( *image: *).*/\1redact/' "$(CONTENT_POLICY)" && rm -f "$(CONTENT_POLICY).bak"
|
||||
@echo "content policy -> image: redact (images stripped; routed via local-model MCP)"
|
||||
|
||||
demo-status:
|
||||
@grep -E '^ *image:' "$(CONTENT_POLICY)" || echo "no image rule found in $(CONTENT_POLICY)"
|
||||
|
||||
test:
|
||||
uv run pytest
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
[project]
|
||||
name = "redact-spike"
|
||||
version = "0.1.0"
|
||||
description = "OpenAI-compatible inference proxy with policy-driven image redaction and MCP fallback."
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"fastapi>=0.115.0",
|
||||
"fastmcp>=2.0.0",
|
||||
"httpx>=0.28.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
"pydantic>=2.10.0",
|
||||
"pyyaml>=6.0",
|
||||
"uvicorn>=0.32.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
redact-proxy = "redact_spike.proxy:main"
|
||||
redact-private-proxy = "redact_spike.private_switch_proxy:main"
|
||||
redact-mcp = "redact_spike.mcp_server:main"
|
||||
redact-mcp-call = "redact_spike.mcp_debug:main"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=8.3.0",
|
||||
"pytest-asyncio>=0.24.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
@@ -0,0 +1,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
__all__ = ["__version__"]
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,75 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
|
||||
def _env_bool(name: str, default: bool) -> bool:
|
||||
value = os.getenv(name)
|
||||
if value is None:
|
||||
return default
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Settings:
|
||||
nvidia_api_key: str | None = None
|
||||
upstream_base_url: str = "https://inference-api.nvidia.com"
|
||||
upstream_model: str = "openai/openai/gpt-5.5"
|
||||
private_upstream_base_url: str = "http://localhost:1234"
|
||||
private_upstream_model: str = "qwen/qwen3.6-27b"
|
||||
private_upstream_api_key: str | None = None
|
||||
private_switch_force_model: bool = True
|
||||
mcp_force_upstream_model: bool = True
|
||||
mcp_transport: str = "http"
|
||||
mcp_host: str = "127.0.0.1"
|
||||
mcp_port: int = 8001
|
||||
proxy_host: str = "127.0.0.1"
|
||||
proxy_port: int = 8000
|
||||
private_proxy_port: int = 8002
|
||||
proxy_log_level: str = "INFO"
|
||||
request_log_dir: Path = Path("request-logs")
|
||||
redaction_enabled: bool = True
|
||||
store_redacted_images: bool = True
|
||||
redacted_image_dir: Path = Path(tempfile.gettempdir()) / "redact-spike-images"
|
||||
content_policy_file: Path | None = None
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "Settings":
|
||||
_load_dotenv_file()
|
||||
content_policy_file = os.getenv("CONTENT_POLICY_FILE")
|
||||
return cls(
|
||||
nvidia_api_key=os.getenv("NVIDIA_API_KEY"),
|
||||
upstream_base_url=os.getenv("UPSTREAM_BASE_URL", cls.upstream_base_url),
|
||||
upstream_model=os.getenv("UPSTREAM_MODEL", cls.upstream_model),
|
||||
private_upstream_base_url=os.getenv("PRIVATE_UPSTREAM_BASE_URL", cls.private_upstream_base_url),
|
||||
private_upstream_model=os.getenv("PRIVATE_UPSTREAM_MODEL", cls.private_upstream_model),
|
||||
private_upstream_api_key=os.getenv("PRIVATE_UPSTREAM_API_KEY"),
|
||||
private_switch_force_model=_env_bool("PRIVATE_SWITCH_FORCE_MODEL", cls.private_switch_force_model),
|
||||
mcp_force_upstream_model=_env_bool("MCP_FORCE_UPSTREAM_MODEL", cls.mcp_force_upstream_model),
|
||||
mcp_transport=os.getenv("MCP_TRANSPORT", cls.mcp_transport),
|
||||
mcp_host=os.getenv("MCP_HOST", cls.mcp_host),
|
||||
mcp_port=int(os.getenv("MCP_PORT", str(cls.mcp_port))),
|
||||
proxy_host=os.getenv("PROXY_HOST", cls.proxy_host),
|
||||
proxy_port=int(os.getenv("PROXY_PORT", str(cls.proxy_port))),
|
||||
private_proxy_port=int(os.getenv("PRIVATE_PROXY_PORT", str(cls.private_proxy_port))),
|
||||
proxy_log_level=os.getenv("PROXY_LOG_LEVEL", cls.proxy_log_level),
|
||||
request_log_dir=Path(os.getenv("REQUEST_LOG_DIR", str(cls.request_log_dir))),
|
||||
redaction_enabled=_env_bool("REDACTION_ENABLED", cls.redaction_enabled),
|
||||
store_redacted_images=_env_bool("STORE_REDACTED_IMAGES", cls.store_redacted_images),
|
||||
redacted_image_dir=Path(os.getenv("REDACTED_IMAGE_DIR", str(cls.redacted_image_dir))),
|
||||
content_policy_file=Path(content_policy_file) if content_policy_file else None,
|
||||
)
|
||||
|
||||
|
||||
def _load_dotenv_file() -> None:
|
||||
env_file = Path(os.getenv("ENV_FILE", ".env"))
|
||||
if env_file.exists():
|
||||
load_dotenv(env_file, override=True)
|
||||
@@ -0,0 +1,72 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageContentRule:
|
||||
"""Resolved `content_policy.content_type` decision for image payloads.
|
||||
|
||||
`action` is "redact" or "allow". `description` is the operator-authored note
|
||||
that the proxy substitutes at the redaction site (the `redact_description`
|
||||
field), telling the main model how to recover the data via the MCP tool.
|
||||
"""
|
||||
|
||||
action: str
|
||||
description: str | None = None
|
||||
|
||||
@property
|
||||
def should_redact(self) -> bool:
|
||||
return self.action == "redact"
|
||||
|
||||
|
||||
# Default when no content policy is in force: images pass through untouched.
|
||||
ALLOW = ImageContentRule(action="allow")
|
||||
|
||||
|
||||
def load_image_content_rule(path: Path | None) -> ImageContentRule | None:
|
||||
"""Read the image content rule from a policy file, re-read on every call.
|
||||
|
||||
Returns None when no decision is expressed (no path, missing file, no
|
||||
`content_policy.content_type.image` key) so the caller can fall back to its
|
||||
own default. This is what makes the demo a live toggle: edit the file and
|
||||
the next request observes the new rule, no restart required.
|
||||
|
||||
The `content_policy` block is an OpenShell-style extension that the OpenShell
|
||||
policy engine ignores today; only this proxy interprets it.
|
||||
"""
|
||||
if path is None:
|
||||
return None
|
||||
|
||||
if not path.exists():
|
||||
logger.info("content_policy_absent path=%s -> no rule", path)
|
||||
return None
|
||||
|
||||
try:
|
||||
document = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
except (OSError, yaml.YAMLError) as error:
|
||||
logger.warning("content_policy_unreadable path=%s error=%s -> no rule", path, error)
|
||||
return None
|
||||
|
||||
content_type = (((document or {}).get("content_policy") or {}).get("content_type")) or {}
|
||||
action = content_type.get("image")
|
||||
if not isinstance(action, str):
|
||||
return None
|
||||
|
||||
description = content_type.get("redact_description")
|
||||
rule = ImageContentRule(
|
||||
action=action.strip().lower(),
|
||||
description=description if isinstance(description, str) else None,
|
||||
)
|
||||
logger.info("content_policy_loaded path=%s image_action=%s", path, rule.action)
|
||||
return rule
|
||||
@@ -0,0 +1,30 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
from redact_spike.mcp_server import call_inference_impl
|
||||
|
||||
|
||||
async def _run(payload: dict[str, Any]) -> None:
|
||||
response = await call_inference_impl(payload)
|
||||
print(json.dumps(response, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Call the MCP inference implementation directly for debugging.")
|
||||
parser.add_argument("--payload", help="OpenAI-style JSON payload. Reads stdin when omitted.")
|
||||
args = parser.parse_args()
|
||||
raw_payload = args.payload if args.payload is not None else sys.stdin.read()
|
||||
payload = json.loads(raw_payload)
|
||||
asyncio.run(_run(payload))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,156 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import datetime as dt
|
||||
import json
|
||||
import logging
|
||||
import mimetypes
|
||||
import base64
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from redact_spike.config import Settings
|
||||
from redact_spike.upstream import NvidiaInferenceClient
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def call_inference_impl(payload: dict[str, Any], settings: Settings | None = None) -> dict[str, Any]:
|
||||
settings = settings or Settings.from_env()
|
||||
call_id = uuid.uuid4().hex[:8]
|
||||
incoming_path = _write_mcp_capture(settings.request_log_dir, call_id, "mcp-incoming-request", payload)
|
||||
outbound_payload = normalize_mcp_payload(payload, settings)
|
||||
outbound_path = _write_mcp_capture(settings.request_log_dir, call_id, "mcp-normalized-request", outbound_payload)
|
||||
logger.info(
|
||||
"mcp_call_inference call_id=%s model=%s original_model=%s message_count=%s incoming=%s normalized=%s",
|
||||
call_id,
|
||||
outbound_payload.get("model"),
|
||||
payload.get("model"),
|
||||
len(outbound_payload.get("messages", [])),
|
||||
incoming_path,
|
||||
outbound_path,
|
||||
)
|
||||
client = NvidiaInferenceClient(
|
||||
settings,
|
||||
base_url=settings.private_upstream_base_url,
|
||||
model=settings.private_upstream_model,
|
||||
api_key=settings.private_upstream_api_key,
|
||||
require_api_key=False,
|
||||
)
|
||||
response = await client.chat_completions(outbound_payload)
|
||||
response_path = _write_mcp_capture(
|
||||
settings.request_log_dir,
|
||||
call_id,
|
||||
"mcp-upstream-response",
|
||||
{"status_code": response.status_code, "body": response.body},
|
||||
)
|
||||
logger.info("mcp_call_inference_response call_id=%s status=%s response=%s", call_id, response.status_code, response_path)
|
||||
return {"status_code": response.status_code, "body": response.body}
|
||||
|
||||
|
||||
def normalize_mcp_payload(payload: dict[str, Any], settings: Settings) -> dict[str, Any]:
|
||||
normalized = copy.deepcopy(payload)
|
||||
if settings.mcp_force_upstream_model or not normalized.get("model"):
|
||||
normalized["model"] = settings.private_upstream_model
|
||||
|
||||
messages = normalized.get("messages")
|
||||
if isinstance(messages, list):
|
||||
normalized["messages"] = [_normalize_message(message) for message in messages]
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalize_message(message: Any) -> Any:
|
||||
if not isinstance(message, dict):
|
||||
return message
|
||||
|
||||
normalized = dict(message)
|
||||
content = normalized.get("content")
|
||||
if isinstance(content, list):
|
||||
normalized["content"] = [_normalize_content_part(part) for part in content]
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalize_content_part(part: Any) -> Any:
|
||||
if isinstance(part, str):
|
||||
return {"type": "text", "text": part}
|
||||
if not isinstance(part, dict):
|
||||
return {"type": "text", "text": f"Unsupported content part: {json.dumps(part, sort_keys=True)}"}
|
||||
|
||||
part_type = part.get("type")
|
||||
if part_type in {"text", "image_url", "input_image"}:
|
||||
return part
|
||||
if part_type == "image_reference":
|
||||
return _image_reference_to_image_url(part)
|
||||
if part_type == "image_placeholder":
|
||||
placeholder_id = str(part.get("id", "UNKNOWN_IMAGE"))
|
||||
return {
|
||||
"type": "text",
|
||||
"text": (
|
||||
f"Image placeholder [{placeholder_id}] was referenced, but no image bytes or local image reference were "
|
||||
"provided to this MCP call. Retry with a text part plus an image_reference part using the local path "
|
||||
"from the redaction context."
|
||||
),
|
||||
}
|
||||
return {"type": "text", "text": f"Unsupported content part: {json.dumps(part, sort_keys=True)}"}
|
||||
|
||||
|
||||
def _image_reference_to_image_url(part: dict[str, Any]) -> dict[str, Any]:
|
||||
reference = part.get("path") or part.get("reference") or part.get("url")
|
||||
if not isinstance(reference, str):
|
||||
return {"type": "text", "text": "Invalid image_reference part: missing path, reference, or url."}
|
||||
|
||||
path = Path(reference.removeprefix("file://"))
|
||||
if not path.exists() or not path.is_file():
|
||||
return {"type": "text", "text": f"Invalid image_reference part: file does not exist at {reference}."}
|
||||
|
||||
mime_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
|
||||
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
|
||||
return {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{encoded}"}}
|
||||
|
||||
|
||||
def create_mcp_server(settings: Settings | None = None) -> FastMCP:
|
||||
mcp = FastMCP("redacted-inference")
|
||||
|
||||
@mcp.tool
|
||||
async def call_inference(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Call the configured OpenAI-compatible inference endpoint.
|
||||
|
||||
Provide an OpenAI-style payload. The server will use the configured model, so do not invent model names.
|
||||
If an image was redacted, include a user message whose content has both a text part and an image part.
|
||||
Use {"type":"image_reference","path":"<local path from redaction context>"} when the redaction context provides a local image path.
|
||||
A placeholder such as REDACTED_IMAGE_1 alone is only an identifier, not image data.
|
||||
"""
|
||||
return await call_inference_impl(payload, settings=settings)
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
mcp = create_mcp_server()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
settings = Settings.from_env()
|
||||
logging.basicConfig(level=settings.proxy_log_level.upper())
|
||||
if settings.mcp_transport == "stdio":
|
||||
mcp.run()
|
||||
else:
|
||||
mcp.run(transport=settings.mcp_transport, host=settings.mcp_host, port=settings.mcp_port)
|
||||
|
||||
|
||||
def _write_mcp_capture(log_dir: Path, call_id: str, label: str, payload: Any) -> str:
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = dt.datetime.now(dt.UTC).strftime("%Y%m%dT%H%M%S.%fZ")
|
||||
path = log_dir / f"{timestamp}-{call_id}-{label}.json"
|
||||
path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
|
||||
return str(path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,22 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uvicorn
|
||||
|
||||
from redact_spike.config import Settings
|
||||
from redact_spike.proxy import create_app
|
||||
|
||||
|
||||
app = create_app(private_model_switch=True)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
settings = Settings.from_env()
|
||||
uvicorn.run(
|
||||
"redact_spike.private_switch_proxy:app",
|
||||
host=settings.proxy_host,
|
||||
port=settings.private_proxy_port,
|
||||
reload=False,
|
||||
)
|
||||
@@ -0,0 +1,363 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
from redact_spike.config import Settings
|
||||
from redact_spike.content_policy import load_image_content_rule
|
||||
from redact_spike.redaction import DiskImageStore, last_message_has_image, redact_images, summarize_payload
|
||||
from redact_spike.upstream import NvidiaInferenceClient, UpstreamResponse
|
||||
|
||||
|
||||
class ChatCompletionsClient(Protocol):
|
||||
async def chat_completions(self, payload: dict[str, Any]) -> UpstreamResponse:
|
||||
pass
|
||||
|
||||
async def stream_chat_completions(self, payload: dict[str, Any]) -> AsyncIterator[bytes]:
|
||||
pass
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_app(
|
||||
settings: Settings | None = None,
|
||||
upstream_client: ChatCompletionsClient | None = None,
|
||||
private_upstream_client: ChatCompletionsClient | None = None,
|
||||
private_model_switch: bool = False,
|
||||
) -> FastAPI:
|
||||
settings = settings or Settings.from_env()
|
||||
logging.basicConfig(level=settings.proxy_log_level.upper())
|
||||
|
||||
app = FastAPI(title="Redacted Inference Proxy", version="0.1.0")
|
||||
app.state.settings = settings
|
||||
app.state.upstream_client = upstream_client or NvidiaInferenceClient(settings)
|
||||
app.state.private_upstream_client = private_upstream_client or NvidiaInferenceClient(
|
||||
settings,
|
||||
base_url=settings.private_upstream_base_url,
|
||||
model=settings.private_upstream_model,
|
||||
api_key=settings.private_upstream_api_key,
|
||||
require_api_key=False,
|
||||
)
|
||||
|
||||
@app.middleware("http")
|
||||
async def log_request(request: Request, call_next):
|
||||
started_at = time.perf_counter()
|
||||
request_id = uuid.uuid4().hex[:8]
|
||||
request.state.request_id = request_id
|
||||
headers = _sanitized_headers(request)
|
||||
body = await request.body()
|
||||
request_log_path = _write_request_file(settings.request_log_dir, request, headers, body, request_id)
|
||||
logger.info(
|
||||
"request_started request_id=%s method=%s path=%s query=%s headers=%s request_log=%s",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
request.url.query,
|
||||
headers,
|
||||
request_log_path,
|
||||
)
|
||||
response = await call_next(request)
|
||||
duration_ms = round((time.perf_counter() - started_at) * 1000, 2)
|
||||
logger.info(
|
||||
"request_finished request_id=%s method=%s path=%s status=%s duration_ms=%s",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
response.status_code,
|
||||
duration_ms,
|
||||
)
|
||||
return response
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
|
||||
@app.get("/v1/models")
|
||||
@app.get("/models")
|
||||
async def models() -> dict[str, Any]:
|
||||
return {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": settings.upstream_model,
|
||||
"object": "model",
|
||||
"owned_by": "nvidia",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
@app.post("/v1/chat/completion")
|
||||
@app.post("/chat/completions")
|
||||
@app.post("/chat/completion")
|
||||
async def chat_completions(request: Request) -> JSONResponse:
|
||||
payload = await request.json()
|
||||
if not isinstance(payload, dict):
|
||||
raise HTTPException(status_code=400, detail="Expected JSON object request body")
|
||||
|
||||
payload_summary = summarize_payload(payload)
|
||||
logger.info(
|
||||
"chat_completion_received model=%s message_count=%s has_image=%s content_part_types=%s",
|
||||
payload_summary["model"],
|
||||
payload_summary["message_count"],
|
||||
payload_summary["has_image"],
|
||||
payload_summary["content_part_types"],
|
||||
)
|
||||
|
||||
outbound_payload, use_private_upstream = _build_outbound_payload(payload, settings, private_model_switch)
|
||||
|
||||
outbound_log_path = _write_json_capture_file(
|
||||
settings.request_log_dir,
|
||||
request.state.request_id,
|
||||
"outbound-redacted-request",
|
||||
outbound_payload,
|
||||
)
|
||||
logger.info("redacted_request_captured request_id=%s path=%s", request.state.request_id, outbound_log_path)
|
||||
|
||||
if outbound_payload.get("stream") is True:
|
||||
logger.info("calling_upstream stream=true private_upstream=%s", use_private_upstream)
|
||||
upstream_client = app.state.private_upstream_client if use_private_upstream else app.state.upstream_client
|
||||
chunks = upstream_client.stream_chat_completions(outbound_payload)
|
||||
if use_private_upstream:
|
||||
chunks = _normalize_reasoning_stream(chunks)
|
||||
stream = _capture_streaming_response(
|
||||
chunks,
|
||||
settings.request_log_dir,
|
||||
request.state.request_id,
|
||||
)
|
||||
return StreamingResponse(
|
||||
stream,
|
||||
media_type="text/event-stream",
|
||||
)
|
||||
|
||||
logger.info("calling_upstream stream=false private_upstream=%s", use_private_upstream)
|
||||
upstream_client = app.state.private_upstream_client if use_private_upstream else app.state.upstream_client
|
||||
response = await upstream_client.chat_completions(outbound_payload)
|
||||
response_body = _normalize_reasoning_response(response.body) if use_private_upstream else response.body
|
||||
response_log_path = _write_json_capture_file(
|
||||
settings.request_log_dir,
|
||||
request.state.request_id,
|
||||
"upstream-response",
|
||||
{"status_code": response.status_code, "body": response_body},
|
||||
)
|
||||
logger.info("upstream_response_captured request_id=%s path=%s", request.state.request_id, response_log_path)
|
||||
return JSONResponse(status_code=response.status_code, content=response_body)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
app = create_app()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
settings = Settings.from_env()
|
||||
uvicorn.run("redact_spike.proxy:app", host=settings.proxy_host, port=settings.proxy_port, reload=False)
|
||||
|
||||
|
||||
def _build_outbound_payload(
|
||||
payload: dict[str, Any],
|
||||
settings: Settings,
|
||||
private_model_switch: bool,
|
||||
) -> tuple[dict[str, Any], bool]:
|
||||
if private_model_switch and last_message_has_image(payload):
|
||||
outbound_payload = dict(payload)
|
||||
if settings.private_switch_force_model:
|
||||
outbound_payload["model"] = settings.private_upstream_model
|
||||
logger.info(
|
||||
"private_model_switch_applied private_base_url=%s private_model=%s force_model=%s",
|
||||
settings.private_upstream_base_url,
|
||||
outbound_payload.get("model"),
|
||||
settings.private_switch_force_model,
|
||||
)
|
||||
return outbound_payload, True
|
||||
|
||||
if private_model_switch:
|
||||
logger.info("private_model_switch_skipped reason=last_message_has_no_image")
|
||||
|
||||
should_redact, redaction_context = _resolve_image_redaction(settings)
|
||||
|
||||
if should_redact:
|
||||
replacement_style = "short" if private_model_switch else "mcp"
|
||||
logger.info(
|
||||
"image_redaction_started store_redacted_images=%s replacement_style=%s custom_context=%s",
|
||||
settings.store_redacted_images,
|
||||
replacement_style,
|
||||
redaction_context is not None,
|
||||
)
|
||||
store = DiskImageStore(settings.redacted_image_dir) if settings.store_redacted_images and not private_model_switch else None
|
||||
redaction_result = redact_images(
|
||||
payload, store=store, replacement_style=replacement_style, redaction_context=redaction_context
|
||||
)
|
||||
if redaction_result.images:
|
||||
logger.info(
|
||||
"image_redaction_applied count=%s placeholders=%s references=%s",
|
||||
len(redaction_result.images),
|
||||
[image.placeholder for image in redaction_result.images],
|
||||
[image.reference for image in redaction_result.images if image.reference],
|
||||
)
|
||||
else:
|
||||
logger.info("image_redaction_skipped reason=no_images_found")
|
||||
return redaction_result.payload, False
|
||||
|
||||
logger.info("image_redaction_skipped reason=allowed_by_policy")
|
||||
return payload, False
|
||||
|
||||
|
||||
def _resolve_image_redaction(settings: Settings) -> tuple[bool, str | None]:
|
||||
"""Decide whether to redact images on this request.
|
||||
|
||||
When a `CONTENT_POLICY_FILE` is configured, the file is the source of truth
|
||||
and is re-read per request (the live demo toggle). Its
|
||||
`content_policy.content_type.image` value (`redact`/`allow`) decides, and
|
||||
`redact_description` overrides the note injected at the redaction site.
|
||||
With no content-policy decision, fall back to the static `REDACTION_ENABLED`.
|
||||
"""
|
||||
rule = load_image_content_rule(settings.content_policy_file)
|
||||
if rule is not None:
|
||||
return rule.should_redact, rule.description
|
||||
return settings.redaction_enabled, None
|
||||
|
||||
|
||||
def _sanitized_headers(request: Request) -> dict[str, str]:
|
||||
allowed_headers = {"accept", "content-type", "user-agent", "host"}
|
||||
secret_headers = {"authorization", "x-api-key", "api-key"}
|
||||
headers: dict[str, str] = {}
|
||||
for key, value in request.headers.items():
|
||||
lower_key = key.lower()
|
||||
if lower_key in secret_headers:
|
||||
headers[lower_key] = "[REDACTED]"
|
||||
elif lower_key in allowed_headers:
|
||||
headers[lower_key] = value
|
||||
return headers
|
||||
|
||||
|
||||
def _write_request_file(log_dir: Path, request: Request, headers: dict[str, str], body: bytes, request_id: str) -> str:
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = dt.datetime.now(dt.UTC).strftime("%Y%m%dT%H%M%S.%fZ")
|
||||
safe_path = re.sub(r"[^a-zA-Z0-9._-]+", "-", request.url.path.strip("/") or "root")
|
||||
path = log_dir / f"{timestamp}-{request_id}-incoming-{request.method.lower()}-{safe_path}.json"
|
||||
|
||||
content_type = request.headers.get("content-type", "")
|
||||
decoded_body = body.decode("utf-8", errors="replace")
|
||||
record: dict[str, Any] = {
|
||||
"method": request.method,
|
||||
"url": str(request.url),
|
||||
"path": request.url.path,
|
||||
"query": request.url.query,
|
||||
"headers": headers,
|
||||
"content_type": content_type,
|
||||
"body_bytes": len(body),
|
||||
"body_text": decoded_body,
|
||||
}
|
||||
|
||||
if "application/json" in content_type.lower() and body:
|
||||
try:
|
||||
record["body_json"] = json.loads(body)
|
||||
except json.JSONDecodeError:
|
||||
record["body_json_parse_error"] = "invalid_json"
|
||||
|
||||
path.write_text(json.dumps(record, indent=2, sort_keys=True), encoding="utf-8")
|
||||
return str(path)
|
||||
|
||||
|
||||
def _write_json_capture_file(log_dir: Path, request_id: str, label: str, payload: Any) -> str:
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = dt.datetime.now(dt.UTC).strftime("%Y%m%dT%H%M%S.%fZ")
|
||||
path = log_dir / f"{timestamp}-{request_id}-{label}.json"
|
||||
path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
|
||||
return str(path)
|
||||
|
||||
|
||||
async def _capture_streaming_response(
|
||||
chunks: AsyncIterator[bytes],
|
||||
log_dir: Path,
|
||||
request_id: str,
|
||||
) -> AsyncIterator[bytes]:
|
||||
captured = bytearray()
|
||||
try:
|
||||
async for chunk in chunks:
|
||||
captured.extend(chunk)
|
||||
yield chunk
|
||||
finally:
|
||||
path = _write_json_capture_file(
|
||||
log_dir,
|
||||
request_id,
|
||||
"upstream-stream-response",
|
||||
{
|
||||
"body_text": captured.decode("utf-8", errors="replace"),
|
||||
"body_bytes": len(captured),
|
||||
},
|
||||
)
|
||||
logger.info("upstream_response_captured request_id=%s path=%s", request_id, path)
|
||||
|
||||
|
||||
async def _normalize_reasoning_stream(chunks: AsyncIterator[bytes]) -> AsyncIterator[bytes]:
|
||||
buffer = ""
|
||||
async for chunk in chunks:
|
||||
buffer += chunk.decode("utf-8", errors="replace")
|
||||
while "\n" in buffer:
|
||||
line, buffer = buffer.split("\n", 1)
|
||||
yield (_normalize_sse_line(line) + "\n").encode("utf-8")
|
||||
if buffer:
|
||||
yield _normalize_sse_line(buffer).encode("utf-8")
|
||||
|
||||
|
||||
def _normalize_sse_line(line: str) -> str:
|
||||
if not line.startswith("data: "):
|
||||
return line
|
||||
|
||||
data = line.removeprefix("data: ")
|
||||
if data.strip() == "[DONE]":
|
||||
return line
|
||||
|
||||
try:
|
||||
payload = json.loads(data)
|
||||
except json.JSONDecodeError:
|
||||
return line
|
||||
|
||||
return "data: " + json.dumps(_normalize_reasoning_response(payload), separators=(",", ":"))
|
||||
|
||||
|
||||
def _normalize_reasoning_response(payload: Any) -> Any:
|
||||
if not isinstance(payload, dict):
|
||||
return payload
|
||||
|
||||
choices = payload.get("choices")
|
||||
if not isinstance(choices, list):
|
||||
return payload
|
||||
|
||||
for choice in choices:
|
||||
if not isinstance(choice, dict):
|
||||
continue
|
||||
for key in ("message", "delta"):
|
||||
message = choice.get(key)
|
||||
if isinstance(message, dict):
|
||||
_copy_reasoning_to_provider_specific_fields(message)
|
||||
return payload
|
||||
|
||||
|
||||
def _copy_reasoning_to_provider_specific_fields(message: dict[str, Any]) -> None:
|
||||
reasoning_content = message.get("reasoning_content")
|
||||
if reasoning_content is None:
|
||||
return
|
||||
|
||||
provider_specific_fields = message.get("provider_specific_fields")
|
||||
if not isinstance(provider_specific_fields, dict):
|
||||
provider_specific_fields = {}
|
||||
message["provider_specific_fields"] = provider_specific_fields
|
||||
provider_specific_fields.setdefault("reasoning_content", reasoning_content)
|
||||
@@ -0,0 +1,233 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
DATA_URL_RE = re.compile(r"^data:(?P<mime>image/[a-zA-Z0-9.+-]+);base64,(?P<data>.+)$", re.DOTALL)
|
||||
IMAGE_KEYWORDS = {"image", "image_url", "input_image"}
|
||||
REDACTION_CONTEXT = (
|
||||
"This image was redacted before this request was sent to you. "
|
||||
"If the user's task requires information from a redacted image, use the pre-installed MCP tool `call_inference` "
|
||||
"to request analysis from the allowed local image-capable model. The MCP request must include a text part plus actual "
|
||||
"image data or an `image_reference` path from this context; a placeholder ID alone is not image data."
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RedactedImage:
|
||||
placeholder: str
|
||||
source: str
|
||||
reference: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RedactionResult:
|
||||
payload: dict[str, Any]
|
||||
images: list[RedactedImage]
|
||||
|
||||
|
||||
class DiskImageStore:
|
||||
def __init__(self, directory: Path) -> None:
|
||||
self.directory = directory
|
||||
|
||||
def store(self, placeholder: str, image_value: Any) -> str:
|
||||
self.directory.mkdir(parents=True, exist_ok=True)
|
||||
serialized = json.dumps(image_value, sort_keys=True).encode("utf-8")
|
||||
digest = hashlib.sha256(serialized).hexdigest()[:16]
|
||||
|
||||
data_url = _find_data_url(image_value)
|
||||
if data_url:
|
||||
match = DATA_URL_RE.match(data_url)
|
||||
if match:
|
||||
extension = _extension_for_mime(match.group("mime"))
|
||||
path = self.directory / f"{placeholder.lower()}-{digest}.{extension}"
|
||||
path.write_bytes(base64.b64decode(match.group("data")))
|
||||
return str(path)
|
||||
|
||||
path = self.directory / f"{placeholder.lower()}-{digest}.json"
|
||||
path.write_text(json.dumps(image_value, indent=2, sort_keys=True), encoding="utf-8")
|
||||
return str(path)
|
||||
|
||||
|
||||
def redact_images(
|
||||
payload: dict[str, Any],
|
||||
store: DiskImageStore | None = None,
|
||||
replacement_style: str = "mcp",
|
||||
redaction_context: str | None = None,
|
||||
) -> RedactionResult:
|
||||
redacted_payload = copy.deepcopy(payload)
|
||||
images: list[RedactedImage] = []
|
||||
|
||||
for message in redacted_payload.get("messages", []):
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
content = message.get("content")
|
||||
if isinstance(content, list):
|
||||
message["content"] = [
|
||||
_redact_content_part(part, images, store, replacement_style, redaction_context) for part in content
|
||||
]
|
||||
elif isinstance(content, str) and _looks_like_image_data_url(content):
|
||||
message["content"] = _replacement_text(
|
||||
_record_image(content, "message.content", images, store), replacement_style, redaction_context
|
||||
)
|
||||
|
||||
return RedactionResult(payload=redacted_payload, images=images)
|
||||
|
||||
|
||||
def summarize_payload(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
messages = payload.get("messages", [])
|
||||
summary: dict[str, Any] = {
|
||||
"model": payload.get("model"),
|
||||
"message_count": len(messages) if isinstance(messages, list) else 0,
|
||||
"roles": [],
|
||||
"content_part_types": [],
|
||||
"has_image": False,
|
||||
}
|
||||
|
||||
if not isinstance(messages, list):
|
||||
return summary
|
||||
|
||||
content_part_types: list[str] = []
|
||||
roles: list[str] = []
|
||||
has_image = False
|
||||
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
roles.append(str(message.get("role", "unknown")))
|
||||
content = message.get("content")
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict):
|
||||
part_type = str(part.get("type", "unknown"))
|
||||
content_part_types.append(part_type)
|
||||
has_image = has_image or _is_image_part(part)
|
||||
elif isinstance(content, str):
|
||||
content_part_types.append("text")
|
||||
has_image = has_image or _looks_like_image_data_url(content)
|
||||
|
||||
summary["roles"] = roles
|
||||
summary["content_part_types"] = content_part_types
|
||||
summary["has_image"] = has_image
|
||||
return summary
|
||||
|
||||
|
||||
def last_message_has_image(payload: dict[str, Any]) -> bool:
|
||||
messages = payload.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return False
|
||||
|
||||
last_message = messages[-1]
|
||||
if not isinstance(last_message, dict):
|
||||
return False
|
||||
|
||||
return message_has_image(last_message)
|
||||
|
||||
|
||||
def message_has_image(message: dict[str, Any]) -> bool:
|
||||
content = message.get("content")
|
||||
if isinstance(content, list):
|
||||
return any(isinstance(part, dict) and _is_image_part(part) for part in content)
|
||||
if isinstance(content, str):
|
||||
return _looks_like_image_data_url(content)
|
||||
return False
|
||||
|
||||
|
||||
def _redact_content_part(
|
||||
part: Any,
|
||||
images: list[RedactedImage],
|
||||
store: DiskImageStore | None,
|
||||
replacement_style: str,
|
||||
redaction_context: str | None = None,
|
||||
) -> Any:
|
||||
if not isinstance(part, dict):
|
||||
return part
|
||||
|
||||
if _is_image_part(part):
|
||||
image = _record_image(part, "message.content[]", images, store)
|
||||
return {"type": "text", "text": _replacement_text(image, replacement_style, redaction_context)}
|
||||
|
||||
return part
|
||||
|
||||
|
||||
def _is_image_part(part: dict[str, Any]) -> bool:
|
||||
part_type = str(part.get("type", "")).lower()
|
||||
if part_type in IMAGE_KEYWORDS:
|
||||
return True
|
||||
if "image_url" in part or "input_image" in part:
|
||||
return True
|
||||
for value in part.values():
|
||||
if isinstance(value, str) and _looks_like_image_data_url(value):
|
||||
return True
|
||||
if isinstance(value, dict) and any(str(key).lower() in IMAGE_KEYWORDS for key in value):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _looks_like_image_data_url(value: str) -> bool:
|
||||
return bool(DATA_URL_RE.match(value.strip()))
|
||||
|
||||
|
||||
def _find_data_url(value: Any) -> str | None:
|
||||
if isinstance(value, str) and _looks_like_image_data_url(value):
|
||||
return value.strip()
|
||||
if isinstance(value, dict):
|
||||
for nested in value.values():
|
||||
data_url = _find_data_url(nested)
|
||||
if data_url:
|
||||
return data_url
|
||||
if isinstance(value, list):
|
||||
for nested in value:
|
||||
data_url = _find_data_url(nested)
|
||||
if data_url:
|
||||
return data_url
|
||||
return None
|
||||
|
||||
|
||||
def _record_image(
|
||||
image_value: Any,
|
||||
source: str,
|
||||
images: list[RedactedImage],
|
||||
store: DiskImageStore | None,
|
||||
) -> RedactedImage:
|
||||
placeholder = f"[REDACTED_IMAGE_{len(images) + 1}]"
|
||||
reference = store.store(placeholder.strip("[]"), image_value) if store else None
|
||||
image = RedactedImage(placeholder=placeholder, source=source, reference=reference)
|
||||
images.append(image)
|
||||
return image
|
||||
|
||||
|
||||
def _replacement_text(image: RedactedImage, replacement_style: str, redaction_context: str | None = None) -> str:
|
||||
text = f"{image.placeholder} image content redacted before upstream inference."
|
||||
if replacement_style == "short":
|
||||
return text
|
||||
|
||||
text += f" {redaction_context or REDACTION_CONTEXT}"
|
||||
if image.reference:
|
||||
text += (
|
||||
f" Local image reference available to MCP for {image.placeholder}: {image.reference}. "
|
||||
"Use MCP payload content like "
|
||||
"[{\"type\":\"text\",\"text\":\"Describe the image.\"}, "
|
||||
f"{{\"type\":\"image_reference\",\"path\":\"{image.reference}\"}}]."
|
||||
)
|
||||
return text
|
||||
|
||||
|
||||
def _extension_for_mime(mime: str) -> str:
|
||||
return {
|
||||
"image/jpeg": "jpg",
|
||||
"image/jpg": "jpg",
|
||||
"image/png": "png",
|
||||
"image/gif": "gif",
|
||||
"image/webp": "webp",
|
||||
}.get(mime.lower(), "bin")
|
||||
@@ -0,0 +1,129 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from redact_spike.config import Settings
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UpstreamResponse:
|
||||
status_code: int
|
||||
body: Any
|
||||
|
||||
|
||||
class NvidiaInferenceClient:
|
||||
def __init__(
|
||||
self,
|
||||
settings: Settings,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
base_url: str | None = None,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
require_api_key: bool = True,
|
||||
) -> None:
|
||||
self.settings = settings
|
||||
self.client = client
|
||||
self.base_url = base_url or settings.upstream_base_url
|
||||
self.model = model or settings.upstream_model
|
||||
self.api_key = api_key if api_key is not None else settings.nvidia_api_key
|
||||
self.require_api_key = require_api_key
|
||||
|
||||
async def chat_completions(self, payload: dict[str, Any]) -> UpstreamResponse:
|
||||
if self.require_api_key and not self.api_key:
|
||||
raise RuntimeError("NVIDIA_API_KEY is required for upstream inference calls")
|
||||
|
||||
body = dict(payload)
|
||||
body.setdefault("model", self.model)
|
||||
|
||||
headers = self._headers()
|
||||
|
||||
close_client = self.client is None
|
||||
client = self.client or httpx.AsyncClient(timeout=60)
|
||||
url = chat_completions_url(self.base_url)
|
||||
logger.info("upstream_request url=%s model=%s stream=%s", url, body.get("model"), body.get("stream", False))
|
||||
try:
|
||||
response = await client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=body,
|
||||
)
|
||||
try:
|
||||
response_body: Any = response.json()
|
||||
except ValueError:
|
||||
response_body = {"error": response.text}
|
||||
logger.info(
|
||||
"upstream_response url=%s status=%s body=%s",
|
||||
url,
|
||||
response.status_code,
|
||||
_sanitize_for_log(response_body),
|
||||
)
|
||||
return UpstreamResponse(status_code=response.status_code, body=response_body)
|
||||
finally:
|
||||
if close_client:
|
||||
await client.aclose()
|
||||
|
||||
async def stream_chat_completions(self, payload: dict[str, Any]) -> AsyncIterator[bytes]:
|
||||
if self.require_api_key and not self.api_key:
|
||||
raise RuntimeError("NVIDIA_API_KEY is required for upstream inference calls")
|
||||
|
||||
body = dict(payload)
|
||||
body.setdefault("model", self.model)
|
||||
|
||||
headers = self._headers(accept="text/event-stream")
|
||||
|
||||
close_client = self.client is None
|
||||
client = self.client or httpx.AsyncClient(timeout=httpx.Timeout(60, read=None))
|
||||
url = chat_completions_url(self.base_url)
|
||||
logger.info("upstream_stream_request url=%s model=%s", url, body.get("model"))
|
||||
try:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
url,
|
||||
headers=headers,
|
||||
json=body,
|
||||
) as response:
|
||||
logger.info("upstream_stream_response url=%s status=%s", url, response.status_code)
|
||||
async for chunk in response.aiter_bytes():
|
||||
yield chunk
|
||||
finally:
|
||||
if close_client:
|
||||
await client.aclose()
|
||||
|
||||
def _headers(self, accept: str | None = None) -> dict[str, str]:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if accept:
|
||||
headers["Accept"] = accept
|
||||
if self.api_key:
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
return headers
|
||||
|
||||
|
||||
def chat_completions_url(base_url: str) -> str:
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/v1"):
|
||||
return f"{base}/chat/completions"
|
||||
return f"{base}/v1/chat/completions"
|
||||
|
||||
|
||||
def _sanitize_for_log(value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {key: _sanitize_for_log(nested) for key, nested in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_sanitize_for_log(item) for item in value]
|
||||
if isinstance(value, str):
|
||||
if value.strip().startswith("data:image/"):
|
||||
return "[REDACTED_IMAGE_DATA_URL]"
|
||||
if len(value) > 500:
|
||||
return f"{value[:500]}...[TRUNCATED {len(value) - 500} chars]"
|
||||
return value
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from redact_spike.config import Settings
|
||||
|
||||
|
||||
def test_settings_loads_dotenv_file(monkeypatch, tmp_path: Path) -> None:
|
||||
env_file = tmp_path / ".env"
|
||||
env_file.write_text(
|
||||
"NVIDIA_API_KEY=from-dotenv\nUPSTREAM_MODEL=model-from-dotenv\nPROXY_PORT=9999\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setenv("ENV_FILE", str(env_file))
|
||||
monkeypatch.delenv("NVIDIA_API_KEY", raising=False)
|
||||
monkeypatch.delenv("UPSTREAM_MODEL", raising=False)
|
||||
monkeypatch.delenv("PROXY_PORT", raising=False)
|
||||
|
||||
settings = Settings.from_env()
|
||||
|
||||
assert settings.nvidia_api_key == "from-dotenv"
|
||||
assert settings.upstream_model == "model-from-dotenv"
|
||||
assert settings.proxy_port == 9999
|
||||
|
||||
|
||||
def test_dotenv_values_override_environment_values(monkeypatch, tmp_path: Path) -> None:
|
||||
env_file = tmp_path / ".env"
|
||||
env_file.write_text("NVIDIA_API_KEY=from-dotenv\n", encoding="utf-8")
|
||||
monkeypatch.setenv("ENV_FILE", str(env_file))
|
||||
monkeypatch.setenv("NVIDIA_API_KEY", "from-environment")
|
||||
|
||||
settings = Settings.from_env()
|
||||
|
||||
assert settings.nvidia_api_key == "from-dotenv"
|
||||
@@ -0,0 +1,52 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from redact_spike.content_policy import load_image_content_rule
|
||||
|
||||
|
||||
def _write(tmp_path: Path, body: str) -> Path:
|
||||
path = tmp_path / "content-policy.yaml"
|
||||
path.write_text(body, encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
def test_returns_none_when_path_missing(tmp_path):
|
||||
assert load_image_content_rule(None) is None
|
||||
assert load_image_content_rule(tmp_path / "nope.yaml") is None
|
||||
|
||||
|
||||
def test_returns_none_without_image_key(tmp_path):
|
||||
path = _write(tmp_path, "version: 1\ncontent_policy:\n content_type: {}\n")
|
||||
assert load_image_content_rule(path) is None
|
||||
|
||||
|
||||
def test_redact_rule_with_description(tmp_path):
|
||||
path = _write(
|
||||
tmp_path,
|
||||
"content_policy:\n"
|
||||
" content_type:\n"
|
||||
" image: redact\n"
|
||||
" redact_description: Use the local_model MCP server.\n",
|
||||
)
|
||||
rule = load_image_content_rule(path)
|
||||
assert rule is not None
|
||||
assert rule.should_redact is True
|
||||
assert rule.description == "Use the local_model MCP server."
|
||||
|
||||
|
||||
def test_allow_rule(tmp_path):
|
||||
path = _write(tmp_path, "content_policy:\n content_type:\n image: allow\n")
|
||||
rule = load_image_content_rule(path)
|
||||
assert rule is not None
|
||||
assert rule.should_redact is False
|
||||
|
||||
|
||||
def test_live_toggle_is_reread(tmp_path):
|
||||
path = _write(tmp_path, "content_policy:\n content_type:\n image: allow\n")
|
||||
assert load_image_content_rule(path).should_redact is False
|
||||
path.write_text("content_policy:\n content_type:\n image: redact\n", encoding="utf-8")
|
||||
assert load_image_content_rule(path).should_redact is True
|
||||
@@ -0,0 +1,110 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from redact_spike import mcp_server
|
||||
from redact_spike.config import Settings
|
||||
from redact_spike.upstream import UpstreamResponse
|
||||
|
||||
|
||||
class FakeNvidiaInferenceClient:
|
||||
def __init__(self, settings: Settings, **kwargs: Any) -> None:
|
||||
self.settings = settings
|
||||
self.kwargs = kwargs
|
||||
|
||||
async def chat_completions(self, payload: dict[str, Any]) -> UpstreamResponse:
|
||||
return UpstreamResponse(status_code=200, body={"echo": payload, "model": self.kwargs["model"], "base_url": self.kwargs["base_url"]})
|
||||
|
||||
|
||||
async def test_call_inference_impl_returns_openai_shaped_payload(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(mcp_server, "NvidiaInferenceClient", FakeNvidiaInferenceClient)
|
||||
|
||||
payload = {"messages": [{"role": "user", "content": "describe image"}]}
|
||||
response = await mcp_server.call_inference_impl(
|
||||
payload,
|
||||
settings=Settings(private_upstream_base_url="http://localhost:1234", private_upstream_model="qwen/qwen3.6-27b"),
|
||||
)
|
||||
|
||||
assert response == {
|
||||
"status_code": 200,
|
||||
"body": {
|
||||
"echo": {"messages": [{"role": "user", "content": "describe image"}], "model": "qwen/qwen3.6-27b"},
|
||||
"model": "qwen/qwen3.6-27b",
|
||||
"base_url": "http://localhost:1234",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_normalize_mcp_payload_forces_configured_model() -> None:
|
||||
payload = {
|
||||
"model": "local-image-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}
|
||||
|
||||
normalized = mcp_server.normalize_mcp_payload(payload, Settings(private_upstream_model="qwen/qwen3.6-27b"))
|
||||
|
||||
assert normalized["model"] == "qwen/qwen3.6-27b"
|
||||
assert payload["model"] == "local-image-model"
|
||||
|
||||
|
||||
def test_normalize_mcp_payload_converts_image_placeholder_to_text() -> None:
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe the image."},
|
||||
{"type": "image_placeholder", "id": "REDACTED_IMAGE_1"},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
normalized = mcp_server.normalize_mcp_payload(payload, Settings(private_upstream_model="qwen/qwen3.6-27b"))
|
||||
|
||||
assert normalized["model"] == "qwen/qwen3.6-27b"
|
||||
assert normalized["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "Describe the image."},
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"Image placeholder [REDACTED_IMAGE_1] was referenced, but no image bytes or local image reference were "
|
||||
"provided to this MCP call. Retry with a text part plus an image_reference part using the local path "
|
||||
"from the redaction context."
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def test_normalize_mcp_payload_converts_image_reference_to_data_url(tmp_path) -> None:
|
||||
image = tmp_path / "redacted.png"
|
||||
image.write_bytes(b"fake image bytes")
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe the image."},
|
||||
{"type": "image_reference", "path": str(image)},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
normalized = mcp_server.normalize_mcp_payload(payload, Settings(private_upstream_model="qwen/qwen3.6-27b"))
|
||||
|
||||
assert normalized["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "Describe the image."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZSBpbWFnZSBieXRlcw=="}},
|
||||
]
|
||||
|
||||
|
||||
def test_create_mcp_server() -> None:
|
||||
server = mcp_server.create_mcp_server(settings=Settings(nvidia_api_key="test"))
|
||||
|
||||
assert server.name == "redacted-inference"
|
||||
@@ -0,0 +1,403 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from redact_spike.config import Settings
|
||||
from redact_spike.proxy import create_app
|
||||
from redact_spike.upstream import UpstreamResponse
|
||||
|
||||
|
||||
class RecordingUpstream:
|
||||
def __init__(self) -> None:
|
||||
self.payloads: list[dict[str, Any]] = []
|
||||
|
||||
async def chat_completions(self, payload: dict[str, Any]) -> UpstreamResponse:
|
||||
self.payloads.append(payload)
|
||||
return UpstreamResponse(status_code=200, body={"id": "chatcmpl-test", "choices": []})
|
||||
|
||||
async def stream_chat_completions(self, payload: dict[str, Any]):
|
||||
self.payloads.append(payload)
|
||||
yield b"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n"
|
||||
yield b"data: [DONE]\n\n"
|
||||
|
||||
|
||||
class ReasoningUpstream:
|
||||
def __init__(self) -> None:
|
||||
self.payloads: list[dict[str, Any]] = []
|
||||
|
||||
async def chat_completions(self, payload: dict[str, Any]) -> UpstreamResponse:
|
||||
self.payloads.append(payload)
|
||||
return UpstreamResponse(
|
||||
status_code=200,
|
||||
body={
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "answer",
|
||||
"reasoning_content": "private reasoning",
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
async def stream_chat_completions(self, payload: dict[str, Any]):
|
||||
self.payloads.append(payload)
|
||||
yield b'data: {"choices":[{"delta":{"role":"assistant","reasoning_content":"private"}}]}\n\n'
|
||||
yield b'data: {"choices":[{"delta":{"reasoning_content":" reasoning"}}]}\n\n'
|
||||
yield b"data: [DONE]\n\n"
|
||||
|
||||
|
||||
async def test_proxy_forwards_chat_completion_payload_when_redaction_disabled() -> None:
|
||||
upstream = RecordingUpstream()
|
||||
settings = Settings(nvidia_api_key="test", redaction_enabled=False)
|
||||
app = create_app(settings=settings, upstream_client=upstream)
|
||||
payload = {
|
||||
"model": "openai/openai/gpt-5.5",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 1024,
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"id": "chatcmpl-test", "choices": []}
|
||||
assert upstream.payloads == [payload]
|
||||
|
||||
|
||||
async def test_proxy_redacts_image_payload_before_forwarding() -> None:
|
||||
upstream = RecordingUpstream()
|
||||
settings = Settings(nvidia_api_key="test", redaction_enabled=True)
|
||||
app = create_app(settings=settings, upstream_client=upstream)
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is in this image?"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
forwarded = upstream.payloads[0]
|
||||
assert "data:image/png" not in str(forwarded)
|
||||
assert "[REDACTED_IMAGE_1]" in str(forwarded)
|
||||
assert "pre-installed MCP tool `call_inference`" in forwarded["messages"][0]["content"][1]["text"]
|
||||
|
||||
|
||||
async def test_proxy_streams_chat_completion_payloads() -> None:
|
||||
upstream = RecordingUpstream()
|
||||
settings = Settings(nvidia_api_key="test", redaction_enabled=False)
|
||||
app = create_app(settings=settings, upstream_client=upstream)
|
||||
payload = {"stream": True, "messages": [{"role": "user", "content": "Hello"}]}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"].startswith("text/event-stream")
|
||||
assert response.content == b"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n"
|
||||
assert upstream.payloads == [payload]
|
||||
|
||||
|
||||
async def test_proxy_supports_chat_completion_route_aliases() -> None:
|
||||
upstream = RecordingUpstream()
|
||||
settings = Settings(nvidia_api_key="test", redaction_enabled=False)
|
||||
app = create_app(settings=settings, upstream_client=upstream)
|
||||
payload = {"messages": [{"role": "user", "content": "Hello"}]}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
for path in ["/chat/completions", "/chat/completion", "/v1/chat/completion"]:
|
||||
response = await client.post(path, json=payload)
|
||||
assert response.status_code == 200
|
||||
|
||||
assert upstream.payloads == [payload, payload, payload]
|
||||
|
||||
|
||||
async def test_proxy_lists_configured_model() -> None:
|
||||
upstream = RecordingUpstream()
|
||||
settings = Settings(nvidia_api_key="test", upstream_model="openai/openai/gpt-5.5")
|
||||
app = create_app(settings=settings, upstream_client=upstream)
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get("/v1/models")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "openai/openai/gpt-5.5",
|
||||
"object": "model",
|
||||
"owned_by": "nvidia",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
async def test_proxy_rejects_non_object_json() -> None:
|
||||
upstream = RecordingUpstream()
|
||||
app = create_app(settings=Settings(nvidia_api_key="test"), upstream_client=upstream)
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=["not", "an", "object"])
|
||||
|
||||
assert response.status_code == 400
|
||||
assert upstream.payloads == []
|
||||
|
||||
|
||||
async def test_proxy_logs_unknown_requests(caplog) -> None:
|
||||
upstream = RecordingUpstream()
|
||||
app = create_app(settings=Settings(nvidia_api_key="test"), upstream_client=upstream)
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
with caplog.at_level(logging.INFO, logger="redact_spike.proxy"):
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.get("/definitely-unknown", headers={"Authorization": "Bearer local"})
|
||||
|
||||
assert response.status_code == 404
|
||||
assert "method=GET path=/definitely-unknown" in caplog.text
|
||||
assert "path=/definitely-unknown status=404" in caplog.text
|
||||
assert "Bearer local" not in caplog.text
|
||||
assert "[REDACTED]" in caplog.text
|
||||
|
||||
|
||||
async def test_proxy_writes_full_request_body_to_file(caplog, tmp_path) -> None:
|
||||
upstream = RecordingUpstream()
|
||||
app = create_app(settings=Settings(nvidia_api_key="test", request_log_dir=tmp_path), upstream_client=upstream)
|
||||
payload = {
|
||||
"model": "bad-model",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
with caplog.at_level(logging.INFO, logger="redact_spike.proxy"):
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "bad-model" in caplog.text
|
||||
assert "data:image/png" not in caplog.text
|
||||
request_log_files = list(tmp_path.glob("*incoming*.json"))
|
||||
outbound_log_files = list(tmp_path.glob("*outbound-redacted-request.json"))
|
||||
response_log_files = list(tmp_path.glob("*upstream-response.json"))
|
||||
assert len(request_log_files) == 1
|
||||
assert len(outbound_log_files) == 1
|
||||
assert len(response_log_files) == 1
|
||||
request_log = json.loads(request_log_files[0].read_text(encoding="utf-8"))
|
||||
outbound_log = json.loads(outbound_log_files[0].read_text(encoding="utf-8"))
|
||||
response_log = json.loads(response_log_files[0].read_text(encoding="utf-8"))
|
||||
assert request_log["path"] == "/v1/chat/completions"
|
||||
assert request_log["body_json"] == payload
|
||||
assert "data:image/png" in request_log["body_text"]
|
||||
assert "data:image/png" not in str(outbound_log)
|
||||
assert "[REDACTED_IMAGE_1]" in str(outbound_log)
|
||||
assert response_log == {"status_code": 200, "body": {"id": "chatcmpl-test", "choices": []}}
|
||||
|
||||
|
||||
async def test_proxy_logs_redaction_and_upstream_lifecycle(caplog, tmp_path) -> None:
|
||||
upstream = RecordingUpstream()
|
||||
app = create_app(settings=Settings(nvidia_api_key="test", request_log_dir=tmp_path), upstream_client=upstream)
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is in this image?"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
with caplog.at_level(logging.INFO, logger="redact_spike.proxy"):
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "chat_completion_received" in caplog.text
|
||||
assert "image_redaction_started" in caplog.text
|
||||
assert "image_redaction_applied count=1" in caplog.text
|
||||
assert "redacted_request_captured" in caplog.text
|
||||
assert "calling_upstream stream=false" in caplog.text
|
||||
assert "upstream_response_captured" in caplog.text
|
||||
|
||||
|
||||
async def test_proxy_captures_streaming_upstream_response(tmp_path) -> None:
|
||||
upstream = RecordingUpstream()
|
||||
app = create_app(
|
||||
settings=Settings(nvidia_api_key="test", redaction_enabled=False, request_log_dir=tmp_path),
|
||||
upstream_client=upstream,
|
||||
)
|
||||
payload = {"stream": True, "messages": [{"role": "user", "content": "Hello"}]}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
response_log_files = list(tmp_path.glob("*upstream-stream-response.json"))
|
||||
assert len(response_log_files) == 1
|
||||
response_log = json.loads(response_log_files[0].read_text(encoding="utf-8"))
|
||||
assert response_log["body_text"] == "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n"
|
||||
|
||||
|
||||
async def test_private_switch_routes_last_message_image_to_private_model_without_redaction(tmp_path: Path) -> None:
|
||||
public_upstream = RecordingUpstream()
|
||||
private_upstream = RecordingUpstream()
|
||||
settings = Settings(
|
||||
nvidia_api_key="test",
|
||||
private_upstream_model="private-image-model",
|
||||
request_log_dir=tmp_path,
|
||||
)
|
||||
app = create_app(
|
||||
settings=settings,
|
||||
upstream_client=public_upstream,
|
||||
private_upstream_client=private_upstream,
|
||||
private_model_switch=True,
|
||||
)
|
||||
payload = {
|
||||
"model": "public-model",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert public_upstream.payloads == []
|
||||
forwarded = private_upstream.payloads[0]
|
||||
assert forwarded["model"] == "private-image-model"
|
||||
assert "data:image/png" in str(forwarded)
|
||||
|
||||
|
||||
async def test_private_switch_redacts_older_images_when_last_message_has_no_image() -> None:
|
||||
upstream = RecordingUpstream()
|
||||
settings = Settings(nvidia_api_key="test", private_upstream_model="private-image-model")
|
||||
app = create_app(settings=settings, upstream_client=upstream, private_upstream_client=RecordingUpstream(), private_model_switch=True)
|
||||
payload = {
|
||||
"model": "public-model",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Old image."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "What should I do?"},
|
||||
{"role": "user", "content": "Answer without using that image."},
|
||||
],
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
forwarded = upstream.payloads[0]
|
||||
assert forwarded["model"] == "public-model"
|
||||
assert "data:image/png" not in str(forwarded)
|
||||
assert "[REDACTED_IMAGE_1] image content redacted before upstream inference." in str(forwarded)
|
||||
assert "pre-installed MCP" not in str(forwarded)
|
||||
|
||||
|
||||
async def test_private_switch_copies_reasoning_content_to_provider_specific_fields(tmp_path: Path) -> None:
|
||||
private_upstream = ReasoningUpstream()
|
||||
app = create_app(
|
||||
settings=Settings(nvidia_api_key="test", private_upstream_model="private-image-model", request_log_dir=tmp_path),
|
||||
upstream_client=RecordingUpstream(),
|
||||
private_upstream_client=private_upstream,
|
||||
private_model_switch=True,
|
||||
)
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
message = response.json()["choices"][0]["message"]
|
||||
assert message["reasoning_content"] == "private reasoning"
|
||||
assert message["provider_specific_fields"]["reasoning_content"] == "private reasoning"
|
||||
|
||||
|
||||
async def test_private_switch_stream_copies_reasoning_content_to_provider_specific_fields(tmp_path: Path) -> None:
|
||||
private_upstream = ReasoningUpstream()
|
||||
app = create_app(
|
||||
settings=Settings(nvidia_api_key="test", private_upstream_model="private-image-model", request_log_dir=tmp_path),
|
||||
upstream_client=RecordingUpstream(),
|
||||
private_upstream_client=private_upstream,
|
||||
private_model_switch=True,
|
||||
)
|
||||
payload = {
|
||||
"stream": True,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post("/v1/chat/completions", json=payload)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert b'"reasoning_content":"private"' in response.content
|
||||
assert b'"provider_specific_fields":{"reasoning_content":"private"}' in response.content
|
||||
assert b'"provider_specific_fields":{"reasoning_content":" reasoning"}' in response.content
|
||||
@@ -0,0 +1,141 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
from redact_spike.redaction import DiskImageStore, last_message_has_image, redact_images, summarize_payload
|
||||
|
||||
|
||||
def image_data_url() -> str:
|
||||
return "data:image/png;base64," + base64.b64encode(b"fake png bytes").decode("ascii")
|
||||
|
||||
|
||||
def test_redacts_image_parts_and_adds_mcp_context_at_image_location() -> None:
|
||||
payload = {
|
||||
"model": "openai/openai/gpt-5.5",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Am I dressed for rain?"},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url()}},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
result = redact_images(payload)
|
||||
|
||||
assert len(result.images) == 1
|
||||
assert result.images[0].placeholder == "[REDACTED_IMAGE_1]"
|
||||
assert image_data_url() not in str(result.payload)
|
||||
assert result.payload["messages"][0] == {"role": "system", "content": "You are helpful."}
|
||||
replacement = result.payload["messages"][1]["content"][1]
|
||||
assert replacement["type"] == "text"
|
||||
assert "This image was redacted" in replacement["text"]
|
||||
assert "pre-installed MCP tool `call_inference`" in replacement["text"]
|
||||
assert "[REDACTED_IMAGE_1]" in replacement["text"]
|
||||
|
||||
|
||||
def test_can_store_redacted_data_urls_on_disk(tmp_path) -> None:
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this."},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url()}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
result = redact_images(payload, store=DiskImageStore(tmp_path))
|
||||
|
||||
assert result.images[0].reference is not None
|
||||
assert result.images[0].reference.endswith(".png")
|
||||
assert (tmp_path / result.images[0].reference.split("/")[-1]).read_bytes() == b"fake png bytes"
|
||||
replacement = result.payload["messages"][0]["content"][1]
|
||||
assert "Local image reference available to MCP" in replacement["text"]
|
||||
assert str(tmp_path) in replacement["text"]
|
||||
assert image_data_url() not in str(result.payload)
|
||||
|
||||
|
||||
def test_short_redaction_omits_mcp_context() -> None:
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this."},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url()}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
result = redact_images(payload, replacement_style="short")
|
||||
|
||||
assert result.payload["messages"][0]["content"][1] == {
|
||||
"type": "text",
|
||||
"text": "[REDACTED_IMAGE_1] image content redacted before upstream inference.",
|
||||
}
|
||||
|
||||
|
||||
def test_last_message_has_image_only_checks_latest_message() -> None:
|
||||
payload = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Old image."},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url()}},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "What should I do with it?"},
|
||||
{"role": "user", "content": "Never mind, answer without the image."},
|
||||
]
|
||||
}
|
||||
|
||||
assert last_message_has_image(payload) is False
|
||||
|
||||
payload["messages"][-1] = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Now describe this."},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url()}},
|
||||
],
|
||||
}
|
||||
|
||||
assert last_message_has_image(payload) is True
|
||||
|
||||
|
||||
def test_summarize_payload_reports_structure_without_prompt_text() -> None:
|
||||
payload = {
|
||||
"model": "model-a",
|
||||
"messages": [
|
||||
{"role": "system", "content": "secret system prompt"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "secret prompt text"},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url()}},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
summary = summarize_payload(payload)
|
||||
|
||||
assert summary == {
|
||||
"model": "model-a",
|
||||
"message_count": 2,
|
||||
"roles": ["system", "user"],
|
||||
"content_part_types": ["text", "text", "image_url"],
|
||||
"has_image": True,
|
||||
}
|
||||
assert "secret" not in str(summary)
|
||||
assert image_data_url() not in str(summary)
|
||||
@@ -0,0 +1,108 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
|
||||
from redact_spike.config import Settings
|
||||
from redact_spike.upstream import NvidiaInferenceClient, chat_completions_url
|
||||
|
||||
|
||||
async def test_upstream_client_adds_auth_header_and_default_model() -> None:
|
||||
seen_request: httpx.Request | None = None
|
||||
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal seen_request
|
||||
seen_request = request
|
||||
return httpx.Response(200, json={"choices": [{"message": {"content": "ok"}}]})
|
||||
|
||||
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
||||
settings = Settings(nvidia_api_key="secret", upstream_base_url="https://upstream.test", upstream_model="default-model")
|
||||
client = NvidiaInferenceClient(settings=settings, client=http_client)
|
||||
|
||||
response = await client.chat_completions({"messages": []})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.body == {"choices": [{"message": {"content": "ok"}}]}
|
||||
assert seen_request is not None
|
||||
assert str(seen_request.url) == "https://upstream.test/v1/chat/completions"
|
||||
assert seen_request.headers["authorization"] == "Bearer secret"
|
||||
assert seen_request.headers["content-type"] == "application/json"
|
||||
assert b'"model":"default-model"' in seen_request.content
|
||||
|
||||
await http_client.aclose()
|
||||
|
||||
|
||||
def test_chat_completions_url_accepts_base_url_with_or_without_v1() -> None:
|
||||
assert chat_completions_url("https://inference-api.nvidia.com") == (
|
||||
"https://inference-api.nvidia.com/v1/chat/completions"
|
||||
)
|
||||
assert chat_completions_url("https://inference-api.nvidia.com/v1") == (
|
||||
"https://inference-api.nvidia.com/v1/chat/completions"
|
||||
)
|
||||
assert chat_completions_url("https://inference-api.nvidia.com/v1/") == (
|
||||
"https://inference-api.nvidia.com/v1/chat/completions"
|
||||
)
|
||||
|
||||
|
||||
async def test_upstream_client_requires_api_key() -> None:
|
||||
client = NvidiaInferenceClient(settings=Settings(nvidia_api_key=None))
|
||||
|
||||
try:
|
||||
await client.chat_completions({"messages": []})
|
||||
except RuntimeError as error:
|
||||
assert "NVIDIA_API_KEY" in str(error)
|
||||
else:
|
||||
raise AssertionError("expected missing API key error")
|
||||
|
||||
|
||||
async def test_upstream_client_can_call_local_openai_compatible_server_without_auth() -> None:
|
||||
seen_request: httpx.Request | None = None
|
||||
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal seen_request
|
||||
seen_request = request
|
||||
return httpx.Response(200, json={"choices": []})
|
||||
|
||||
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
||||
client = NvidiaInferenceClient(
|
||||
settings=Settings(),
|
||||
client=http_client,
|
||||
base_url="http://localhost:1234",
|
||||
model="qwen/qwen3.6-27b",
|
||||
api_key=None,
|
||||
require_api_key=False,
|
||||
)
|
||||
|
||||
response = await client.chat_completions({"messages": []})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert seen_request is not None
|
||||
assert str(seen_request.url) == "http://localhost:1234/v1/chat/completions"
|
||||
assert "authorization" not in seen_request.headers
|
||||
assert b'"model":"qwen/qwen3.6-27b"' in seen_request.content
|
||||
|
||||
await http_client.aclose()
|
||||
|
||||
|
||||
async def test_upstream_client_logs_url_status_and_body(caplog) -> None:
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(404, json={"error": {"message": "not found"}})
|
||||
|
||||
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
||||
settings = Settings(nvidia_api_key="secret", upstream_base_url="https://upstream.test/v1")
|
||||
client = NvidiaInferenceClient(settings=settings, client=http_client)
|
||||
|
||||
with caplog.at_level(logging.INFO, logger="redact_spike.upstream"):
|
||||
response = await client.chat_completions({"model": "openai/openai/gpt-5.5", "messages": []})
|
||||
|
||||
assert response.status_code == 404
|
||||
assert "upstream_request url=https://upstream.test/v1/chat/completions" in caplog.text
|
||||
assert "model=openai/openai/gpt-5.5" in caplog.text
|
||||
assert "upstream_response url=https://upstream.test/v1/chat/completions status=404" in caplog.text
|
||||
assert "not found" in caplog.text
|
||||
|
||||
await http_client.aclose()
|
||||
@@ -0,0 +1,33 @@
|
||||
#!/bin/sh
|
||||
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Provision opencode inside an OpenShell sandbox, wired to inference.local plus
|
||||
# the host-side local-model MCP. opencode is preinstalled on the base image.
|
||||
# The opencode.json config is fed on stdin (ssh '...' < opencode.json).
|
||||
set -e
|
||||
export HOME=/sandbox
|
||||
|
||||
CFG_DIR="$HOME/.config/opencode"
|
||||
mkdir -p "$CFG_DIR"
|
||||
|
||||
# opencode treats its working directory as the project root and needs one to
|
||||
# start a session in.
|
||||
WORKSPACE="$HOME/workspace"
|
||||
echo "== creating opencode workspace ($WORKSPACE) =="
|
||||
mkdir -p "$WORKSPACE"
|
||||
|
||||
echo "== writing opencode config ($CFG_DIR/opencode.json) =="
|
||||
cat > "$CFG_DIR/opencode.json"
|
||||
|
||||
echo "== opencode version =="
|
||||
opencode --version
|
||||
|
||||
echo "== configured providers/models =="
|
||||
opencode models 2>&1 | grep -i inference-local || true
|
||||
|
||||
echo "== registered MCP servers =="
|
||||
opencode mcp list 2>&1 || true
|
||||
|
||||
echo "== provision complete =="
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"$schema": "https://opencode.ai/config.json",
|
||||
"autoupdate": false,
|
||||
"share": "disabled",
|
||||
"model": "inference-local/openai/openai/gpt-5.5",
|
||||
"small_model": "inference-local/openai/openai/gpt-5.5",
|
||||
"provider": {
|
||||
"inference-local": {
|
||||
"npm": "@ai-sdk/openai-compatible",
|
||||
"name": "OpenShell inference.local",
|
||||
"options": {
|
||||
"baseURL": "https://inference.local/v1",
|
||||
"apiKey": "dummy"
|
||||
},
|
||||
"models": {
|
||||
"openai/openai/gpt-5.5": {
|
||||
"name": "gpt-5.5 (inference.local)",
|
||||
"modalities": {
|
||||
"input": ["text", "image", "pdf"],
|
||||
"output": ["text"]
|
||||
},
|
||||
"limit": {
|
||||
"context": 131072,
|
||||
"output": 32768
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"mcp": {
|
||||
"local_model": {
|
||||
"type": "remote",
|
||||
"url": "http://host.openshell.internal:8001/mcp",
|
||||
"enabled": true,
|
||||
"oauth": false
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Create an OpenShell sandbox running opencode, wired to inference.local and the
|
||||
# host-side local-model MCP, then start the opencode web UI and expose it.
|
||||
#
|
||||
# Prereqs (on the host):
|
||||
# - An OpenShell gateway is registered and selected (openshell gateway list).
|
||||
# - The redaction proxy (:8000) and MCP (:8001) are running: see ../proxy.
|
||||
# - inference is routed to the proxy (see README "One-time gateway wiring").
|
||||
set -euo pipefail
|
||||
|
||||
SANDBOX_NAME="${SANDBOX_NAME:-opencode}"
|
||||
WEB_PORT="${WEB_PORT:-9119}"
|
||||
SERVICE_NAME="${SERVICE_NAME:-web}"
|
||||
|
||||
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
POLICY="$HERE/../policies/opencode-policy.yaml"
|
||||
CONFIG="$HERE/opencode.json"
|
||||
PROVISION="$HERE/opencode-provision.sh"
|
||||
SSH_CONFIG="/tmp/${SANDBOX_NAME}-ssh-config"
|
||||
|
||||
echo "== creating sandbox '$SANDBOX_NAME' =="
|
||||
openshell sandbox create --name "$SANDBOX_NAME" --policy "$POLICY" --no-tty -- echo sandbox-ready
|
||||
|
||||
echo "== fetching ssh config =="
|
||||
openshell sandbox ssh-config "$SANDBOX_NAME" > "$SSH_CONFIG"
|
||||
SSH=(ssh -F "$SSH_CONFIG" "openshell-${SANDBOX_NAME}")
|
||||
|
||||
echo "== provisioning opencode config =="
|
||||
"${SSH[@]}" 'cat > /tmp/opencode-provision.sh' < "$PROVISION"
|
||||
"${SSH[@]}" 'cat > /tmp/opencode.json' < "$CONFIG"
|
||||
"${SSH[@]}" 'sh /tmp/opencode-provision.sh < /tmp/opencode.json'
|
||||
|
||||
echo "== starting opencode web on 127.0.0.1:${WEB_PORT} (inside sandbox) =="
|
||||
# setsid+nohup so the server survives the ssh session. Logs to /sandbox/web.log.
|
||||
# NB: do not pkill by "opencode web" from a one-liner whose text contains that
|
||||
# string - pkill -f would match its own shell. Kill by port/PID instead.
|
||||
"${SSH[@]}" "export HOME=/sandbox; cd /sandbox/workspace; rm -f /sandbox/web.log; \
|
||||
setsid nohup opencode web --hostname 127.0.0.1 --port ${WEB_PORT} > /sandbox/web.log 2>&1 < /dev/null & echo launched"
|
||||
sleep 6
|
||||
"${SSH[@]}" "curl -s -m 6 -o /dev/null -w 'sandbox web / = HTTP %{http_code}\n' http://127.0.0.1:${WEB_PORT}/" || true
|
||||
|
||||
echo "== exposing web service via the gateway =="
|
||||
openshell service expose "$SANDBOX_NAME" "$WEB_PORT" "$SERVICE_NAME"
|
||||
|
||||
echo
|
||||
echo "Done. Open the URL printed above (http://${SANDBOX_NAME}--${SERVICE_NAME}.openshell.localhost:<gateway-port>/)."
|
||||
echo "Verify the wiring:"
|
||||
echo " ${SSH[*]} 'export HOME=/sandbox; cd /sandbox/workspace; opencode run --model inference-local/openai/openai/gpt-5.5 \"say hi\"'"
|
||||
Reference in New Issue
Block a user