diff --git a/examples/gateway/README.md b/examples/gateway/README.md new file mode 100644 index 00000000..e2e7ca92 --- /dev/null +++ b/examples/gateway/README.md @@ -0,0 +1,267 @@ +# Gateway Debug Launcher + +`debug_launcher.py` starts the existing gateway without training, creates one +external session, prints connection URLs, and writes finalized trajectories when +the session is released. It is meant for local gateway debugging and interaction +collection. It also writes a debug snapshot with the normalized message history +and pre-finalize session state so token-level trajectories can be interpreted +after the run. + +It does not add a gateway HTTP control plane, does not create a training +framework mode, and does not add a sessionless `/v1/messages` facade. + +## Local Backend + +For real local inference, start an OpenAI-compatible completions backend first. +For example: + +```bash +CUDA_VISIBLE_DEVICES=0 vllm serve /data1/models/Qwen/Qwen3.5-4B \ + --host 127.0.0.1 \ + --port 38315 \ + --gpu-memory-utilization 0.75 \ + --max-model-len 4096 \ + --trust-remote-code +``` + +The launcher calls `{backend_base_url}/completions` with token-id prompts and +`return_token_ids=true`. The response must include `choices[0].token_ids`. +Missing token ids are a hard failure because trajectory response tokens must +come from backend `TokenOutput`, not from text re-tokenization. + +Before starting the launcher, verify that the backend is ready: + +```bash +python - <<'PY' +import httpx + +r = httpx.get("http://127.0.0.1:38315/v1/models", timeout=5) +print(r.status_code, r.text[:300]) +PY +``` + +The launcher should point at the same host and port: + +```text +--backend-base-url http://127.0.0.1:38315/v1 +``` + +## Manual Claude Code Session + +Start the launcher without `--run-claude`: + +```bash +python examples/gateway/debug_launcher.py \ + --backend openai-completions \ + --backend-base-url http://127.0.0.1:38315/v1 \ + --backend-model /data1/models/Qwen/Qwen3.5-4B \ + --tokenizer /data1/models/Qwen/Qwen3.5-4B \ + --session-id manual-cc-001 \ + --output-dir /tmp/uni-agent-gateway-debug \ + --response-length 64 +``` + +The launcher prints per-session URLs: + +```text +Claude Code: + export ANTHROPIC_BASE_URL=http://host:port/sessions/manual-cc-001 + export ANTHROPIC_API_KEY=dummy-local-key +OpenAI-compatible: + base_url=http://host:port/sessions/manual-cc-001/v1 +``` + +In another terminal, use the printed values: + +```bash +export ANTHROPIC_BASE_URL=http://host:port/sessions/manual-cc-001 +export ANTHROPIC_API_KEY=dummy-local-key +unset ANTHROPIC_AUTH_TOKEN +unset CLAUDE_CODE_API_BASE_URL +unset CLAUDE_CODE_API_KEY +export NO_PROXY=host,127.0.0.1,localhost + +claude --bare --model claude-sonnet-4-5 +``` + +Do not use `-p` / `--print` for a manual interactive session. `-p` sends one +prompt, prints the response, and exits. + +When finished, return to the launcher terminal and press `Enter` or `Ctrl-C`. +This is the normal finalize path: the launcher finalizes the session and writes +trajectories, session metadata, and the debug snapshot. + +## One-Shot Claude Code Smoke + +Use `--run-claude` when you want the launcher to run one Claude Code prompt and +then finalize automatically: + +```bash +python examples/gateway/debug_launcher.py \ + --backend openai-completions \ + --backend-base-url http://127.0.0.1:38315/v1 \ + --backend-model /data1/models/Qwen/Qwen3.5-4B \ + --tokenizer /data1/models/Qwen/Qwen3.5-4B \ + --session-id claude-smoke-001 \ + --output-dir /tmp/uni-agent-gateway-debug \ + --run-claude \ + --claude-prompt "Reply with OK only." \ + --claude-model claude-sonnet-4-5 \ + --claude-timeout 180 \ + --response-length 64 +``` + +The equivalent manual one-shot command is: + +```bash +claude --bare --no-session-persistence \ + -p "Reply with OK only." \ + --output-format text \ + --model claude-sonnet-4-5 +``` + +`--response-length` is important for small local backends. Claude Code can ask +for a large `max_tokens` value such as `32000`; the session response budget +clamps that request before it reaches the backend. + +## Direct API Smoke + +You can also call the Anthropic Messages endpoint directly: + +```bash +curl "$ANTHROPIC_BASE_URL/v1/messages" \ + -H "x-api-key: dummy-local-key" \ + -H "content-type: application/json" \ + -d '{ + "model": "claude-sonnet-4-5", + "max_tokens": 32, + "messages": [{"role": "user", "content": "Reply with OK only."}] + }' +``` + +The OpenAI-compatible URL printed by the launcher keeps the `/v1` suffix: + +```text +http://host:port/sessions/manual-cc-001/v1 +``` + +The Claude Code base URL must not include `/v1`: + +```text +http://host:port/sessions/manual-cc-001 +``` + +Do not use `CLAUDE_CODE_API_BASE_URL` for this Anthropic Messages proxy path. + +## Scope and Known Gaps + +The launcher is a local debug entrypoint, not a general standalone gateway +service. It creates one session up front and writes artifacts only when that +session is finalized. + +The documented smoke paths cover connection setup, text generation, trajectory +writing, and Claude Code CLI availability. Gateway-level automated tests cover +OpenAI tool calls and Anthropic `tool_use` round trips, but this launcher does +not currently provide a full Claude Code tool-workflow recipe. Use a dedicated +Claude Code SWE recipe or manual smoke when validating complex tool behavior. + +## Model Arguments + +`--model claude-sonnet-4-5` in a manual Claude Code command, and +`--claude-model claude-sonnet-4-5` in one-shot mode, are Claude Code frontend +model strings. They are not the local inference backend model. + +For local collection: + +- Claude Code uses that value when constructing the Anthropic Messages request. +- The gateway forwards the request through the local backend adapter. +- The actual local inference model is selected by `--backend-model`. + +`claude-sonnet-4-5` was accepted by the locally tested Claude Code 2.1.177 +client. If a different Claude Code version rejects that name before sending the +request, use a model name accepted by that CLI version. The local backend will +still be controlled by `--backend-model`. + +## Output Files + +For `--output-dir /tmp/uni-agent-gateway-debug` and +`--session-id manual-cc-001`, outputs are written under: + +```text +/tmp/uni-agent-gateway-debug/manual-cc-001/ +``` + +Files: + +- `trajectories.jsonl`: one JSON record per finalized trajectory. +- `session_metadata.json`: launcher metadata, provider URLs, Claude return code, + artifact paths, and trajectory count. This exists even when zero trajectories + were recorded. +- `debug_snapshot.json`: provider URLs, normalized message history, + active tool schemas, and session state captured immediately before + finalization. This is a diagnostic artifact, not a training trajectory schema. +- `claude.stdout`: captured Claude Code stdout in one-shot mode. +- `claude.stderr`: captured Claude Code stderr in one-shot mode. +- `claude-debug.log`: Claude Code debug log in one-shot mode. + +Each trajectory record has `schema_version: 1` and stores the finalized +`Trajectory` fields: + +- `prompt_ids` +- `response_ids` +- `response_mask` +- `response_logprobs` +- `reward_info` +- `reward_score` +- `num_turns` +- `multi_modal_data` +- `extra_fields` + +The trajectory records intentionally store token-level training data. Use +`debug_snapshot.json` when you need the normalized conversation history that +produced those tokens. + +## Fake Backend + +For a no-GPU smoke: + +```bash +python examples/gateway/debug_launcher.py \ + --backend fake \ + --session-id fake-debug-001 \ + --output-dir /tmp/uni-agent-gateway-debug +``` + +The fake backend always returns token ids for `OK`. It is useful for checking +Claude Code connectivity and trajectory writing without spending Anthropic API +quota or starting a local inference server. + +## Common Failures + +If Claude Code does not hit the gateway: + +- Use `ANTHROPIC_BASE_URL=http://host:port/sessions/{session_id}`. +- Do not append `/v1` to `ANTHROPIC_BASE_URL`. +- Unset `CLAUDE_CODE_API_BASE_URL`, `CLAUDE_CODE_API_KEY`, and + `ANTHROPIC_AUTH_TOKEN`. +- Set `NO_PROXY` for the printed host, `127.0.0.1`, and `localhost`. + +If vLLM rejects `max_tokens`: + +- Add `--response-length 64` or another local budget appropriate for the model. + +If Claude Code prints `API Error: 500 ConnectError: All connection attempts failed`: + +- Claude Code reached the gateway. +- The gateway could not connect to `--backend-base-url`. +- Start vLLM, or fix the host/port, then verify `/v1/models` returns 200. + +If `trajectories.jsonl` is empty: + +- Check `session_metadata.json` first. +- Inspect `debug_snapshot.json` to confirm which messages reached the gateway. +- Then inspect `claude.stdout`, `claude.stderr`, and `claude-debug.log` if + one-shot mode was used. + +Never print or commit real API keys. Use `dummy-local-key` for local gateway +debugging. diff --git a/examples/gateway/debug_launcher.py b/examples/gateway/debug_launcher.py new file mode 100644 index 00000000..b9a12fb8 --- /dev/null +++ b/examples/gateway/debug_launcher.py @@ -0,0 +1,692 @@ +"""Standalone gateway debug launcher for Claude Code and OpenAI-compatible clients.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import signal +import subprocess +import sys +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import httpx + +from uni_agent.gateway.config import GatewayActorConfig +from uni_agent.gateway.gateway import _GatewayActor +from uni_agent.gateway.session import SessionHandle, Trajectory +from verl.workers.rollout.replica import TokenOutput + +DEFAULT_CLAUDE_PROMPT = "Reply with OK only." +DEFAULT_CLAUDE_MODEL = "claude-sonnet-4-5" +DEFAULT_ANTHROPIC_API_KEY = "dummy-local-key" + + +class DebugFakeTokenizer: + """Small tokenizer that keeps this launcher independent of test helpers.""" + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + tokenize: bool = True, + add_generation_prompt: bool = True, + tools: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> list[int] | str: + parts = [] + for message in messages: + parts.append(f"{message['role']}:{self._normalize_content(message.get('content', ''))}\n") + if tools: + parts.append(f"tools:{json.dumps(tools, sort_keys=True)}\n") + if add_generation_prompt: + parts.append("assistant:") + text = "".join(parts) + if tokenize: + return self.encode(text, add_special_tokens=False) + return text + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return [ord(char) for char in text] + + def decode(self, token_ids: Any, skip_special_tokens: bool = True) -> str: + if hasattr(token_ids, "tolist"): + token_ids = token_ids.tolist() + return "".join(chr(int(token_id.item() if hasattr(token_id, "item") else token_id)) for token_id in token_ids) + + def _normalize_content(self, content: Any) -> str: + if isinstance(content, list): + parts = [] + for part in content: + if isinstance(part, dict): + parts.append(str(part.get("text", ""))) + else: + parts.append(str(part)) + return "".join(parts) + if content is None: + return "" + return str(content) + + +class DebugFakeBackend: + """Debug backend that always returns token ids for the text OK.""" + + def __init__(self): + self.calls: list[dict[str, Any]] = [] + + async def generate( + self, + request_id: str, + *, + prompt_ids: list[int], + sampling_params: dict[str, Any], + image_data: list[Any] | None = None, + video_data: list[Any] | None = None, + ) -> TokenOutput: + self.calls.append( + { + "request_id": request_id, + "prompt_ids": list(prompt_ids), + "sampling_params": dict(sampling_params), + "image_data": image_data, + "video_data": video_data, + } + ) + return TokenOutput( + token_ids=[ord("O"), ord("K")], + log_probs=[0.0, 0.0], + stop_reason="completed", + ) + + +def _token_ids_from_template_result(result: Any) -> Any: + if hasattr(result, "ids"): + return list(result.ids) + if isinstance(result, list): + token_ids: list[int] = [] + saw_encoding = False + for item in result: + if hasattr(item, "ids"): + token_ids.extend(list(item.ids)) + saw_encoding = True + else: + token_ids.append(item) + if saw_encoding: + return token_ids + return result + + +class TemplateResultTokenIdsWrapper: + """Tokenizer proxy that flattens tokenizers Encoding template results.""" + + def __init__(self, tokenizer: Any): + self._tokenizer = tokenizer + + def apply_chat_template(self, *args: Any, **kwargs: Any) -> Any: + result = self._tokenizer.apply_chat_template(*args, **kwargs) + if kwargs.get("tokenize", True): + return _token_ids_from_template_result(result) + return result + + def __getattr__(self, name: str) -> Any: + return getattr(self._tokenizer, name) + + +class OpenAICompletionsBackend: + """Backend adapter for vLLM-style OpenAI completions with token ids enabled.""" + + def __init__(self, *, backend_base_url: str, backend_model: str, timeout: float = 60.0): + self._completions_url = f"{backend_base_url.rstrip('/')}/completions" + self._backend_model = backend_model + self._timeout = timeout + + async def generate( + self, + request_id: str, + *, + prompt_ids: list[int], + sampling_params: dict[str, Any], + image_data: list[Any] | None = None, + video_data: list[Any] | None = None, + ) -> TokenOutput: + payload: dict[str, Any] = { + "model": self._backend_model, + "prompt": list(prompt_ids), + "max_tokens": sampling_params.get("max_tokens", 16), + "request_id": request_id, + "return_token_ids": True, + "logprobs": 1, + "add_special_tokens": False, + } + for key in ("temperature", "top_p", "top_k", "stop"): + if key in sampling_params: + payload[key] = sampling_params[key] + + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.post(self._completions_url, json=payload) + if getattr(response, "is_error", False): + status_code = getattr(response, "status_code", "unknown") + body = getattr(response, "text", "") + raise RuntimeError( + f"OpenAI completions backend request {request_id} failed with HTTP {status_code}: {body}" + ) + response.raise_for_status() + body = response.json() + choices = body.get("choices") if isinstance(body, dict) else None + if not choices or not isinstance(choices, list): + raise RuntimeError("OpenAI completions backend returned no choices") + choice = choices[0] + token_ids = choice.get("token_ids") if isinstance(choice, dict) else None + if token_ids is None: + raise RuntimeError("OpenAI completions backend must support vLLM return_token_ids=true") + + log_probs = None + logprobs = choice.get("logprobs") + if isinstance(logprobs, dict): + token_logprobs = logprobs.get("token_logprobs") + if isinstance(token_logprobs, list) and len(token_logprobs) == len(token_ids): + log_probs = list(token_logprobs) + stop_reason = "length" if choice.get("finish_reason") == "length" else "completed" + return TokenOutput(token_ids=list(token_ids), log_probs=log_probs, stop_reason=stop_reason) + + +@dataclass(frozen=True) +class DebugSessionResult: + session_id: str + output_path: Path + metadata_path: Path + debug_snapshot_path: Path + trajectories_count: int + provider_urls: dict[str, str] + claude_returncode: int | None = None + + +def utc_now() -> str: + return datetime.now(UTC).isoformat().replace("+00:00", "Z") + + +def trajectory_to_record( + *, + session_id: str, + trajectory_index: int, + trajectory: Trajectory, + metadata: dict[str, Any] | None = None, + created_at: str | None = None, +) -> dict[str, Any]: + return { + "schema_version": 1, + "session_id": session_id, + "trajectory_index": trajectory_index, + "created_at": created_at or utc_now(), + "metadata": dict(metadata or {}), + "trajectory": { + "prompt_ids": list(trajectory.prompt_ids), + "response_ids": list(trajectory.response_ids), + "response_mask": list(trajectory.response_mask), + "response_logprobs": ( + list(trajectory.response_logprobs) if trajectory.response_logprobs is not None else None + ), + "reward_info": dict(trajectory.reward_info), + "reward_score": trajectory.reward_score, + "num_turns": trajectory.num_turns, + "multi_modal_data": trajectory.multi_modal_data, + "extra_fields": dict(trajectory.extra_fields), + }, + } + + +def write_trajectories_jsonl( + *, + output_dir: Path, + session_id: str, + trajectories: list[Trajectory], + metadata: dict[str, Any] | None = None, + created_at: str | None = None, +) -> Path: + session_dir = output_dir / session_id + session_dir.mkdir(parents=True, exist_ok=True) + path = session_dir / "trajectories.jsonl" + timestamp = created_at or utc_now() + with path.open("w", encoding="utf-8") as f: + for index, trajectory in enumerate(trajectories): + record = trajectory_to_record( + session_id=session_id, + trajectory_index=index, + trajectory=trajectory, + metadata=metadata, + created_at=timestamp, + ) + f.write(json.dumps(record, ensure_ascii=False, sort_keys=True)) + f.write("\n") + return path + + +def write_session_metadata_json( + *, + output_dir: Path, + session_id: str, + metadata: dict[str, Any], + provider_urls: dict[str, str], + trajectories_count: int, + trajectories_path: Path, + debug_snapshot_path: Path, + created_at: str | None = None, +) -> Path: + session_dir = output_dir / session_id + session_dir.mkdir(parents=True, exist_ok=True) + path = session_dir / "session_metadata.json" + record = { + "schema_version": 1, + "session_id": session_id, + "created_at": created_at or utc_now(), + "metadata": dict(metadata), + "provider_urls": dict(provider_urls), + "trajectories_count": trajectories_count, + "trajectories_path": str(trajectories_path), + "debug_snapshot_path": str(debug_snapshot_path), + } + path.write_text(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8") + return path + + +def write_debug_snapshot_json( + *, + output_dir: Path, + session_id: str, + snapshot: dict[str, Any], + created_at: str | None = None, +) -> Path: + session_dir = output_dir / session_id + session_dir.mkdir(parents=True, exist_ok=True) + path = session_dir / "debug_snapshot.json" + record = { + **snapshot, + "schema_version": 1, + "session_id": session_id, + "created_at": created_at or utc_now(), + } + path.write_text(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8") + return path + + +def provider_urls(handle: SessionHandle) -> dict[str, str]: + if not handle.base_url: + raise RuntimeError("Session handle does not include a base_url") + base_url = handle.base_url.rstrip("/") + anthropic_base_url = base_url.removesuffix("/v1") + return { + "anthropic_base_url": anthropic_base_url, + "openai_base_url": base_url, + "reward_info_url": handle.reward_info_url, + } + + +def build_claude_env( + base_env: dict[str, str], + *, + anthropic_base_url: str, + anthropic_api_key: str, +) -> dict[str, str]: + env = {} + for key, value in base_env.items(): + lower_key = key.lower() + if key in {"ANTHROPIC_AUTH_TOKEN", "CLAUDE_CODE_API_BASE_URL", "CLAUDE_CODE_API_KEY"}: + continue + if "proxy" in lower_key: + continue + env[key] = value + + host = urlparse(anthropic_base_url).hostname or "localhost" + env.update( + { + "ANTHROPIC_BASE_URL": anthropic_base_url, + "ANTHROPIC_API_KEY": anthropic_api_key, + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "NO_PROXY": f"{host},127.0.0.1,localhost", + } + ) + return env + + +def build_claude_command(*, prompt: str, model: str, debug_file: Path) -> list[str]: + return [ + "claude", + "--bare", + "--no-session-persistence", + "--debug-file", + str(debug_file), + "-p", + prompt, + "--output-format", + "text", + "--model", + model, + ] + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend", choices=("fake", "openai-completions"), default="fake") + parser.add_argument("--session-id", required=True) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--run-claude", action="store_true") + parser.add_argument("--claude-prompt", default=DEFAULT_CLAUDE_PROMPT) + parser.add_argument("--claude-model", default=DEFAULT_CLAUDE_MODEL) + parser.add_argument("--claude-timeout", default=300.0, type=float) + parser.add_argument("--anthropic-api-key", default=DEFAULT_ANTHROPIC_API_KEY) + parser.add_argument("--backend-base-url") + parser.add_argument("--backend-model") + parser.add_argument("--tokenizer") + parser.add_argument("--prompt-length", type=int) + parser.add_argument("--response-length", type=int) + args = parser.parse_args(argv) + + if args.backend == "openai-completions": + missing = [ + flag + for flag, value in ( + ("--backend-base-url", args.backend_base_url), + ("--backend-model", args.backend_model), + ("--tokenizer", args.tokenizer), + ) + if not value + ] + if missing: + parser.error(f"--backend openai-completions requires {', '.join(missing)}") + return args + + +def build_tokenizer(args: argparse.Namespace): + if args.backend == "fake": + return DebugFakeTokenizer() + + from transformers import AutoTokenizer + + return TemplateResultTokenIdsWrapper(AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)) + + +def build_backend(args: argparse.Namespace): + if args.backend == "fake": + return DebugFakeBackend() + return OpenAICompletionsBackend( + backend_base_url=args.backend_base_url, + backend_model=args.backend_model, + ) + + +async def wait_for_manual_finalize(stdin: Any = sys.stdin) -> str: + loop = asyncio.get_running_loop() + stop_event = asyncio.Event() + trigger = "signal" + + def _stop(source: str) -> None: + nonlocal trigger + if not stop_event.is_set(): + trigger = source + stop_event.set() + + registered_signals = [] + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, _stop, "signal") + registered_signals.append(sig) + except NotImplementedError: + pass + stdin_task: asyncio.Task | None = None + stdin_fd = None + try: + if stdin is not None and stdin.isatty(): + print("Press Enter or Ctrl-C to finalize the session and write trajectories.") + try: + stdin_fd = stdin.fileno() + loop.add_reader(stdin_fd, _stop, "stdin") + except (AttributeError, OSError, NotImplementedError): + stdin_task = asyncio.create_task(asyncio.to_thread(stdin.readline)) + stdin_task.add_done_callback(lambda _: _stop("stdin")) + await stop_event.wait() + return trigger + finally: + if stdin_fd is not None: + try: + loop.remove_reader(stdin_fd) + except (OSError, NotImplementedError): + pass + if stdin_task is not None: + stdin_task.cancel() + for sig in registered_signals: + loop.remove_signal_handler(sig) + + +def capture_debug_snapshot(actor: _GatewayActor, session_id: str, urls: dict[str, str]) -> dict[str, Any]: + session = actor._sessions.get(session_id) + if session is None: + return { + "session_state": None, + "message_history": [], + "active_tool_schemas": None, + "provider_urls": dict(urls), + "note": "session was not available when debug snapshot was captured", + } + return { + "session_state": session.snapshot_state(), + "message_history": list(session.message_history), + "active_tool_schemas": session.active_tool_schemas, + "provider_urls": dict(urls), + } + + +async def run_claude_once( + *, + urls: dict[str, str], + output_dir: Path, + session_id: str, + claude_prompt: str, + claude_model: str, + anthropic_api_key: str, + claude_timeout: float, +) -> dict[str, Any]: + session_dir = output_dir / session_id + session_dir.mkdir(parents=True, exist_ok=True) + debug_file = session_dir / "claude-debug.log" + command = build_claude_command(prompt=claude_prompt, model=claude_model, debug_file=debug_file) + env = build_claude_env( + dict(os.environ), + anthropic_base_url=urls["anthropic_base_url"], + anthropic_api_key=anthropic_api_key, + ) + stdout_path = session_dir / "claude.stdout" + stderr_path = session_dir / "claude.stderr" + timed_out = False + try: + completed = await asyncio.to_thread( + subprocess.run, + command, + env=env, + check=False, + text=True, + capture_output=True, + timeout=claude_timeout, + ) + returncode = completed.returncode + stdout = completed.stdout or "" + stderr = completed.stderr or "" + except subprocess.TimeoutExpired as exc: + timed_out = True + returncode = 124 + stdout = exc.stdout if exc.stdout is not None else exc.output + stderr = exc.stderr + if isinstance(stdout, bytes): + stdout = stdout.decode(errors="replace") + elif stdout is None: + stdout = "" + else: + stdout = str(stdout) + if isinstance(stderr, bytes): + stderr = stderr.decode(errors="replace") + elif stderr is None: + stderr = "" + else: + stderr = str(stderr) + stdout_path.write_text(stdout, encoding="utf-8") + stderr_path.write_text(stderr, encoding="utf-8") + return { + "enabled": True, + "returncode": returncode, + "timed_out": timed_out, + "timeout_seconds": claude_timeout, + "stdout_path": str(stdout_path), + "stderr_path": str(stderr_path), + "debug_log_path": str(debug_file), + } + + +def print_provider_urls(urls: dict[str, str], *, anthropic_api_key: str) -> None: + api_key_display = anthropic_api_key if anthropic_api_key == DEFAULT_ANTHROPIC_API_KEY else "" + print("Claude Code:") + print(f" export ANTHROPIC_BASE_URL={urls['anthropic_base_url']}") + print(f" export ANTHROPIC_API_KEY={api_key_display}") + print("OpenAI-compatible:") + print(f" base_url={urls['openai_base_url']}") + print(f"Reward info: {urls['reward_info_url']}") + + +async def run_debug_session_once( + *, + session_id: str, + output_dir: Path, + backend: str = "fake", + run_claude: bool = False, + claude_prompt: str = DEFAULT_CLAUDE_PROMPT, + claude_model: str = DEFAULT_CLAUDE_MODEL, + anthropic_api_key: str = DEFAULT_ANTHROPIC_API_KEY, + claude_timeout: float = 300.0, + backend_base_url: str | None = None, + backend_model: str | None = None, + tokenizer: str | None = None, + prompt_length: int | None = None, + response_length: int | None = None, +) -> DebugSessionResult: + args = argparse.Namespace( + backend=backend, + backend_base_url=backend_base_url, + backend_model=backend_model, + tokenizer=tokenizer, + ) + gateway_backend = build_backend(args) + gateway_tokenizer = build_tokenizer(args) + actor = _GatewayActor( + GatewayActorConfig( + tokenizer=gateway_tokenizer, + prompt_length=prompt_length, + response_length=response_length, + ), + gateway_backend, + ) + + await actor.start() + trajectories: list[Trajectory] = [] + urls: dict[str, str] = {} + debug_snapshot: dict[str, Any] = { + "session_state": None, + "message_history": [], + "active_tool_schemas": None, + "provider_urls": {}, + "note": "session was not created before debug snapshot was captured", + } + claude_metadata: dict[str, Any] = { + "enabled": run_claude, + "returncode": None, + "timed_out": False, + "timeout_seconds": claude_timeout, + "stdout_path": None, + "stderr_path": None, + "debug_log_path": None, + } + try: + handle = await actor.create_session(session_id) + urls = provider_urls(handle) + print_provider_urls(urls, anthropic_api_key=anthropic_api_key) + if run_claude: + claude_metadata = await run_claude_once( + urls=urls, + output_dir=output_dir, + session_id=session_id, + claude_prompt=claude_prompt, + claude_model=claude_model, + anthropic_api_key=anthropic_api_key, + claude_timeout=claude_timeout, + ) + else: + await wait_for_manual_finalize() + debug_snapshot = capture_debug_snapshot(actor, session_id, urls) + trajectories = await actor.finalize_session(session_id) + finally: + await actor.shutdown() + + metadata = {"backend": backend, "claude": claude_metadata} + output_path = write_trajectories_jsonl( + output_dir=output_dir, + session_id=session_id, + trajectories=trajectories, + metadata=metadata, + ) + debug_snapshot_path = write_debug_snapshot_json( + output_dir=output_dir, + session_id=session_id, + snapshot=debug_snapshot, + ) + metadata_path = write_session_metadata_json( + output_dir=output_dir, + session_id=session_id, + metadata=metadata, + provider_urls=urls, + trajectories_count=len(trajectories), + trajectories_path=output_path, + debug_snapshot_path=debug_snapshot_path, + ) + return DebugSessionResult( + session_id=session_id, + output_path=output_path, + metadata_path=metadata_path, + debug_snapshot_path=debug_snapshot_path, + trajectories_count=len(trajectories), + provider_urls=urls, + claude_returncode=claude_metadata["returncode"], + ) + + +async def async_main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + result = await run_debug_session_once( + session_id=args.session_id, + output_dir=args.output_dir, + backend=args.backend, + run_claude=args.run_claude, + claude_prompt=args.claude_prompt, + claude_model=args.claude_model, + claude_timeout=args.claude_timeout, + anthropic_api_key=args.anthropic_api_key, + backend_base_url=args.backend_base_url, + backend_model=args.backend_model, + tokenizer=args.tokenizer, + prompt_length=args.prompt_length, + response_length=args.response_length, + ) + print(f"Wrote {result.trajectories_count} trajectories to {result.output_path}") + print(f"Wrote debug snapshot to {result.debug_snapshot_path}") + print(f"Wrote session metadata to {result.metadata_path}") + if result.claude_returncode not in (None, 0): + return result.claude_returncode + return 0 + + +def main(argv: list[str] | None = None) -> int: + return asyncio.run(async_main(argv)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/uni_agent/gateway/adapters/test_anthropic_adapter.py b/tests/uni_agent/gateway/adapters/test_anthropic_adapter.py new file mode 100644 index 00000000..a6944272 --- /dev/null +++ b/tests/uni_agent/gateway/adapters/test_anthropic_adapter.py @@ -0,0 +1,593 @@ +import json + +import pytest + +from uni_agent.gateway.adapters.anthropic import anthropic_to_internal +from uni_agent.gateway.adapters.types import MalformedRequestError + +ALLOWED_SAMPLING_KEYS = frozenset({"temperature", "top_p", "top_k", "max_tokens", "stop"}) +BASE = dict(base_sampling_params={}, allowed_sampling_keys=ALLOWED_SAMPLING_KEYS) + + +def test_system_string_becomes_first_message(): + """Top-level Anthropic system text becomes the leading internal system + message and max_tokens is forwarded into sampling params.""" + req = anthropic_to_internal( + {"system": "Be concise.", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}, + **BASE, + ) + assert req["messages"][0] == {"role": "system", "content": "Be concise."} + assert req["sampling_params"]["max_tokens"] == 16 + + +def test_user_text_image_text_order_preserved(): + """Mixed Anthropic user text/image blocks lower to OpenAI-compatible parts + without reordering the multimodal content.""" + req = anthropic_to_internal( + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "a"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}}, + {"type": "text", "text": "b"}, + ], + } + ], + "max_tokens": 8, + }, + **BASE, + ) + assert req["messages"][0]["content"] == [ + {"type": "text", "text": "a"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + {"type": "text", "text": "b"}, + ] + + +def test_assistant_tool_use_to_tool_calls_dict_args(): + """Anthropic assistant tool_use blocks become internal OpenAI-style + tool_calls while preserving dict arguments.""" + req = anthropic_to_internal( + { + "messages": [ + {"role": "user", "content": "go"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {"q": "x"}}], + }, + ], + "max_tokens": 8, + }, + **BASE, + ) + tc = req["messages"][-1]["tool_calls"][0] + assert tc["function"]["name"] == "lookup" + assert tc["function"]["arguments"] == {"q": "x"} + + +def test_malformed_anthropic_content_blocks_rejected(): + """Malformed or unsupported Anthropic content blocks fail at the adapter + boundary before the session codec sees corrupted history.""" + payloads = [ + {"messages": [{"role": "user", "content": [{"type": "audio", "data": "x"}]}], "max_tokens": 8}, + { + "messages": [ + {"role": "assistant", "content": [{"type": "tool_use", "name": "lookup", "input": {"q": "x"}}]} + ], + "max_tokens": 8, + }, + { + "messages": [ + {"role": "assistant", "content": [{"type": "tool_use", "id": 123, "name": "lookup", "input": {"q": "x"}}]} + ], + "max_tokens": 8, + }, + { + "messages": [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": "bad"}], + } + ], + "max_tokens": 8, + }, + {"messages": [{"role": "assistant", "content": [{"type": "citation", "text": "x"}]}], "max_tokens": 8}, + ] + + for payload in payloads: + with pytest.raises(MalformedRequestError): + anthropic_to_internal(payload, **BASE) + + +def test_user_tool_result_text_becomes_tool_message(): + """A text-only Anthropic tool_result lowers to an internal tool message with + the original tool_use_id.""" + req = anthropic_to_internal( + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [{"type": "text", "text": "found"}], + } + ], + } + ], + "max_tokens": 8, + }, + **BASE, + ) + assert req["messages"][0]["role"] == "tool" + assert req["messages"][0]["content"] == "found" + + +def test_malformed_tool_result_blocks_rejected(): + """Tool results without a usable id or with unsupported nested blocks are + rejected instead of corrupting tool-return history.""" + payloads = [ + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "", "content": [{"type": "text", "text": "found"}]} + ], + } + ], + "max_tokens": 8, + }, + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [{"type": "json", "value": {}}], + } + ], + } + ], + "max_tokens": 8, + }, + ] + + for payload in payloads: + with pytest.raises(MalformedRequestError): + anthropic_to_internal(payload, **BASE) + + +def test_tool_result_image_appends_user_message(): + """Tool-result images are preserved as a following user image message while + text remains in the correlated tool message.""" + req = anthropic_to_internal( + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t", + "content": [ + {"type": "text", "text": "see"}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "BBBB"}, + }, + ], + } + ], + } + ], + "max_tokens": 8, + }, + **BASE, + ) + assert req["messages"][0]["role"] == "tool" + assert req["messages"][1]["role"] == "user" + assert req["messages"][1]["content"][0]["type"] == "image_url" + + +def test_tool_result_image_only_preserves_empty_tool_message(): + """Image-only tool_result content still emits an empty tool message sentinel + before the downgraded user image message.""" + req = anthropic_to_internal( + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t", + "content": [ + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "BBBB"}, + } + ], + } + ], + } + ], + "max_tokens": 8, + }, + **BASE, + ) + assert req["messages"][0] == {"role": "tool", "tool_call_id": "t", "content": ""} + assert req["messages"][1]["role"] == "user" + assert req["messages"][1]["content"][0]["type"] == "image_url" + + +def test_anthropic_request_control_fields_lower_or_reject(): + """Small Anthropic request-level controls stay adapter-owned: unsupported + tool choices are rejected, while stop_sequences lowers to sampling stop.""" + with pytest.raises(MalformedRequestError): + anthropic_to_internal( + {"messages": [{"role": "user", "content": "hi"}], "tool_choice": {"type": "any"}, "max_tokens": 8}, + **BASE, + ) + + req = anthropic_to_internal( + {"messages": [{"role": "user", "content": "hi"}], "stop_sequences": [""], "max_tokens": 8}, + **BASE, + ) + assert req["sampling_params"]["stop"] == [""] + + +def test_mid_list_system_folded_into_user(): + """Mid-conversation system reminders are folded into user content so chat + templates never see a system role after the first message.""" + req = anthropic_to_internal( + { + "messages": [ + {"role": "user", "content": "a"}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": "b"}, + ], + "max_tokens": 8, + }, + **BASE, + ) + roles = [m["role"] for m in req["messages"]] + assert len(req["messages"]) == 2 + assert roles == ["user", "user"] + assert "system" not in roles[1:] + assert any("reminder" in str(m.get("content")) for m in req["messages"]) + assert not any(m.get("content") == "\nreminder\n" for m in req["messages"]) + + +def test_mid_list_system_before_tool_result_does_not_cross_tool_message(): + """A system reminder before a tool_result folds into the prior user message + and does not break the assistant/tool message adjacency.""" + req = anthropic_to_internal( + { + "messages": [ + {"role": "user", "content": "a"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {"q": "x"}}], + }, + {"role": "system", "content": "reminder"}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [{"type": "text", "text": "found"}], + } + ], + }, + ], + "max_tokens": 8, + }, + **BASE, + ) + assert req["messages"][0]["role"] == "user" + assert "reminder" in req["messages"][0]["content"] + assert req["messages"][-1]["role"] == "tool" + assert not any( + m["role"] == "user" and m.get("content") == "\nreminder\n" + for m in req["messages"][1:] + ) + + +def test_billing_header_stripped_from_system(): + """Claude Code billing headers embedded in system text are stripped before + prompt construction.""" + req = anthropic_to_internal( + { + "system": "x-anthropic-billing-header: cch=abc\nBe concise.", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8, + }, + **BASE, + ) + assert req["messages"][0] == {"role": "system", "content": "Be concise."} + + +def test_provider_specific_tool_metadata_lowers_to_template_function_schema(): + """Provider-specific Anthropic tool metadata is ignored while the tool name + remains available to the local chat template as an OpenAI function schema.""" + req = anthropic_to_internal( + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8, + "tools": [ + { + "type": "web_search_20250305", + "name": "web_search", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + **BASE, + ) + + assert req["tools"] == [{"type": "function", "function": {"name": "web_search", "parameters": {}}}] + + +def test_tool_input_schema_passes_through_json_schema(): + """Anthropic JSON schema is copied into OpenAI function parameters without + gateway-specific parser normalization.""" + schema = { + "type": "object", + "properties": { + "target": {"anyOf": [{"const": "file"}, {"type": "string"}]}, + "mode": {"anyOf": [{"const": "read"}, {"const": "write"}]}, + "nested": {"type": "object", "properties": {"x": {"type": "integer"}}}, + }, + "required": ["target"], + } + req = anthropic_to_internal( + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8, + "tools": [{"name": "edit", "description": "edit files", "input_schema": schema}], + }, + **BASE, + ) + params = req["tools"][0]["function"]["parameters"] + assert params is not schema + assert params == schema + assert "type" not in params["properties"]["target"] + assert "enum" not in params["properties"]["target"] + + +def test_tool_input_schema_anyof_preserves_heterogeneous_branches_without_inferred_type(): + """Heterogeneous anyOf branches keep their distinct JSON types instead of + being collapsed into a compatibility hint for a specific parser.""" + schema = { + "type": "object", + "properties": { + "value": {"anyOf": [{"type": "integer"}, {"type": "number"}]}, + }, + } + req = anthropic_to_internal( + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8, + "tools": [{"name": "edit", "description": "edit files", "input_schema": schema}], + }, + **BASE, + ) + assert req["tools"][0]["function"]["parameters"]["properties"]["value"] == { + "anyOf": [{"type": "integer"}, {"type": "number"}] + } + + +def test_thinking_block_dropped_with_warning(caplog): + """Inbound Anthropic thinking blocks are dropped with a warning while + neighboring assistant text remains in the prompt history.""" + with caplog.at_level("WARNING", logger="gateway"): + req = anthropic_to_internal( + { + "messages": [ + {"role": "user", "content": "go"}, + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "reasoning..."}, {"type": "text", "text": "done"}], + }, + ], + "max_tokens": 8, + }, + **BASE, + ) + assert req["messages"][-1]["content"] == "done" + assert any("thinking" in r.message for r in caplog.records) + + +def test_unsupported_system_and_redacted_thinking_blocks_rejected(): + """System non-text blocks and encrypted redacted_thinking blocks are rejected + because neither has a faithful local template representation.""" + with pytest.raises(MalformedRequestError): + anthropic_to_internal( + { + "system": [{"type": "audio", "data": "x"}], + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8, + }, + **BASE, + ) + + with pytest.raises(MalformedRequestError): + anthropic_to_internal( + { + "messages": [ + {"role": "user", "content": "go"}, + { + "role": "assistant", + "content": [{"type": "redacted_thinking", "data": "xxx"}, {"type": "text", "text": "done"}], + }, + ], + "max_tokens": 8, + }, + **BASE, + ) + + +def test_anthropic_build_response_text(): + """A text GenerationOutcome serializes to the Anthropic Messages response + shape with end_turn stop reason and usage counts.""" + from uni_agent.gateway.adapters.anthropic import anthropic_build_response + from uni_agent.gateway.session.session import GenerationOutcome + + body = anthropic_build_response( + GenerationOutcome( + assistant_msg={"role": "assistant", "content": "hi"}, + finish_reason="stop", + prompt_tokens=3, + completion_tokens=1, + ), + model="claude-x", + ) + assert body["type"] == "message" + assert body["role"] == "assistant" + assert body["content"] == [{"type": "text", "text": "hi"}] + assert body["stop_reason"] == "end_turn" + assert body["usage"] == {"input_tokens": 3, "output_tokens": 1} + + +def test_anthropic_build_response_tool_use(): + """Internal OpenAI-style tool_calls serialize back to Anthropic tool_use + blocks and map finish_reason to tool_use.""" + from uni_agent.gateway.adapters.anthropic import anthropic_build_response + from uni_agent.gateway.session.session import GenerationOutcome + + body = anthropic_build_response( + GenerationOutcome( + assistant_msg={ + "role": "assistant", + "content": "", + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": {"a": 1}}}], + }, + finish_reason="tool_calls", + prompt_tokens=2, + completion_tokens=4, + ), + model="claude-x", + ) + tu = [b for b in body["content"] if b["type"] == "tool_use"][0] + assert tu == {"type": "tool_use", "id": "c1", "name": "f", "input": {"a": 1}} + assert body["stop_reason"] == "tool_use" + + +def test_anthropic_build_response_length_maps_to_max_tokens(): + """Internal length finish reasons are exposed as Anthropic max_tokens stop + reasons.""" + from uni_agent.gateway.adapters.anthropic import anthropic_build_response + from uni_agent.gateway.session.session import GenerationOutcome + + body = anthropic_build_response( + GenerationOutcome( + assistant_msg={"role": "assistant", "content": ""}, + finish_reason="length", + prompt_tokens=1, + completion_tokens=0, + ), + model="claude-x", + ) + assert body["stop_reason"] == "max_tokens" + + +def test_anthropic_build_response_normalizes_tool_use_input(): + """Tool parser arguments are normalized for Anthropic clients: JSON-string + objects become dicts and invalid strings become empty input.""" + from uni_agent.gateway.adapters.anthropic import anthropic_build_response + from uni_agent.gateway.session.session import GenerationOutcome + + for arguments, expected_input in (('{"a": 1}', {"a": 1}), ("not json", {})): + body = anthropic_build_response( + GenerationOutcome( + assistant_msg={ + "role": "assistant", + "content": "", + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": arguments}}], + }, + finish_reason="tool_calls", + prompt_tokens=2, + completion_tokens=4, + ), + model="claude-x", + ) + tu = [b for b in body["content"] if b["type"] == "tool_use"][0] + assert tu["input"] == expected_input + + +@pytest.mark.asyncio +async def test_anthropic_stream_event_sequence(): + """A synthesized text response follows the Anthropic Messages SSE event + order and reports stop reason plus final usage in message_delta.""" + from uni_agent.gateway.adapters.anthropic import anthropic_stream_response + from uni_agent.gateway.session.session import GenerationOutcome + + resp = anthropic_stream_response( + GenerationOutcome( + assistant_msg={"role": "assistant", "content": "hi"}, + finish_reason="stop", + prompt_tokens=3, + completion_tokens=1, + ), + model="claude-x", + ) + assert resp.headers["cache-control"] == "no-cache" + assert resp.headers["connection"] == "keep-alive" + assert resp.headers["x-accel-buffering"] == "no" + + text = (b"".join([c async for c in resp.body_iterator])).decode() + assert "event: message_start" in text + assert "event: content_block_start" in text + assert "text_delta" in text + assert "event: content_block_stop" in text + assert "event: message_delta" in text + assert '"stop_reason": "end_turn"' in text or '"stop_reason":"end_turn"' in text + assert ( + '"usage": {"input_tokens": 3, "output_tokens": 1}' in text + or '"usage":{"input_tokens":3,"output_tokens":1}' in text + ) + assert "event: message_stop" in text + + +@pytest.mark.asyncio +async def test_anthropic_stream_response_emits_tool_use_delta(): + """A synthesized tool-call response emits Anthropic tool_use start data and + streams the tool input through an input_json_delta event.""" + from uni_agent.gateway.adapters.anthropic import anthropic_stream_response + from uni_agent.gateway.session.session import GenerationOutcome + + resp = anthropic_stream_response( + GenerationOutcome( + assistant_msg={ + "role": "assistant", + "content": "", + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": {"a": 1}}}], + }, + finish_reason="tool_calls", + prompt_tokens=2, + completion_tokens=4, + ), + model="claude-x", + ) + + text = (b"".join([c async for c in resp.body_iterator])).decode() + events = [] + for block in text.strip().split("\n\n"): + lines = block.splitlines() + events.append((lines[0].removeprefix("event: "), json.loads(lines[1].removeprefix("data: ")))) + + start = next(data for event, data in events if event == "content_block_start") + delta = next(data for event, data in events if event == "content_block_delta") + message_delta = next(data for event, data in events if event == "message_delta") + + assert start["content_block"] == {"type": "tool_use", "id": "c1", "name": "f", "input": {}} + assert delta["delta"]["type"] == "input_json_delta" + assert json.loads(delta["delta"]["partial_json"]) == {"a": 1} + assert message_delta["delta"]["stop_reason"] == "tool_use" diff --git a/tests/uni_agent/gateway/adapters/test_openai_adapter.py b/tests/uni_agent/gateway/adapters/test_openai_adapter.py new file mode 100644 index 00000000..cff4ad9e --- /dev/null +++ b/tests/uni_agent/gateway/adapters/test_openai_adapter.py @@ -0,0 +1,129 @@ +import json + +import pytest + +ALLOWED_SAMPLING_KEYS = frozenset({"temperature", "top_p", "top_k", "max_tokens", "stop"}) + + +def test_openai_build_response_shape(): + """A completed internal outcome serializes to the basic OpenAI chat + completion envelope with ids, model, choices, and token usage.""" + from uni_agent.gateway.adapters.openai import openai_build_response + from uni_agent.gateway.session.session import GenerationOutcome + + outcome = GenerationOutcome( + assistant_msg={"role": "assistant", "content": "hello"}, + finish_reason="stop", + prompt_tokens=3, + completion_tokens=2, + ) + body = openai_build_response(outcome, model="m") + assert body["id"].startswith("chatcmpl-") + assert body["object"] == "chat.completion" + assert isinstance(body["created"], int) + assert body["choices"][0]["message"]["content"] == "hello" + assert body["choices"][0]["finish_reason"] == "stop" + assert body["usage"]["total_tokens"] == 5 + assert body["model"] == "m" + + +@pytest.mark.asyncio +async def test_openai_stream_response_emits_compatible_sse_chunks(): + """A completed internal outcome is synthesized into OpenAI-compatible SSE + chunks for role, reasoning, content, tool calls, final usage, and DONE.""" + from uni_agent.gateway.adapters.openai import openai_stream_response + from uni_agent.gateway.session.session import GenerationOutcome + + resp = openai_stream_response( + GenerationOutcome( + assistant_msg={ + "role": "assistant", + "reasoning_content": "thinking", + "content": "hello", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "f", "arguments": '{"a":1}'}, + } + ], + }, + finish_reason="tool_calls", + prompt_tokens=3, + completion_tokens=2, + ), + model="m", + ) + + assert resp.headers["cache-control"] == "no-cache" + assert resp.headers["connection"] == "keep-alive" + assert resp.headers["x-accel-buffering"] == "no" + + text = (b"".join([chunk async for chunk in resp.body_iterator])).decode() + data_lines = [line.removeprefix("data: ") for line in text.splitlines() if line.startswith("data: ")] + assert data_lines[-1] == "[DONE]" + chunks = [json.loads(line) for line in data_lines[:-1]] + + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant"} + assert chunks[1]["choices"][0]["delta"] == {"reasoning_content": "thinking"} + assert chunks[2]["choices"][0]["delta"] == {"content": "hello"} + assert chunks[3]["choices"][0]["delta"]["tool_calls"][0]["index"] == 0 + assert chunks[3]["choices"][0]["delta"]["tool_calls"][0]["id"] == "call-1" + assert chunks[4]["choices"][0]["finish_reason"] == "tool_calls" + assert chunks[4]["usage"] == { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + } + + +def test_openai_to_internal_normalizes_messages_sampling_and_tools(): + """OpenAI wire requests lower to the internal shape: sampling params use the + gateway allowlist, JSON tool arguments are parsed, malformed argument + strings are preserved, and tool_choice=none clears tool schemas.""" + from uni_agent.gateway.adapters.openai import openai_to_internal + + payload = { + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + {"id": "x", "type": "function", "function": {"name": "f", "arguments": '{"x": 1}'}}, + {"id": "y", "type": "function", "function": {"name": "g", "arguments": "not json"}}, + ], + }, + ], + "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}], + "max_tokens": 32, + "temperature": 0.7, + "stop": [""], + "ignored_field": 1, + } + req = openai_to_internal( + payload, + base_sampling_params={"top_p": 0.9}, + allowed_sampling_keys=ALLOWED_SAMPLING_KEYS, + ) + assert set(req) == {"messages", "tools", "chat_template_kwargs", "sampling_params"} + assert req["messages"][0] == {"role": "user", "content": "hi"} + assert req["messages"][1]["tool_calls"][0]["function"]["arguments"] == {"x": 1} + assert req["messages"][1]["tool_calls"][1]["function"]["arguments"] == "not json" + assert req["tools"][0]["function"]["name"] == "f" + assert req["chat_template_kwargs"] == {} + assert req["sampling_params"]["max_tokens"] == 32 + assert req["sampling_params"]["temperature"] == 0.7 + assert req["sampling_params"]["top_p"] == 0.9 + assert req["sampling_params"]["stop"] == [""] + assert "ignored_field" not in req["sampling_params"] + + without_tools = openai_to_internal( + { + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}], + "tool_choice": "none", + }, + base_sampling_params={}, + allowed_sampling_keys=ALLOWED_SAMPLING_KEYS, + ) + assert without_tools["tools"] is None diff --git a/tests/uni_agent/gateway/test_debug_launcher.py b/tests/uni_agent/gateway/test_debug_launcher.py new file mode 100644 index 00000000..d6526fb9 --- /dev/null +++ b/tests/uni_agent/gateway/test_debug_launcher.py @@ -0,0 +1,588 @@ +import asyncio +import json +import os +import subprocess +from pathlib import Path + +import httpx +import pytest + +from examples.gateway import debug_launcher +from uni_agent.gateway.session import SessionHandle, Trajectory + + +def test_trajectory_to_record_is_json_serializable_and_preserves_fields(): + """Trajectory output records keep every training-visible field and can be + written as plain JSON.""" + trajectory = Trajectory( + prompt_ids=[1, 2], + response_ids=[3, 4], + response_mask=[1, 0], + response_logprobs=[-0.1, 0.0], + reward_info={"ok": True}, + reward_score=1.5, + num_turns=2, + multi_modal_data={"images": ["image://one"]}, + extra_fields={"finish_reason": "completed"}, + ) + + record = debug_launcher.trajectory_to_record( + session_id="debug-session", + trajectory_index=7, + trajectory=trajectory, + metadata={"backend": "fake"}, + created_at="2026-06-24T00:00:00Z", + ) + + assert json.loads(json.dumps(record)) == record + assert record == { + "schema_version": 1, + "session_id": "debug-session", + "trajectory_index": 7, + "created_at": "2026-06-24T00:00:00Z", + "metadata": {"backend": "fake"}, + "trajectory": { + "prompt_ids": [1, 2], + "response_ids": [3, 4], + "response_mask": [1, 0], + "response_logprobs": [-0.1, 0.0], + "reward_info": {"ok": True}, + "reward_score": 1.5, + "num_turns": 2, + "multi_modal_data": {"images": ["image://one"]}, + "extra_fields": {"finish_reason": "completed"}, + }, + } + + +def test_write_trajectories_jsonl_writes_expected_file(tmp_path): + """Finalized trajectories are written under the session output directory as + one JSON record per line with stable trajectory indexes.""" + trajectories = [ + Trajectory(prompt_ids=[1], response_ids=[2], response_mask=[1], num_turns=1), + Trajectory(prompt_ids=[3], response_ids=[4], response_mask=[1], response_logprobs=[-0.2], num_turns=2), + ] + + path = debug_launcher.write_trajectories_jsonl( + output_dir=tmp_path, + session_id="s1", + trajectories=trajectories, + metadata={"backend": "fake"}, + created_at="2026-06-24T01:02:03Z", + ) + + assert path == tmp_path / "s1" / "trajectories.jsonl" + lines = path.read_text().splitlines() + assert len(lines) == 2 + records = [json.loads(line) for line in lines] + assert records[0]["session_id"] == "s1" + assert records[0]["trajectory_index"] == 0 + assert records[0]["trajectory"]["prompt_ids"] == [1] + assert records[1]["trajectory_index"] == 1 + assert records[1]["trajectory"]["response_logprobs"] == [-0.2] + + +def test_write_session_metadata_json_writes_run_context_without_trajectories(tmp_path): + """Session metadata is written even when no trajectories were captured so a + failed smoke still leaves inspectable run context.""" + path = debug_launcher.write_session_metadata_json( + output_dir=tmp_path, + session_id="s-empty", + metadata={ + "backend": "fake", + "claude": { + "enabled": True, + "returncode": 1, + "stdout_path": "/tmp/stdout", + "stderr_path": "/tmp/stderr", + "debug_log_path": "/tmp/debug", + }, + }, + provider_urls={"anthropic_base_url": "http://127.0.0.1:9000/sessions/s-empty"}, + trajectories_count=0, + trajectories_path=tmp_path / "s-empty" / "trajectories.jsonl", + debug_snapshot_path=tmp_path / "s-empty" / "debug_snapshot.json", + created_at="2026-06-24T01:02:03Z", + ) + + assert path == tmp_path / "s-empty" / "session_metadata.json" + record = json.loads(path.read_text()) + assert record["schema_version"] == 1 + assert record["session_id"] == "s-empty" + assert record["metadata"]["claude"]["returncode"] == 1 + assert record["trajectories_count"] == 0 + assert record["trajectories_path"].endswith("trajectories.jsonl") + assert record["debug_snapshot_path"].endswith("debug_snapshot.json") + + +def test_write_debug_snapshot_json_writes_message_history_and_session_state(tmp_path): + """Debug snapshots preserve normalized conversation state separately from + token-level trajectories for post-run diagnosis.""" + path = debug_launcher.write_debug_snapshot_json( + output_dir=tmp_path, + session_id="s-debug", + snapshot={ + "session_state": {"phase": "ACTIVE", "has_active_trajectory": True}, + "message_history": [{"role": "user", "content": "hi"}], + }, + created_at="2026-06-24T01:02:03Z", + ) + + assert path == tmp_path / "s-debug" / "debug_snapshot.json" + record = json.loads(path.read_text()) + assert record["schema_version"] == 1 + assert record["session_id"] == "s-debug" + assert record["created_at"] == "2026-06-24T01:02:03Z" + assert record["session_state"]["has_active_trajectory"] is True + assert record["message_history"] == [{"role": "user", "content": "hi"}] + + +def test_provider_urls_strips_v1_for_anthropic_and_keeps_openai_base_url(): + """The printed Anthropic base URL omits /v1 for Claude Code while the + OpenAI-compatible URL keeps the session-scoped /v1 suffix.""" + handle = SessionHandle( + session_id="abc", + base_url="http://127.0.0.1:8000/sessions/abc/v1", + reward_info_url="http://127.0.0.1:8000/sessions/abc/reward_info", + ) + + assert debug_launcher.provider_urls(handle) == { + "anthropic_base_url": "http://127.0.0.1:8000/sessions/abc", + "openai_base_url": "http://127.0.0.1:8000/sessions/abc/v1", + "reward_info_url": "http://127.0.0.1:8000/sessions/abc/reward_info", + } + + +def test_build_claude_env_removes_auth_and_proxy_vars_and_sets_expected_vars(): + """Claude Code smoke env construction removes conflicting auth/proxy + variables and injects the local gateway Anthropic endpoint.""" + base_env = { + "PATH": "/bin", + "ANTHROPIC_AUTH_TOKEN": "real-token", + "CLAUDE_CODE_API_BASE_URL": "https://real.example", + "CLAUDE_CODE_API_KEY": "real-key", + "HTTP_PROXY": "http://proxy.example", + "custom_proxy_setting": "proxy", + "ANTHROPIC_API_KEY": "old-key", + } + + env = debug_launcher.build_claude_env( + base_env, + anthropic_base_url="http://10.1.2.3:9000/sessions/s1", + anthropic_api_key="dummy-local-key", + ) + + assert env["PATH"] == "/bin" + assert "ANTHROPIC_AUTH_TOKEN" not in env + assert "CLAUDE_CODE_API_BASE_URL" not in env + assert "CLAUDE_CODE_API_KEY" not in env + assert "HTTP_PROXY" not in env + assert "custom_proxy_setting" not in env + assert env["ANTHROPIC_BASE_URL"] == "http://10.1.2.3:9000/sessions/s1" + assert env["ANTHROPIC_API_KEY"] == "dummy-local-key" + assert env["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"] == "1" + assert env["NO_PROXY"] == "10.1.2.3,127.0.0.1,localhost" + + +def test_parse_args_minimal_and_openai_required_args_validation(tmp_path): + """The fake backend has minimal CLI requirements, while openai-completions + requires backend URL, backend model, and tokenizer path.""" + args = debug_launcher.parse_args(["--session-id", "s1", "--output-dir", str(tmp_path)]) + + assert args.backend == "fake" + assert args.session_id == "s1" + assert args.output_dir == tmp_path + assert args.anthropic_api_key == "dummy-local-key" + assert args.claude_timeout == 300.0 + + with pytest.raises(SystemExit): + debug_launcher.parse_args( + [ + "--backend", + "openai-completions", + "--session-id", + "s1", + "--output-dir", + str(tmp_path), + ] + ) + + args = debug_launcher.parse_args( + [ + "--backend", + "openai-completions", + "--session-id", + "s1", + "--output-dir", + str(tmp_path), + "--backend-base-url", + "http://127.0.0.1:8080", + "--backend-model", + "debug-model", + "--tokenizer", + "/tmp/tokenizer", + ] + ) + assert args.backend_base_url == "http://127.0.0.1:8080" + assert args.backend_model == "debug-model" + assert args.tokenizer == "/tmp/tokenizer" + + +def test_real_tokenizer_wrapper_flattens_encoding_template_results(): + """Real tokenizer wrappers flatten tokenizers Encoding objects returned by + apply_chat_template while delegating other tokenizer methods.""" + class FakeEncoding: + def __init__(self, ids): + self.ids = ids + + class FakeTokenizer: + def apply_chat_template(self, *args, **kwargs): + return [FakeEncoding([1, 2]), FakeEncoding([3])] + + def decode(self, token_ids, skip_special_tokens=True): + return "decoded" + + wrapper = debug_launcher.TemplateResultTokenIdsWrapper(FakeTokenizer()) + + assert wrapper.apply_chat_template([{"role": "user", "content": "hi"}], tokenize=True) == [1, 2, 3] + assert wrapper.decode([1, 2, 3]) == "decoded" + + +def test_build_claude_command(tmp_path): + """One-shot Claude Code smoke uses the bare, non-persistent CLI mode with + captured debug output, prompt text output, and the requested frontend model.""" + debug_file = tmp_path / "claude-debug.log" + + command = debug_launcher.build_claude_command( + prompt="Reply with OK only.", + model="claude-sonnet-4-5", + debug_file=debug_file, + ) + + assert command[0] == "claude" + assert "--bare" in command + assert "--no-session-persistence" in command + assert command[command.index("--debug-file") + 1] == str(debug_file) + assert command[command.index("-p") + 1] == "Reply with OK only." + assert command[command.index("--output-format") + 1] == "text" + assert command[command.index("--model") + 1] == "claude-sonnet-4-5" + + +@pytest.mark.asyncio +async def test_openai_completions_backend_generate_sends_prompt_token_ids_and_parses_token_ids_logprobs( + monkeypatch, +): + """The OpenAI completions debug backend sends token-id prompts with + return_token_ids enabled and maps token ids, logprobs, and finish reason.""" + requests = [] + + class FakeResponse: + def raise_for_status(self): + pass + + def json(self): + return { + "choices": [ + { + "token_ids": [79, 75], + "logprobs": {"token_logprobs": [-0.3, -0.4]}, + "finish_reason": "length", + } + ] + } + + class FakeAsyncClient: + def __init__(self, timeout): + self.timeout = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + async def post(self, url, json): + requests.append({"url": url, "json": json, "timeout": self.timeout}) + return FakeResponse() + + monkeypatch.setattr(debug_launcher.httpx, "AsyncClient", FakeAsyncClient) + backend = debug_launcher.OpenAICompletionsBackend( + backend_base_url="http://backend.example/", + backend_model="debug-model", + ) + + output = await backend.generate( + request_id="s1", + prompt_ids=[10, 11, 12], + sampling_params={ + "max_tokens": 8, + "temperature": 0.2, + "top_p": 0.9, + "top_k": 40, + "stop": ["END"], + "ignored": True, + }, + ) + + assert requests == [ + { + "url": "http://backend.example/completions", + "json": { + "model": "debug-model", + "prompt": [10, 11, 12], + "max_tokens": 8, + "request_id": "s1", + "temperature": 0.2, + "top_p": 0.9, + "top_k": 40, + "stop": ["END"], + "return_token_ids": True, + "logprobs": 1, + "add_special_tokens": False, + }, + "timeout": 60.0, + } + ] + assert output.token_ids == [79, 75] + assert output.log_probs == [-0.3, -0.4] + assert output.stop_reason == "length" + + +@pytest.mark.asyncio +async def test_openai_completions_backend_requires_token_ids(monkeypatch): + """A completions backend response without token_ids is rejected because + trajectory response tokens must come from backend token truth.""" + class FakeResponse: + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"text": "OK", "finish_reason": "stop"}]} + + class FakeAsyncClient: + def __init__(self, timeout): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + async def post(self, url, json): + return FakeResponse() + + monkeypatch.setattr(debug_launcher.httpx, "AsyncClient", FakeAsyncClient) + backend = debug_launcher.OpenAICompletionsBackend( + backend_base_url="http://backend.example", + backend_model="debug-model", + ) + + with pytest.raises(RuntimeError, match="return_token_ids=true"): + await backend.generate(request_id="s1", prompt_ids=[1], sampling_params={}) + + +@pytest.mark.asyncio +async def test_openai_completions_backend_http_error_includes_request_id_and_body(monkeypatch): + """HTTP errors from the completions backend include the gateway request id, + status code, and response body in the raised diagnostic.""" + class FakeResponse: + is_error = True + status_code = 400 + text = "invalid token ids" + + def raise_for_status(self): + pass + + def json(self): + return {} + + class FakeAsyncClient: + def __init__(self, timeout): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + async def post(self, url, json): + return FakeResponse() + + monkeypatch.setattr(debug_launcher.httpx, "AsyncClient", FakeAsyncClient) + backend = debug_launcher.OpenAICompletionsBackend( + backend_base_url="http://backend.example", + backend_model="debug-model", + ) + + with pytest.raises(RuntimeError, match="s1.*400.*invalid token ids"): + await backend.generate(request_id="s1", prompt_ids=[1], sampling_params={}) + + +@pytest.mark.asyncio +async def test_run_claude_once_records_timeout_metadata(monkeypatch, tmp_path): + """A timed-out one-shot Claude subprocess records timeout metadata and + preserves captured stdout/stderr files for diagnosis.""" + def fake_run(cmd, **kwargs): + raise subprocess.TimeoutExpired(cmd=cmd, timeout=kwargs["timeout"], output="partial stdout", stderr="slow") + + monkeypatch.setattr(debug_launcher.subprocess, "run", fake_run) + + metadata = await debug_launcher.run_claude_once( + urls={"anthropic_base_url": "http://127.0.0.1:9000/sessions/s-timeout"}, + output_dir=tmp_path, + session_id="s-timeout", + claude_prompt="Reply with OK only.", + claude_model="claude-sonnet-4-5", + anthropic_api_key="dummy-local-key", + claude_timeout=0.01, + ) + + assert metadata["returncode"] == 124 + assert metadata["timed_out"] is True + assert metadata["timeout_seconds"] == 0.01 + assert Path(metadata["stdout_path"]).read_text() == "partial stdout" + assert Path(metadata["stderr_path"]).read_text() == "slow" + + +@pytest.mark.asyncio +async def test_wait_for_manual_finalize_accepts_enter_from_tty_stdin(): + """Manual launcher mode can be finalized with Enter instead of requiring a + signal-only shutdown path.""" + read_fd, write_fd = os.pipe() + reader = os.fdopen(read_fd, "r", encoding="utf-8") + + class EnterStdin: + def isatty(self): + return True + + def fileno(self): + return reader.fileno() + + def readline(self): + return reader.readline() + + try: + task = asyncio.create_task(debug_launcher.wait_for_manual_finalize(stdin=EnterStdin())) + await asyncio.sleep(0) + os.write(write_fd, b"\n") + trigger = await asyncio.wait_for(task, timeout=1.0) + finally: + reader.close() + os.close(write_fd) + + assert trigger == "stdin" + + +@pytest.mark.asyncio +async def test_run_debug_session_once_fake_creates_finalizes_and_writes_trajectories( + monkeypatch, + tmp_path, +): + """The fake debug session starts the gateway, simulates Claude hitting the + Anthropic route, finalizes the session, and writes trajectory artifacts.""" + post_urls = [] + captured_run = {} + + original_post = httpx.AsyncClient.post + + async def recording_post(self, url, *args, **kwargs): + post_urls.append(url) + return await original_post(self, url, *args, **kwargs) + + def fake_run(cmd, *, env, check, text, capture_output, timeout): + captured_run["cmd"] = cmd + captured_run["env"] = env + captured_run["check"] = check + captured_run["timeout"] = timeout + probe_response = httpx.head(env["ANTHROPIC_BASE_URL"], timeout=5.0) + assert probe_response.status_code == 404 + url = f"{env['ANTHROPIC_BASE_URL']}/v1/messages" + post_urls.append(url) + response = httpx.post( + url, + params={"beta": "true"}, + headers={"x-api-key": env["ANTHROPIC_API_KEY"]}, + json={ + "model": "claude-sonnet-4-5", + "max_tokens": 8, + "messages": [{"role": "user", "content": "Reply with OK only."}], + }, + timeout=5.0, + ) + response.raise_for_status() + return subprocess.CompletedProcess(args=cmd, returncode=0, stdout="OK\n", stderr="debug log\n") + + monkeypatch.setattr(httpx.AsyncClient, "post", recording_post) + monkeypatch.setattr(debug_launcher.subprocess, "run", fake_run) + + result = await debug_launcher.run_debug_session_once( + session_id="fake-session", + output_dir=tmp_path, + backend="fake", + run_claude=True, + claude_prompt="Reply with OK only.", + claude_model="claude-sonnet-4-5", + anthropic_api_key="dummy-local-key", + ) + + assert result.session_id == "fake-session" + assert result.output_path == tmp_path / "fake-session" / "trajectories.jsonl" + assert result.metadata_path == tmp_path / "fake-session" / "session_metadata.json" + assert result.debug_snapshot_path == tmp_path / "fake-session" / "debug_snapshot.json" + assert result.trajectories_count == 1 + assert captured_run["cmd"][0] == "claude" + assert captured_run["check"] is False + assert captured_run["timeout"] == 300.0 + assert captured_run["env"]["ANTHROPIC_BASE_URL"].endswith("/sessions/fake-session") + assert captured_run["env"]["ANTHROPIC_API_KEY"] == "dummy-local-key" + assert any(url.endswith("/sessions/fake-session/v1/messages") for url in post_urls) + records = [json.loads(line) for line in result.output_path.read_text().splitlines()] + session_metadata = json.loads(result.metadata_path.read_text()) + assert session_metadata["metadata"]["claude"]["returncode"] == 0 + assert session_metadata["trajectories_count"] == 1 + assert session_metadata["debug_snapshot_path"].endswith("debug_snapshot.json") + assert records[0]["session_id"] == "fake-session" + assert records[0]["metadata"]["backend"] == "fake" + assert records[0]["metadata"]["claude"]["enabled"] is True + assert records[0]["metadata"]["claude"]["returncode"] == 0 + stdout_path = Path(records[0]["metadata"]["claude"]["stdout_path"]) + stderr_path = Path(records[0]["metadata"]["claude"]["stderr_path"]) + assert stdout_path.read_text() == "OK\n" + assert stderr_path.read_text() == "debug log\n" + assert records[0]["metadata"]["claude"]["debug_log_path"].endswith("claude-debug.log") + assert records[0]["trajectory"]["response_ids"] == [79, 75] + debug_snapshot = json.loads(result.debug_snapshot_path.read_text()) + assert debug_snapshot["session_state"]["has_active_trajectory"] is True + assert debug_snapshot["message_history"][0] == {"role": "user", "content": "Reply with OK only."} + + +@pytest.mark.asyncio +async def test_async_main_returns_nonzero_when_claude_fails(monkeypatch, tmp_path): + """The CLI exits with Claude's non-zero return code when one-shot Claude + execution fails after artifacts are written.""" + async def fake_run_debug_session_once(**kwargs): + return debug_launcher.DebugSessionResult( + session_id=kwargs["session_id"], + output_path=tmp_path / kwargs["session_id"] / "trajectories.jsonl", + metadata_path=tmp_path / kwargs["session_id"] / "session_metadata.json", + debug_snapshot_path=tmp_path / kwargs["session_id"] / "debug_snapshot.json", + trajectories_count=0, + provider_urls={}, + claude_returncode=1, + ) + + monkeypatch.setattr(debug_launcher, "run_debug_session_once", fake_run_debug_session_once) + + exit_code = await debug_launcher.async_main( + [ + "--session-id", + "failed-claude", + "--output-dir", + str(tmp_path), + "--run-claude", + ] + ) + + assert exit_code == 1 diff --git a/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py b/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py index 0eccce9c..9517af02 100644 --- a/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py +++ b/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py @@ -1,6 +1,7 @@ import asyncio import json import time +from types import SimpleNamespace import httpx import pytest @@ -20,6 +21,15 @@ fake_vision_info_extractor, ) +ALLOWED_SAMPLING_KEYS = frozenset({"temperature", "top_p", "top_k", "max_tokens", "stop"}) + + +def fake_tool_call_dispatch(text, tools, parser_name, tokenizer): + if "" not in text: + return text, [] + arguments = '{"query":"crop"}' if "crop" in text else '{"query":"weather"}' + return "", [SimpleNamespace(name="search", arguments=arguments)] + @pytest.fixture(scope="session") def ray_runtime(): @@ -43,6 +53,8 @@ async def test_gateway_actor_max_tokens_clamped_to_remaining_response_budget(): tokenizer=FakeTokenizer(), prompt_length=2048, response_length=100, + base_sampling_params={"top_p": 0.8}, + allowed_request_sampling_param_keys={"max_tokens"}, ), InspectingBackend(), ) @@ -57,9 +69,10 @@ async def test_gateway_actor_max_tokens_clamped_to_remaining_response_budget(): payload = {"messages": [{"role": "user", "content": "hi"}], "max_tokens": 200} actor._sessions["s1"].message_history = list(payload["messages"]) - await actor._handle_chat_completions("s1", payload) + await actor._handle_openai_chat_completions("s1", payload) assert actor._backend.calls[-1]["sampling_params"]["max_tokens"] == 40 + assert actor._backend.calls[-1]["sampling_params"]["top_p"] == 0.8 finally: await actor.shutdown() @@ -108,7 +121,7 @@ async def test_gateway_actor_continuation_budget_exhausted_materializes_length_s } backend.calls.clear() - response = await actor._handle_chat_completions("s1", payload) + response = await actor._handle_openai_chat_completions("s1", payload) body = json.loads(response.body) assert body["choices"][0]["finish_reason"] == "length" @@ -141,7 +154,7 @@ async def test_backend_value_error_raises_400(): await actor.create_session("s1") backend.next_error = ValueError("Prompt length (123456) exceeds the model's maximum context length (8192).") with pytest.raises(HTTPException) as exc_info: - await actor._handle_chat_completions("s1", {"messages": [{"role": "user", "content": "hi"}]}) + await actor._handle_openai_chat_completions("s1", {"messages": [{"role": "user", "content": "hi"}]}) assert exc_info.value.status_code == 400 assert "exceeds the model's maximum context length" in str(exc_info.value.detail) @@ -162,47 +175,79 @@ async def test_unknown_session_raises_404(): await actor.start() try: with pytest.raises(HTTPException) as exc_info: - await actor._handle_chat_completions("does-not-exist", {"messages": [{"role": "user", "content": "hi"}]}) + await actor._handle_openai_chat_completions( + "does-not-exist", {"messages": [{"role": "user", "content": "hi"}]} + ) assert exc_info.value.status_code == 404 finally: await actor.shutdown() -@pytest.mark.parametrize( - ("raw_arguments", "expected_arguments"), - [ - # Valid JSON string is parsed to a dict so Qwen-style chat templates that - # iterate with ``|items`` receive the expected type. - ('{"x": 1}', {"x": 1}), - # Invalid JSON is preserved as the raw string rather than raising or - # silently corrupting it. - ("not json", "not json"), - ], -) -def test_message_normalization_tool_call_arguments(raw_arguments, expected_arguments): - """``MessageCodec.normalize_request`` parses valid JSON tool-call arguments - into a dict and leaves invalid JSON as the original string.""" - from uni_agent.gateway.session import MessageCodec +def test_prefix_canonicalization_ignores_provider_ids_and_normalizes_arguments(): + """Prefix comparison ignores provider-generated tool ids and normalizes + JSON-equivalent tool-call arguments without hiding invalid raw strings.""" + from uni_agent.gateway.session.codec import MessageCodec - result = MessageCodec(FakeTokenizer()).normalize_request( - { - "messages": [ - { - "role": "assistant", - "tool_calls": [ - { - "id": "x", - "type": "function", - "function": {"name": "f", "arguments": raw_arguments}, - } - ], - } - ] - } - )["messages"][0] + codec = MessageCodec(FakeTokenizer()) + a = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_AAA", + "type": "function", + "function": {"name": "f", "arguments": {"x": 1}}, + } + ], + } + b = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_BBB", + "type": "function", + "function": {"name": "f", "arguments": {"x": 1}}, + } + ], + } + assert codec.canonicalize_message_for_prefix_comparison(a) == codec.canonicalize_message_for_prefix_comparison(b) + tc = codec.canonicalize_message_for_prefix_comparison(a)["tool_calls"][0] + assert "id" not in tc + + assert codec.canonicalize_message_for_prefix_comparison( + {"role": "tool", "tool_call_id": "call_AAA", "content": "found"} + ) == codec.canonicalize_message_for_prefix_comparison( + {"role": "tool", "tool_call_id": "call_BBB", "content": "found"} + ) + assert codec.canonicalize_message_for_prefix_comparison( + {"role": "assistant", "tool_call_id": "call_AAA", "content": ""} + ) == {"role": "assistant", "tool_call_id": "call_AAA", "content": ""} - assert result["tool_calls"][0]["function"]["arguments"] == expected_arguments + argument_cases = [ + ({"b": 2, "a": 1}, '{"a": 1, "b": 2}', True), + ("{b: 2, a: 1}", "{a: 1, b: 2}", False), + ] + for arguments_a, arguments_b, expect_equal in argument_cases: + msg_a = { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "call-1", "type": "function", "function": {"name": "search", "arguments": arguments_a}} + ], + } + msg_b = { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "call-2", "type": "function", "function": {"name": "search", "arguments": arguments_b}} + ], + } + assert ( + codec.canonicalize_message_for_prefix_comparison(msg_a) + == codec.canonicalize_message_for_prefix_comparison(msg_b) + ) is expect_equal @pytest.mark.asyncio @@ -233,7 +278,7 @@ def _spy(tokenizer, messages, **kwargs): await actor.start() try: await actor.create_session("s1") - await actor._handle_chat_completions( + await actor._handle_openai_chat_completions( "s1", { "messages": [{"role": "user", "content": "hi"}], @@ -248,19 +293,7 @@ def _spy(tokenizer, messages, **kwargs): @pytest.mark.asyncio -@pytest.mark.parametrize( - ("payload_extra", "expected_message_substr"), - [ - ({"n": 2}, "n=2 is not supported"), - ({"response_format": {"type": "json_object"}}, "response_format is not supported"), - ({"tool_choice": "required"}, 'tool_choice="required"'), - ( - {"tool_choice": {"type": "function", "function": {"name": "foo"}}}, - "tool_choice", - ), - ], -) -async def test_unsupported_capabilities_rejected_with_400(payload_extra, expected_message_substr): +async def test_unsupported_capabilities_rejected_with_400(): """OpenAI capabilities that the gateway does not support (``n > 1``, ``response_format``, ``tool_choice="required"``, and per-function ``tool_choice``) are rejected with HTTP 400 before reaching the session @@ -274,42 +307,61 @@ async def test_unsupported_capabilities_rejected_with_400(payload_extra, expecte await actor.start() try: await actor.create_session("s1") - with pytest.raises(HTTPException) as exc_info: - await actor._handle_chat_completions( - "s1", {"messages": [{"role": "user", "content": "hi"}], **payload_extra} - ) - - assert exc_info.value.status_code == 400 - assert expected_message_substr in str(exc_info.value.detail) + cases = [ + ({"n": 2}, "n=2 is not supported"), + ({"response_format": {"type": "json_object"}}, "response_format is not supported"), + ({"tool_choice": "required"}, 'tool_choice="required"'), + ({"tool_choice": {"type": "function", "function": {"name": "foo"}}}, "tool_choice"), + ] + for payload_extra, expected_message_substr in cases: + with pytest.raises(HTTPException) as exc_info: + await actor._handle_openai_chat_completions( + "s1", {"messages": [{"role": "user", "content": "hi"}], **payload_extra} + ) + assert exc_info.value.status_code == 400 + assert expected_message_substr in str(exc_info.value.detail) finally: await actor.shutdown() @pytest.mark.asyncio -async def test_stream_true_softly_falls_back_to_non_streaming(caplog): - """``stream=true`` is not supported; the gateway logs a warning and - returns a non-streaming response (soft fallback).""" +async def test_provider_stream_true_returns_sse(): + """OpenAI and Anthropic requests with stream=true return provider-shaped + StreamingResponse bodies.""" + from fastapi.responses import StreamingResponse + from uni_agent.gateway.config import GatewayActorConfig from uni_agent.gateway.gateway import _GatewayActor actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), InspectingBackend()) await actor.start() try: - await actor.create_session("s1") - with caplog.at_level("WARNING", logger="gateway"): - response = await actor._handle_chat_completions( - "s1", {"messages": [{"role": "user", "content": "hi"}], "stream": True} - ) - - assert response.status_code == 200 - assert any("stream=true" in record.getMessage() for record in caplog.records) + await actor.create_session("s-oa-stream") + openai_resp = await actor._handle_openai_chat_completions( + "s-oa-stream", {"messages": [{"role": "user", "content": "hi"}], "stream": True} + ) + assert isinstance(openai_resp, StreamingResponse) + openai_text = (b"".join([chunk async for chunk in openai_resp.body_iterator])).decode() + assert "chat.completion.chunk" in openai_text + assert "data: [DONE]" in openai_text + + await actor.create_session("s-an-stream") + anthropic_resp = await actor._handle_anthropic_messages( + "s-an-stream", + {"max_tokens": 8, "stream": True, "messages": [{"role": "user", "content": "hi"}]}, + ) + assert isinstance(anthropic_resp, StreamingResponse) + anthropic_text = (b"".join([c async for c in anthropic_resp.body_iterator])).decode() + assert "event: message_start" in anthropic_text + assert "event: message_stop" in anthropic_text finally: await actor.shutdown() @pytest.mark.asyncio -async def test_chat_completion_response_includes_created_and_model_fields(): - """Response body carries OpenAI-standard ``created`` and ``model`` fields.""" +async def test_openai_chat_completion_response_uses_request_model_or_unknown(): + """The route passes the request model into the OpenAI response envelope and + falls back to ``unknown`` when the request omits it.""" from uni_agent.gateway.config import GatewayActorConfig from uni_agent.gateway.gateway import _GatewayActor @@ -318,23 +370,18 @@ async def test_chat_completion_response_includes_created_and_model_fields(): try: await actor.create_session("s-model") before = int(time.time()) - response = await actor._handle_chat_completions( + response = await actor._handle_openai_chat_completions( "s-model", {"model": "dummy-model", "messages": [{"role": "user", "content": "hi"}]}, ) after = int(time.time()) body = json.loads(response.body) - assert body["id"].startswith("chatcmpl-") - assert body["object"] == "chat.completion" assert before <= body["created"] <= after assert body["model"] == "dummy-model" - assert body["choices"][0]["index"] == 0 - assert body["choices"][0]["finish_reason"] == "stop" - assert body["usage"]["total_tokens"] == body["usage"]["prompt_tokens"] + body["usage"]["completion_tokens"] await actor.create_session("s-fallback") - fallback_response = await actor._handle_chat_completions( + fallback_response = await actor._handle_openai_chat_completions( "s-fallback", {"messages": [{"role": "user", "content": "hi"}]}, ) @@ -373,7 +420,7 @@ def _spy(tokenizer, messages, **kwargs): await actor.start() try: await actor.create_session("s1") - response = await actor._handle_chat_completions( + response = await actor._handle_openai_chat_completions( "s1", { "messages": [{"role": "user", "content": "hi"}], @@ -395,9 +442,9 @@ async def test_gateway_actor_forwards_image_data_on_initial_multimodal_request(r """On the first turn of a multimodal session, ``image_data`` extracted from the request is forwarded to the backend and recorded in the resulting ``Trajectory.multi_modal_data``.""" + from uni_agent.gateway.adapters.openai import openai_to_internal from uni_agent.gateway.config import GatewayActorConfig from uni_agent.gateway.gateway import GatewayActor - from uni_agent.gateway.session import MessageCodec processor = FakeProcessor() actor = GatewayActor.remote( @@ -425,7 +472,11 @@ async def test_gateway_actor_forwards_image_data_on_initial_multimodal_request(r ], } - normalized = MessageCodec(FakeTokenizer()).normalize_request(payload) + normalized = openai_to_internal( + payload, + base_sampling_params={}, + allowed_sampling_keys=ALLOWED_SAMPLING_KEYS, + ) raw_prompt = processor.apply_chat_template( normalized["messages"], tokenize=False, @@ -617,18 +668,20 @@ async def test_gateway_actor_multimodal_reference_change_splits_trajectory(ray_r @pytest.mark.asyncio -async def test_gateway_actor_continuation_with_tool_returned_image_appends_media(ray_runtime): +async def test_gateway_actor_continuation_with_tool_returned_image_appends_media(monkeypatch): """When a tool-call continuation brings a new image (e.g. a zoomed crop), the new image is appended to the session media accumulator. The full ``prompt_ids`` sequence (initial prompt + tool-call tokens + incremental prompt) is verified token-by-token.""" from uni_agent.gateway.config import GatewayActorConfig - from uni_agent.gateway.gateway import GatewayActor + from uni_agent.gateway.gateway import _GatewayActor + import uni_agent.gateway.session.codec as codec_mod from verl.utils.chat_template import apply_chat_template, initialize_system_prompt + monkeypatch.setattr(codec_mod, "_extract_tool_calls_with_sglang_or_vllm", fake_tool_call_dispatch) processor = FakeProcessor() tool_call_text = '\n{"name": "search", "arguments": {"query": "crop"}}\n' - actor = GatewayActor.remote( + actor = _GatewayActor( GatewayActorConfig( tokenizer=FakeTokenizer(), processor=processor, @@ -637,8 +690,8 @@ async def test_gateway_actor_continuation_with_tool_returned_image_appends_media ), InspectingSequencedBackend([tool_call_text, "__inspect__"]), ) - ray.get(actor.start.remote()) - session = ray.get(actor.create_session.remote("session-mm-tool-image")) + await actor.start() + await actor.create_session("session-mm-tool-image") tools = [{"type": "function", "function": {"name": "search", "parameters": {"type": "object"}}}] initial_message = { @@ -649,40 +702,39 @@ async def test_gateway_actor_continuation_with_tool_returned_image_appends_media ], } - async with httpx.AsyncClient(timeout=5.0) as client: - first = await client.post( - f"{session.base_url}/chat/completions", - json={ - "model": "dummy-model", - "tools": tools, - "messages": [initial_message], - }, - ) - assert first.status_code == 200 - assistant_message = first.json()["choices"][0]["message"] - tool_message = { - "role": "tool", - "tool_call_id": assistant_message["tool_calls"][0]["id"], - "content": [ - {"type": "image_url", "image_url": {"url": "image://tool-b.png"}}, - {"type": "text", "text": "zoomed crop"}, - ], - } + first = await actor._handle_openai_chat_completions( + "session-mm-tool-image", + { + "model": "dummy-model", + "tools": tools, + "messages": [initial_message], + }, + ) + assert first.status_code == 200 + assistant_message = json.loads(first.body)["choices"][0]["message"] + tool_message = { + "role": "tool", + "tool_call_id": assistant_message["tool_calls"][0]["id"], + "content": [ + {"type": "image_url", "image_url": {"url": "image://tool-b.png"}}, + {"type": "text", "text": "zoomed crop"}, + ], + } - second = await client.post( - f"{session.base_url}/chat/completions", - json={ - "model": "dummy-model", - "tools": tools, - "messages": [initial_message, assistant_message, tool_message], - }, - ) + second = await actor._handle_openai_chat_completions( + "session-mm-tool-image", + { + "model": "dummy-model", + "tools": tools, + "messages": [initial_message, assistant_message, tool_message], + }, + ) - trajectories = ray.get(actor.finalize_session.remote("session-mm-tool-image")) - ray.get(actor.shutdown.remote()) + trajectories = await actor.finalize_session("session-mm-tool-image") + await actor.shutdown() assert second.status_code == 200 - second_call = json.loads(second.json()["choices"][0]["message"]["content"]) + second_call = json.loads(json.loads(second.body)["choices"][0]["message"]["content"]) assert second_call["image_data"] == ["image://a.png", "image://tool-b.png"] assert len(trajectories) == 1 assert trajectories[0].multi_modal_data == { @@ -898,41 +950,6 @@ async def test_gateway_actor_continuation_preserves_prompt_and_generation_masks( assert trajectories[0].response_mask[-len("SECOND") :] == [1] * len("SECOND") -@pytest.mark.parametrize( - ("arguments_a", "arguments_b", "expect_equal"), - [ - # Valid JSON: a dict and an equivalent JSON string (same keys, different - # order) canonicalize equal, so tool-argument JSON round-trip drift - # between turns does not spuriously split the trajectory. - ({"b": 2, "a": 1}, '{"a": 1, "b": 2}', True), - # Invalid JSON (unquoted keys): comparison falls back to raw string - # comparison, which is order-sensitive — the same keys in a different - # order stay not-equal (contrast with the JSON path above). - ("{b: 2, a: 1}", "{a: 1, b: 2}", False), - ], -) -def test_canonicalize_tool_call_arguments_for_prefix_comparison(arguments_a, arguments_b, expect_equal): - """``MessageCodec.canonicalize_message_for_prefix_comparison`` normalizes - tool-call arguments so that JSON-equivalent values match, and falls back to - raw string comparison when the arguments are not valid JSON.""" - from uni_agent.gateway.session import MessageCodec - - def _message(arguments): - return { - "role": "assistant", - "content": "", - "tool_calls": [ - {"id": "call-1", "type": "function", "function": {"name": "search", "arguments": arguments}} - ], - } - - codec = MessageCodec(FakeTokenizer()) - canonical_a = codec.canonicalize_message_for_prefix_comparison(_message(arguments_a)) - canonical_b = codec.canonicalize_message_for_prefix_comparison(_message(arguments_b)) - - assert (canonical_a == canonical_b) is expect_equal - - @pytest.mark.asyncio async def test_gateway_actor_serializes_same_session_concurrent_requests(ray_runtime): """Two concurrent requests to the same session are serialized by @@ -973,9 +990,19 @@ async def send_request(): assert trajectories[1].response_mask == [1] * len("SECOND") -@pytest.mark.parametrize( - ("payload", "detail_fragment"), - [ +@pytest.mark.asyncio +async def test_gateway_actor_rejects_malformed_requests_with_bad_request(ray_runtime): + """Malformed request payloads (empty messages, bad types, invalid + tool_calls/tools structure) are rejected with HTTP 400 and an + OpenAI-style error envelope.""" + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import GatewayActor + + actor = GatewayActor.remote(GatewayActorConfig(tokenizer=FakeTokenizer()), QueuedBackend(["DONE"])) + ray.get(actor.start.remote()) + session = ray.get(actor.create_session.remote("session-validation")) + + cases = [ ({"model": "dummy-model", "messages": []}, "messages must be non-empty"), ( {"model": "dummy-model", "messages": [{"role": "user", "name": 123, "content": "hello"}]}, @@ -1000,30 +1027,25 @@ async def send_request(): }, "tools must be a list", ), - ], -) -@pytest.mark.asyncio -async def test_gateway_actor_rejects_malformed_requests_with_bad_request(ray_runtime, payload, detail_fragment): - """Malformed request payloads (empty messages, bad types, invalid - tool_calls/tools structure) are rejected with HTTP 400 and an - OpenAI-style error envelope.""" - from uni_agent.gateway.config import GatewayActorConfig - from uni_agent.gateway.gateway import GatewayActor - - actor = GatewayActor.remote(GatewayActorConfig(tokenizer=FakeTokenizer()), QueuedBackend(["DONE"])) - ray.get(actor.start.remote()) - session = ray.get(actor.create_session.remote("session-validation")) - + ] async with httpx.AsyncClient(timeout=5.0) as client: - response = await client.post( - f"{session.base_url}/chat/completions", - json=payload, - ) + responses = [] + for payload, detail_fragment in cases: + response = await client.post( + f"{session.base_url}/chat/completions", + json=payload, + ) + responses.append((response, detail_fragment)) ray.get(actor.shutdown.remote()) - assert response.status_code == 400 - assert detail_fragment in response.text + for response, detail_fragment in responses: + assert response.status_code == 400 + body = response.json() + assert body["error"]["type"] == "invalid_request_error" + assert body["error"]["code"] is None + assert body["error"]["param"] is None + assert detail_fragment in body["error"]["message"] @pytest.mark.asyncio @@ -1049,6 +1071,7 @@ async def test_gateway_actor_backend_failure_does_not_commit_partial_state(ray_r ray.get(actor.shutdown.remote()) assert response.status_code == 500 + assert response.json()["error"]["type"] == "internal_server_error" assert state["num_trajectories"] == 0 assert state["has_active_trajectory"] is False @@ -1105,35 +1128,36 @@ async def test_gateway_actor_backend_failure_after_tool_mismatch_does_not_split( @pytest.mark.asyncio -async def test_gateway_actor_tool_call_decode_returns_openai_format(ray_runtime): +async def test_gateway_actor_tool_call_decode_returns_openai_format(monkeypatch): """When tool_parser_name is set and model outputs tool call tokens, the HTTP response should contain tool_calls in OpenAI format.""" from uni_agent.gateway.config import GatewayActorConfig - from uni_agent.gateway.gateway import GatewayActor + from uni_agent.gateway.gateway import _GatewayActor + import uni_agent.gateway.session.codec as codec_mod + monkeypatch.setattr(codec_mod, "_extract_tool_calls_with_sglang_or_vllm", fake_tool_call_dispatch) tool_call_text = '\n{"name": "search", "arguments": {"query": "weather"}}\n' - actor = GatewayActor.remote( + actor = _GatewayActor( GatewayActorConfig( tokenizer=FakeTokenizer(), tool_parser_name="hermes", ), QueuedBackend([tool_call_text, "sunny today"]), ) - ray.get(actor.start.remote()) - session = ray.get(actor.create_session.remote("session-tool-call")) - - async with httpx.AsyncClient(timeout=5.0) as client: + await actor.start() + await actor.create_session("session-tool-call") + try: # First request: model returns a tool call - first = await client.post( - f"{session.base_url}/chat/completions", - json={ + first = await actor._handle_openai_chat_completions( + "session-tool-call", + { "model": "dummy-model", "tools": [{"type": "function", "function": {"name": "search", "parameters": {"type": "object"}}}], "messages": [{"role": "user", "content": "what is the weather?"}], }, ) assert first.status_code == 200 - first_data = first.json() + first_data = json.loads(first.body) assert first_data["choices"][0]["finish_reason"] == "tool_calls" tool_calls = first_data["choices"][0]["message"].get("tool_calls") assert tool_calls is not None @@ -1145,9 +1169,9 @@ async def test_gateway_actor_tool_call_decode_returns_openai_format(ray_runtime) assert isinstance(tool_calls[0]["function"]["arguments"], str) # Second request: agent sends back tool result as continuation - second = await client.post( - f"{session.base_url}/chat/completions", - json={ + second = await actor._handle_openai_chat_completions( + "session-tool-call", + { "model": "dummy-model", "tools": [{"type": "function", "function": {"name": "search", "parameters": {"type": "object"}}}], "messages": [ @@ -1158,12 +1182,231 @@ async def test_gateway_actor_tool_call_decode_returns_openai_format(ray_runtime) }, ) assert second.status_code == 200 - assert second.json()["choices"][0]["message"]["content"] == "sunny today" + assert json.loads(second.body)["choices"][0]["message"]["content"] == "sunny today" - trajectories = ray.get(actor.finalize_session.remote("session-tool-call")) - ray.get(actor.shutdown.remote()) + trajectories = await actor.finalize_session("session-tool-call") + finally: + await actor.shutdown() assert len(trajectories) == 1 # Should have both mask=0 (incremental) and mask=1 (model output) tokens assert 0 in trajectories[0].response_mask assert 1 in trajectories[0].response_mask + + +@pytest.mark.asyncio +async def test_anthropic_messages_end_to_end(): + """The Anthropic Messages route lowers a simple request, runs one + generation, and serializes an Anthropic message response.""" + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), InspectingBackend()) + await actor.start() + try: + await actor.create_session("s-anth") + resp = await actor._handle_anthropic_messages( + "s-anth", + {"model": "claude-x", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}, + ) + body = json.loads(resp.body) + assert body["type"] == "message" + assert body["content"][0]["type"] == "text" + assert "input_tokens" in body["usage"] + finally: + await actor.shutdown() + + +@pytest.mark.asyncio +async def test_anthropic_and_openai_produce_identical_trajectory(): + """An Anthropic request and its equivalent OpenAI request yield identical + token-truth (prompt_ids/response_ids/response_mask/response_logprobs).""" + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), QueuedBackend(["same", "same"])) + await actor.start() + try: + await actor.create_session("s-oa") + await actor._handle_openai_chat_completions( + "s-oa", + { + "messages": [{"role": "system", "content": "Be brief."}, {"role": "user", "content": "hi"}], + "max_tokens": 16, + }, + ) + oa = (await actor.finalize_session("s-oa"))[0] + + await actor.create_session("s-an") + await actor._handle_anthropic_messages( + "s-an", {"system": "Be brief.", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]} + ) + an = (await actor.finalize_session("s-an"))[0] + + assert oa.prompt_ids == an.prompt_ids + assert oa.response_ids == an.response_ids + assert oa.response_mask == an.response_mask + assert oa.response_logprobs == an.response_logprobs + finally: + await actor.shutdown() + + +@pytest.mark.asyncio +async def test_anthropic_tool_turn_round_trip_extends_not_reencodes(monkeypatch): + """When an Anthropic agent echoes a previous assistant tool_use turn back as + history, conversion must reproduce stored normalized message sufficiently for + prefix check to pass and extend one trajectory instead of re-encoding a new one.""" + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + import uni_agent.gateway.session.codec as codec_mod + + monkeypatch.setattr(codec_mod, "_extract_tool_calls_with_sglang_or_vllm", fake_tool_call_dispatch) + tool_call_text = '\n{"name": "search", "arguments": {"query": "weather"}}\n' + actor = _GatewayActor( + GatewayActorConfig( + tokenizer=FakeTokenizer(), + tool_parser_name="hermes", + ), + QueuedBackend([tool_call_text, "sunny today"]), + ) + await actor.start() + try: + await actor.create_session("s-rt") + first = await actor._handle_anthropic_messages( + "s-rt", + { + "max_tokens": 16, + "tools": [{"name": "search", "input_schema": {"type": "object"}}], + "messages": [{"role": "user", "content": "go"}], + }, + ) + first_body = json.loads(first.body) + assert first_body["content"][0]["type"] == "tool_use" + echoed = [ + {"role": "user", "content": "go"}, + {"role": "assistant", "content": first_body["content"]}, + {"role": "user", "content": "next"}, + ] + await actor._handle_anthropic_messages( + "s-rt", + { + "max_tokens": 16, + "tools": [{"name": "search", "input_schema": {"type": "object"}}], + "messages": echoed, + }, + ) + trajectories = await actor.finalize_session("s-rt") + assert len(trajectories) == 1 + finally: + await actor.shutdown() + + +@pytest.mark.asyncio +async def test_anthropic_error_envelope_shape(): + """Malformed Anthropic requests return Anthropic-shaped error bodies rather + than OpenAI or FastAPI default error envelopes.""" + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), InspectingBackend()) + await actor.start() + try: + await actor.create_session("s-err") + transport = httpx.ASGITransport(app=actor._app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as client: + r = await client.post( + "/sessions/s-err/v1/messages", + json={"max_tokens": 8, "messages": [{"role": "user", "content": "hi"}], "tool_choice": {"type": "any"}}, + ) + assert r.status_code == 400 + body = r.json() + assert body["type"] == "error" + assert body["error"]["type"] == "invalid_request_error" + assert "message" in body["error"] + finally: + await actor.shutdown() + + +@pytest.mark.asyncio +async def test_provider_malformed_json_uses_error_envelope(): + """Invalid JSON is reported through the matching provider error envelope.""" + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), InspectingBackend()) + await actor.start() + try: + await actor.create_session("s-json-anth") + await actor.create_session("s-json-openai") + transport = httpx.ASGITransport(app=actor._app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as client: + anthropic = await client.post( + "/sessions/s-json-anth/v1/messages", + content="{bad", + headers={"content-type": "application/json"}, + ) + openai = await client.post( + "/sessions/s-json-openai/v1/chat/completions", + content="{bad", + headers={"content-type": "application/json"}, + ) + assert anthropic.status_code == 400 + anthropic_body = anthropic.json() + assert anthropic_body["type"] == "error" + assert anthropic_body["error"]["type"] == "invalid_request_error" + assert "message" in anthropic_body["error"] + + assert openai.status_code == 400 + openai_body = openai.json() + assert openai_body["error"]["type"] == "invalid_request_error" + assert "message" in openai_body["error"] + finally: + await actor.shutdown() + + +@pytest.mark.asyncio +async def test_unhandled_exception_uses_provider_error_envelope(monkeypatch): + """Unhandled route exceptions still produce provider-shaped error bodies, + not FastAPI's default {"detail": ...} envelope.""" + from uni_agent.gateway.config import GatewayActorConfig + from uni_agent.gateway.gateway import _GatewayActor + + def _explode(*args, **kwargs): + raise KeyError("simulated programmer bug") + + cases = [ + ( + "uni_agent.gateway.gateway.anthropic_to_internal", + "s-unhandled-anth", + "/sessions/s-unhandled-anth/v1/messages", + {"max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]}, + "api_error", + True, + ), + ( + "uni_agent.gateway.gateway.openai_to_internal", + "s-unhandled-oa", + "/sessions/s-unhandled-oa/v1/chat/completions", + {"messages": [{"role": "user", "content": "hi"}]}, + "internal_server_error", + False, + ), + ] + for patch_target, session_id, path, payload, expected_error_type, has_top_level_type in cases: + with monkeypatch.context() as m: + m.setattr(patch_target, _explode) + actor = _GatewayActor(GatewayActorConfig(tokenizer=FakeTokenizer()), InspectingBackend()) + await actor.start() + try: + await actor.create_session(session_id) + transport = httpx.ASGITransport(app=actor._app, raise_app_exceptions=False) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as client: + r = await client.post(path, json=payload) + assert r.status_code == 500 + body = r.json() + if has_top_level_type: + assert body["type"] == "error" + assert body["error"]["type"] == expected_error_type + assert body["error"]["message"] == "Internal server error" + finally: + await actor.shutdown() diff --git a/tests/uni_agent/gateway/test_message_codec_tool_dispatch.py b/tests/uni_agent/gateway/test_message_codec_tool_dispatch.py new file mode 100644 index 00000000..dd7f855f --- /dev/null +++ b/tests/uni_agent/gateway/test_message_codec_tool_dispatch.py @@ -0,0 +1,122 @@ +from types import SimpleNamespace + +import pytest + +from tests.uni_agent.support import FakeTokenizer + + +TOOLS = [ + { + "type": "function", + "function": { + "name": "search", + "description": "search docs", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } +] + + +def test_tool_call_dispatch_prefers_sglang(monkeypatch): + import uni_agent.gateway.session.codec as codec_mod + + seen = {} + + def fake_sglang(text, tools, parser_name): + seen["sglang"] = (text, tools, parser_name) + return "visible", [SimpleNamespace(name="search", arguments='{"query":"x"}')] + + def fail_vllm(*args, **kwargs): + raise AssertionError("vLLM should not run when SGLang succeeds") + + monkeypatch.setattr(codec_mod, "_process_tool_calls_sglang", fake_sglang, raising=False) + monkeypatch.setattr(codec_mod, "_process_tool_calls_vllm", fail_vllm, raising=False) + + content, calls = codec_mod._extract_tool_calls_with_sglang_or_vllm("raw", TOOLS, "hermes", FakeTokenizer()) + + assert content == "visible" + assert calls[0].name == "search" + assert seen["sglang"] == ("raw", TOOLS, "hermes") + + +def test_tool_call_dispatch_falls_back_to_vllm_with_name_mapping(monkeypatch): + import uni_agent.gateway.session.codec as codec_mod + + seen = {} + + def missing_sglang(*args, **kwargs): + raise ModuleNotFoundError("sglang") + + def fake_vllm(text, tools, parser_name, tokenizer): + seen["vllm"] = (text, tools, parser_name, tokenizer) + return "", [SimpleNamespace(name="search", arguments='{"query":"x"}')] + + monkeypatch.setattr(codec_mod, "_process_tool_calls_sglang", missing_sglang, raising=False) + monkeypatch.setattr(codec_mod, "_process_tool_calls_vllm", fake_vllm, raising=False) + + tokenizer = FakeTokenizer() + content, calls = codec_mod._extract_tool_calls_with_sglang_or_vllm("raw", TOOLS, "qwen25", tokenizer) + + assert content == "" + assert calls[0].arguments == '{"query":"x"}' + assert seen["vllm"] == ("raw", TOOLS, "qwen3_xml", tokenizer) + + +def test_tool_call_dispatch_returns_text_when_backends_unavailable(monkeypatch): + import uni_agent.gateway.session.codec as codec_mod + + def missing_backend(*args, **kwargs): + raise ModuleNotFoundError("tool parser backend") + + monkeypatch.setattr(codec_mod, "_process_tool_calls_sglang", missing_backend, raising=False) + monkeypatch.setattr(codec_mod, "_process_tool_calls_vllm", missing_backend, raising=False) + + content, calls = codec_mod._extract_tool_calls_with_sglang_or_vllm("plain text", TOOLS, "hermes", FakeTokenizer()) + + assert content == "plain text" + assert calls == [] + + +@pytest.mark.asyncio +async def test_decode_response_uses_gateway_dispatcher_for_tool_calls(monkeypatch): + import uni_agent.gateway.session.codec as codec_mod + from uni_agent.gateway.session.codec import MessageCodec + + seen = {} + + def fake_dispatch(text, tools, parser_name, tokenizer): + seen["dispatch"] = (text, tools, parser_name, tokenizer) + return "", [SimpleNamespace(name="search", arguments='{"query":"weather"}')] + + monkeypatch.setattr(codec_mod, "_extract_tool_calls_with_sglang_or_vllm", fake_dispatch, raising=False) + + tokenizer = FakeTokenizer() + codec = MessageCodec(tokenizer, tool_parser_name="qwen3_xml") + message, finish_reason = await codec.decode_response( + [ord(char) for char in "ignored"], + tools=[ + { + "type": "function", + "function": { + "name": "search", + "description": "search docs", + "parameters": { + "type": "object", + "properties": { + "target": {"anyOf": [{"const": "file"}, {"type": "string"}]}, + }, + }, + }, + } + ], + stop_reason="stop", + ) + + assert finish_reason == "tool_calls" + assert message["content"] == "" + assert message["tool_calls"][0]["type"] == "function" + assert message["tool_calls"][0]["function"] == {"name": "search", "arguments": '{"query":"weather"}'} + assert seen["dispatch"][2] == "qwen3_xml" diff --git a/uni_agent/gateway/adapters/__init__.py b/uni_agent/gateway/adapters/__init__.py new file mode 100644 index 00000000..d0973d71 --- /dev/null +++ b/uni_agent/gateway/adapters/__init__.py @@ -0,0 +1,20 @@ +from .anthropic import anthropic_build_response, anthropic_error_body, anthropic_stream_response, anthropic_to_internal +from .openai import ( + openai_build_response, + openai_error_body, + openai_stream_response, + openai_to_internal, +) +from .types import MalformedRequestError + +__all__ = [ + "anthropic_build_response", + "anthropic_error_body", + "anthropic_stream_response", + "anthropic_to_internal", + "MalformedRequestError", + "openai_build_response", + "openai_error_body", + "openai_stream_response", + "openai_to_internal", +] diff --git a/uni_agent/gateway/adapters/anthropic.py b/uni_agent/gateway/adapters/anthropic.py new file mode 100644 index 00000000..ffd51b25 --- /dev/null +++ b/uni_agent/gateway/adapters/anthropic.py @@ -0,0 +1,553 @@ +"""Anthropic Messages wire -> InternalGenerationRequest.""" + +from __future__ import annotations + +import copy +import json +import logging +import re +import secrets +from collections.abc import AsyncIterator +from typing import Any + +from fastapi.responses import StreamingResponse + +from uni_agent.gateway.session.session import GenerationOutcome +from uni_agent.gateway.session.types import InternalGenerationRequest + +from .types import MalformedRequestError + +_BILLING_HEADER_RE = re.compile(r"^\s*x-anthropic-billing-header:[^\n]*\n?", re.IGNORECASE | re.MULTILINE) +# content_filter/function_call are defensive: current internal finish_reason +# values are stop/length/tool_calls. +_STOP_REASON_MAP = { + "stop": "end_turn", + "length": "max_tokens", + "tool_calls": "tool_use", + "function_call": "tool_use", + "content_filter": "refusal", +} + +logger = logging.getLogger("gateway") + +_SSE_HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", +} + + +# Only status codes the gateway actually emits today (gateway.py / session.py: +# 400 malformed/JSON, 409 concurrent session generation, 500 internal). Other +# Anthropic types (authentication_error / permission_error / not_found_error / +# request_too_large / rate_limit_error / overloaded_error) are intentionally +# omitted: this gateway is not an Anthropic proxy and has no auth / rate-limit +# / capacity paths that would surface those codes. Anything not listed falls +# back to "api_error" via .get() default. +ANTHROPIC_ERROR_TYPE_BY_STATUS = { + 400: "invalid_request_error", + 409: "invalid_request_error", + 500: "api_error", +} + + +def anthropic_error_body(status_code: int, message: str) -> dict[str, Any]: + return { + "type": "error", + "error": { + "type": ANTHROPIC_ERROR_TYPE_BY_STATUS.get(status_code, "api_error"), + "message": message, + }, + } + + +def _tool_call_input(arguments: Any) -> dict[str, Any]: + if isinstance(arguments, dict): + return arguments + if isinstance(arguments, str): + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + logger.warning("Invalid Anthropic tool call JSON arguments; using empty input") + return {} + if isinstance(parsed, dict): + return parsed + logger.warning("Anthropic tool call JSON arguments were not an object; using empty input") + return {} + logger.warning("Anthropic tool call arguments were not an object or JSON string; using empty input") + return {} + + +def _outcome_to_blocks(assistant_msg: dict[str, Any]) -> list[dict[str, Any]]: + blocks: list[dict[str, Any]] = [] + # Outbound thinking/reasoning is intentionally not synthesized here: + # `assistant_msg` from session/decode currently has no `reasoning_content` + # field, and adding outbound wrapping before there is a producer would be + # defensive code at a trusted boundary (AGENTS §4 rule 3). When session + # decode grows reasoning_content, mirror OpenAI streaming and emit Anthropic + # thinking blocks at that point. + content = assistant_msg.get("content") + if isinstance(content, str) and content: + blocks.append({"type": "text", "text": content}) + + tool_calls = assistant_msg.get("tool_calls", []) + if isinstance(tool_calls, list): + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function", {}) + if not isinstance(function, dict): + function = {} + blocks.append( + { + "type": "tool_use", + "id": tool_call.get("id") or f"toolu_{secrets.token_hex(8)}", + "name": function.get("name", "tool"), + "input": _tool_call_input(function.get("arguments")), + } + ) + + if not blocks: + blocks.append({"type": "text", "text": ""}) + return blocks + + +def anthropic_build_response(outcome: GenerationOutcome, *, model: str) -> dict[str, Any]: + return { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "model": model, + "content": _outcome_to_blocks(outcome.assistant_msg), + "stop_reason": _STOP_REASON_MAP.get(outcome.finish_reason, "end_turn"), + "stop_sequence": None, + "usage": {"input_tokens": outcome.prompt_tokens, "output_tokens": outcome.completion_tokens}, + } + + +def anthropic_stream_response(outcome: GenerationOutcome, *, model: str) -> StreamingResponse: + """Synthesize an Anthropic Messages SSE stream from a completed outcome. + + Whole-turn synthesis (backend is not token-streaming): message_start, + per-block start/delta/stop, message_delta with stop_reason+usage, + message_stop. + """ + blocks = _outcome_to_blocks(outcome.assistant_msg) + stop_reason = _STOP_REASON_MAP.get(outcome.finish_reason, "end_turn") + msg_id = f"msg_{secrets.token_hex(12)}" + + def _event(name: str, data: dict[str, Any]) -> bytes: + return f"event: {name}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n".encode() + + async def _gen() -> AsyncIterator[bytes]: + yield _event( + "message_start", + { + "type": "message_start", + "message": { + "id": msg_id, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": outcome.prompt_tokens, "output_tokens": 0}, + }, + }, + ) + for idx, block in enumerate(blocks): + if block["type"] == "text": + start = {"type": "text", "text": ""} + delta = {"type": "text_delta", "text": block["text"]} + else: + start = {"type": "tool_use", "id": block["id"], "name": block["name"], "input": {}} + delta = {"type": "input_json_delta", "partial_json": json.dumps(block["input"], ensure_ascii=False)} + yield _event("content_block_start", {"type": "content_block_start", "index": idx, "content_block": start}) + yield _event("content_block_delta", {"type": "content_block_delta", "index": idx, "delta": delta}) + yield _event("content_block_stop", {"type": "content_block_stop", "index": idx}) + yield _event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": {"input_tokens": outcome.prompt_tokens, "output_tokens": outcome.completion_tokens}, + }, + ) + yield _event("message_stop", {"type": "message_stop"}) + + return StreamingResponse(_gen(), media_type="text/event-stream", headers=_SSE_HEADERS) + + +def _strip_billing_header(text: str) -> str: + # Claude Code can send this request-scoped billing header in system text; if + # kept, it changes prompt prefixes and causes cache/template prefix drift. + return _BILLING_HEADER_RE.sub("", text).strip() + + +def _system_to_text(system: Any) -> str: + if system is None: + return "" + if isinstance(system, str): + return _strip_billing_header(system) + if not isinstance(system, list): + raise MalformedRequestError("system must be a string or text block list") + + parts: list[str] = [] + for block in system: + if not isinstance(block, dict): + raise MalformedRequestError("system blocks must be objects") + if block.get("type") == "text": + text = block.get("text", "") + if not isinstance(text, str): + raise MalformedRequestError("system text must be a string") + parts.append(text) + else: + raise MalformedRequestError(f"Unsupported system block type: {block.get('type')}") + return _strip_billing_header("\n".join(part for part in parts if part)) + + +def _image_block_to_openai_part(block: dict[str, Any]) -> dict[str, Any]: + source = block.get("source") + if not isinstance(source, dict): + raise MalformedRequestError("image.source must be an object") + source_type = source.get("type") + if source_type == "base64": + media_type = source.get("media_type") + data = source.get("data") + if not isinstance(media_type, str) or not isinstance(data, str): + raise MalformedRequestError("base64 image source requires media_type and data") + url = f"data:{media_type};base64,{data}" + elif source_type == "url": + url = source.get("url") + if not isinstance(url, str): + raise MalformedRequestError("url image source requires url") + else: + raise MalformedRequestError(f"Unsupported image source type: {source_type}") + return {"type": "image_url", "image_url": {"url": url}} + + +def _tool_result_messages(tool_call_id: str, content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"role": "tool", "tool_call_id": tool_call_id, "content": content}] + if content is None: + return [{"role": "tool", "tool_call_id": tool_call_id, "content": ""}] + if not isinstance(content, list): + raise MalformedRequestError("tool_result.content must be a string or block list") + + messages: list[dict[str, Any]] = [] + texts: list[str] = [] + emitted_tool_message = False + + def flush_text() -> None: + nonlocal emitted_tool_message + if not texts: + return + messages.append({"role": "tool", "tool_call_id": tool_call_id, "content": "\n".join(texts)}) + emitted_tool_message = True + texts.clear() + + for block in content: + if not isinstance(block, dict): + raise MalformedRequestError("tool_result content blocks must be objects") + block_type = block.get("type") + if block_type == "text": + text = block.get("text", "") + if not isinstance(text, str): + raise MalformedRequestError("tool_result text must be a string") + texts.append(text) + elif block_type == "image": + # OpenAI/vLLM tool-role messages cannot carry images, so preserve + # visual payloads as a following user message downgrade. The + # assistant's tool_calls must be closed by at least one tool + # message before any user follows; emit an empty tool sentinel if + # neither buffered text nor a prior tool message has done that. + if not texts and not emitted_tool_message: + messages.append({"role": "tool", "tool_call_id": tool_call_id, "content": ""}) + emitted_tool_message = True + flush_text() + messages.append({"role": "user", "content": [_image_block_to_openai_part(block)]}) + else: + raise MalformedRequestError(f"Unsupported tool_result block type: {block_type}") + flush_text() + if not messages: + messages.append({"role": "tool", "tool_call_id": tool_call_id, "content": ""}) + return messages + + +def _user_messages_from_content(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"role": "user", "content": content}] + if not isinstance(content, list): + raise MalformedRequestError("user content must be a string or block list") + + messages: list[dict[str, Any]] = [] + parts: list[dict[str, Any]] = [] # A single openai message can carry multiple text/image parts. + + def flush_user() -> None: + # Collapse pure-text parts to a single string so downstream chat + # templates take their well-trodden string path; keep list form only + # when an image is present (image_url has no string fallback in the + # OpenAI protocol). Claude Code splits content into multiple text + # blocks to attach cache_control per block, so join with \n to keep + # those blocks from running together. + if not parts: + return + if any(part.get("type") != "text" for part in parts): + content = list(parts) + else: + content = "\n".join(part["text"] for part in parts) + messages.append({"role": "user", "content": content}) + parts.clear() + + for block in content: + if not isinstance(block, dict): + raise MalformedRequestError("user content blocks must be objects") + block_type = block.get("type") + if block_type == "text": + text = block.get("text", "") + if not isinstance(text, str): + raise MalformedRequestError("user text must be a string") + parts.append({"type": "text", "text": text}) + elif block_type == "image": + parts.append(_image_block_to_openai_part(block)) + elif block_type == "tool_result": + flush_user() + tool_call_id = block.get("tool_use_id") + if not isinstance(tool_call_id, str) or not tool_call_id: + raise MalformedRequestError("tool_result.tool_use_id must be a non-empty string") + messages.extend(_tool_result_messages(tool_call_id, block.get("content", ""))) + else: + raise MalformedRequestError(f"Unsupported user block type: {block_type}") + flush_user() + return messages + + +def _assistant_message_from_content(content: Any) -> dict[str, Any]: + if isinstance(content, str): + return {"role": "assistant", "content": content} + if not isinstance(content, list): + raise MalformedRequestError("assistant content must be a string or block list") + + texts: list[str] = [] + tool_calls: list[dict[str, Any]] = [] + for block in content: + if not isinstance(block, dict): + raise MalformedRequestError("assistant content blocks must be objects") + block_type = block.get("type") + if block_type == "text": + text = block.get("text", "") + if not isinstance(text, str): + raise MalformedRequestError("assistant text must be a string") + texts.append(text) + elif block_type == "tool_use": + tool_id = block.get("id") + if not isinstance(tool_id, str) or not tool_id: + raise MalformedRequestError("tool_use.id must be a non-empty string") + name = block.get("name") + if not isinstance(name, str): + raise MalformedRequestError("tool_use.name must be a string") + arguments = block.get("input", {}) + if not isinstance(arguments, dict): + raise MalformedRequestError("tool_use.input must be an object") + tool_calls.append( + { + "id": tool_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + ) + elif block_type == "thinking": + # Self-hosted SGLang never emits Anthropic thinking blocks, so seeing + # one in inbound history means the client mixed sources. Drop with a + # warning rather than crash; ... tokens, if present, + # are already preserved as plain text on the same assistant message. + logger.warning("Dropping Anthropic thinking block from assistant history; multi-turn prefix may drift") + continue + elif block_type == "redacted_thinking": + # Anthropic-encrypted reasoning signature with byte-level fidelity; + # no faithful representation in our chat-template path, and silently + # mapping it to plain text would corrupt training distribution. + raise MalformedRequestError("redacted_thinking blocks are not supported") + else: + raise MalformedRequestError(f"Unsupported assistant block type: {block_type}") + + message: dict[str, Any] = {"role": "assistant", "content": "\n".join(texts)} + if tool_calls: + message["tool_calls"] = tool_calls + return message + + +def _fold_reminder_into_user(message: dict[str, Any], text: str, *, prepend: bool = False) -> None: + """Fold a system-text reminder into an existing user message in place. + + Content is str | list of Anthropic blocks; both forms are handled. Mid-list + `` text is anchored to a neighbouring user message rather + than hoisted to index 0, since many chat templates reject system messages + beyond index 0. `text` is the fully-wrapped reminder string — the caller is + responsible for adding the `...` envelope + before passing it in. + """ + content = message.get("content", "") + if isinstance(content, list): + part = {"type": "text", "text": text} + if prepend: + content.insert(0, part) + else: + content.append(part) + return + if not content: + message["content"] = text + elif prepend: + message["content"] = f"{text}\n{content}" + else: + message["content"] = f"{content}\n{text}" + + +def _fold_mid_list_system_into_user(messages: Any) -> Any: + """Fold every `messages[]`-level system entry into a neighbouring user + message as a `` block, returning a system-free list. + + Anthropic carries the real system prompt in the top-level `system` field; a + `role: system` entry inside `messages` is a mid-conversation reminder that + clients like Claude Code inject. Many chat templates reject any system + message past index 0 *and* reject consecutive same-role messages, so a + reminder must be merged into an adjacent user message rather than emitted + standalone. We fold index-0 too: keeping it as a leading system would be + pushed to index 1 by the top-level system prompt, which is exactly what + templates reject. Non-list input is passed through for the caller to reject. + + Strategy: single forward pass, preferring the preceding user (append). When + no preceding user exists yet, buffer the reminder and prepend it into the + next user; if neither exists, emit a standalone user at the end (last + resort: violates same-role-adjacency only when the whole conversation has + no user message, which is itself ill-formed). + """ + if not isinstance(messages, list): + return messages + if not any(isinstance(m, dict) and m.get("role") == "system" for m in messages): + return messages + + out: list[Any] = [] + pending = "" + for message in messages: + if isinstance(message, dict) and message.get("role") == "system": + system_text = _system_to_text(message.get("content", "")) + if not system_text: + continue + reminder = f"\n{system_text}\n" + prev_user = next( + (m for m in reversed(out) if isinstance(m, dict) and m.get("role") == "user"), + None, + ) + if prev_user is not None: + _fold_reminder_into_user(prev_user, reminder) + else: + pending = f"{pending}\n{reminder}" if pending else reminder + continue + if pending and isinstance(message, dict) and message.get("role") == "user": + _fold_reminder_into_user(message, pending, prepend=True) + pending = "" + out.append(message) + if pending: + out.append({"role": "user", "content": pending}) + return out + + +def _messages_to_internal(messages: Any) -> list[dict[str, Any]]: + # Caller folds mid-list system into user first, so only user/assistant + # remain here (a stray system role is therefore unsupported). + if not isinstance(messages, list) or not messages: + raise MalformedRequestError("messages must be non-empty") + + result: list[dict[str, Any]] = [] + for message in messages: + if not isinstance(message, dict): + raise MalformedRequestError("messages entries must be objects") + role = message.get("role") + content = message.get("content", "") + if role == "user": + result.extend(_user_messages_from_content(content)) + elif role == "assistant": + result.append(_assistant_message_from_content(content)) + else: + raise MalformedRequestError(f"Unsupported message role: {role}") + return result + + +def _convert_tools(tools: Any) -> list[dict[str, Any]] | None: + if tools is None: + return None + if not isinstance(tools, list): + raise MalformedRequestError("tools must be a list") + + converted: list[dict[str, Any]] = [] + for tool in tools: + if not isinstance(tool, dict): + raise MalformedRequestError("tools entries must be objects") + name = tool.get("name") + if not isinstance(name, str): + raise MalformedRequestError("tool.name must be a string") + input_schema = copy.deepcopy(tool.get("input_schema", {})) + function: dict[str, Any] = {"name": name, "parameters": input_schema} + description = tool.get("description") + if isinstance(description, str): + function["description"] = description + # cache_control is request/cache metadata; internal tools keep only the + # OpenAI function schema. + converted.append({"type": "function", "function": function}) + return converted + + +def _apply_tool_choice(payload: dict[str, Any], tools: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None: + tool_choice = payload.get("tool_choice", "auto") + if isinstance(tool_choice, str): + if tool_choice == "auto": + return tools + if tool_choice == "none": + return None + raise MalformedRequestError(f"Unsupported tool_choice: {tool_choice}") + if isinstance(tool_choice, dict): + choice_type = tool_choice.get("type") + if choice_type == "auto": + return tools + if choice_type == "none": + return None + raise MalformedRequestError(f"Unsupported tool_choice type: {choice_type}") + raise MalformedRequestError("tool_choice must be a string or object") + + +def anthropic_to_internal( + payload: dict, + *, + base_sampling_params: dict, + allowed_sampling_keys: frozenset[str], +) -> InternalGenerationRequest: + """Lower an Anthropic Messages request into the internal request shape.""" + if not isinstance(payload, dict): + raise MalformedRequestError("payload must be an object") + + # Message/system lowering handles Anthropic-specific compatibility downgrades + # before the session sees the OpenAI-like template-facing canonical. + messages = _messages_to_internal(_fold_mid_list_system_into_user(payload.get("messages"))) + system_text = _system_to_text(payload.get("system")) + if system_text: + messages.insert(0, {"role": "system", "content": system_text}) + + # Tool conversion and sampling stay adapter-owned; the session consumes only + # internal tools plus codec-keyed sampling params. + tools = _apply_tool_choice(payload, _convert_tools(payload.get("tools"))) + sampling_params = dict(base_sampling_params) + for key in ("max_tokens", "temperature", "top_p", "top_k"): + if key in allowed_sampling_keys and key in payload: + sampling_params[key] = payload[key] + if "stop" in allowed_sampling_keys and "stop_sequences" in payload: + sampling_params["stop"] = payload["stop_sequences"] + # Anthropic cache_control can appear on request blocks; it is intentionally + # ignored because generation sampling params have no equivalent field. + return { + "messages": messages, + "tools": tools, + "chat_template_kwargs": {}, + "sampling_params": sampling_params, + } diff --git a/uni_agent/gateway/adapters/openai.py b/uni_agent/gateway/adapters/openai.py new file mode 100644 index 00000000..299bf475 --- /dev/null +++ b/uni_agent/gateway/adapters/openai.py @@ -0,0 +1,225 @@ +"""OpenAI Chat Completions wire <-> InternalGenerationRequest. + +Actor calls this before session; MessageCodec never sees wire shape. +""" + +from __future__ import annotations + +import json +import time +from collections.abc import AsyncIterator +from typing import Any +from uuid import uuid4 + +from fastapi.responses import StreamingResponse + +from uni_agent.gateway.session.session import GenerationOutcome +from uni_agent.gateway.session.types import InternalGenerationRequest + +from .types import MalformedRequestError + +_SSE_HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", +} + +_OPENAI_ERROR_TYPE_BY_STATUS = { + # Only status codes the gateway actually emits today (gateway.py / + # session.py: 400 malformed/JSON, 409 concurrent session generation, + # 500 internal). 401/403/404/422/429 are intentionally omitted: this + # gateway has no auth / rate-limit / routing paths that surface them. + 400: "invalid_request_error", + 409: "conflict_error", +} + + +def openai_error_body(status_code: int, message: str) -> dict[str, Any]: + error_type = _OPENAI_ERROR_TYPE_BY_STATUS.get( + status_code, + "invalid_request_error" if 400 <= status_code < 500 else "internal_server_error", + ) + return { + "error": { + "message": message, + "type": error_type, + "code": None, + "param": None, + } + } + + +def openai_build_response(outcome: GenerationOutcome, *, model: str) -> dict[str, Any]: + return { + "id": f"chatcmpl-{uuid4().hex}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [{"index": 0, "message": outcome.assistant_msg, "finish_reason": outcome.finish_reason}], + "usage": { + "prompt_tokens": outcome.prompt_tokens, + "completion_tokens": outcome.completion_tokens, + "total_tokens": outcome.prompt_tokens + outcome.completion_tokens, + }, + } + + +def openai_stream_response(outcome: GenerationOutcome, *, model: str) -> StreamingResponse: + """Synthesize an OpenAI chat.completion.chunk SSE stream from a completed outcome.""" + chunk_id = f"chatcmpl-{uuid4().hex}" + created = int(time.time()) + + def _chunk(delta: dict[str, Any], finish: str | None, usage: dict[str, int] | None = None) -> str: + body = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + if usage is not None: + body["usage"] = usage + return f"data: {json.dumps(body, ensure_ascii=False)}\n\n" + + async def _gen() -> AsyncIterator[bytes]: + msg = outcome.assistant_msg + yield _chunk({"role": "assistant"}, None).encode() + if isinstance(msg.get("reasoning_content"), str) and msg["reasoning_content"]: + yield _chunk({"reasoning_content": msg["reasoning_content"]}, None).encode() + if isinstance(msg.get("content"), str) and msg["content"]: + yield _chunk({"content": msg["content"]}, None).encode() + if msg.get("tool_calls"): + tool_calls = [{**tool_call, "index": idx} for idx, tool_call in enumerate(msg["tool_calls"])] + yield _chunk({"tool_calls": tool_calls}, None).encode() + usage = { + "prompt_tokens": outcome.prompt_tokens, + "completion_tokens": outcome.completion_tokens, + "total_tokens": outcome.prompt_tokens + outcome.completion_tokens, + } + yield _chunk({}, outcome.finish_reason, usage).encode() + yield b"data: [DONE]\n\n" + + return StreamingResponse(_gen(), media_type="text/event-stream", headers=_SSE_HEADERS) + + +def _normalize_message_content(content: Any) -> Any: + """Normalize message content: coerce None to empty string, validate type.""" + if isinstance(content, list | dict | str): + return content + if content is None: + return "" + raise MalformedRequestError(f"Unsupported content type: {type(content).__name__}") + + +def _normalize_tool_calls(tool_calls: Any) -> list[dict[str, Any]]: + """Validate tool_calls and parse JSON-string function arguments.""" + if not isinstance(tool_calls, list): + raise MalformedRequestError("tool_calls must be a list") + result = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + raise MalformedRequestError("tool_calls entries must be objects") + function = tool_call.get("function") + if not isinstance(function, dict): + raise MalformedRequestError("tool_call.function must be an object") + + normalized_tool_call = dict(tool_call) + normalized_function = dict(function) + arguments = normalized_function.get("arguments") + if isinstance(arguments, str): + try: + normalized_function["arguments"] = json.loads(arguments) + except (json.JSONDecodeError, TypeError): + pass + normalized_tool_call["function"] = normalized_function + result.append(normalized_tool_call) + return result + + +def _normalize_message(message: Any) -> dict[str, Any]: + """Normalize a single message and filter to known OpenAI chat fields.""" + if not isinstance(message, dict): + raise MalformedRequestError("messages entries must be objects") + + role = message.get("role") + if not isinstance(role, str) or not role: + raise MalformedRequestError("message.role must be a non-empty string") + + normalized: dict[str, Any] = { + "role": role, + "content": _normalize_message_content(message.get("content", "")), + } + if "name" in message: + name = message["name"] + if not isinstance(name, str): + raise MalformedRequestError("message.name must be a string") + normalized["name"] = name + if "tool_calls" in message: + normalized["tool_calls"] = _normalize_tool_calls(message["tool_calls"]) + if "tool_call_id" in message: + normalized["tool_call_id"] = str(message["tool_call_id"]) + if "reasoning_content" in message: + reasoning_content = message["reasoning_content"] + if reasoning_content is not None and not isinstance(reasoning_content, str): + raise MalformedRequestError("message.reasoning_content must be a string or null") + normalized["reasoning_content"] = reasoning_content + return normalized + + +def openai_to_internal( + payload: dict, + *, + base_sampling_params: dict, + allowed_sampling_keys: frozenset[str], +) -> InternalGenerationRequest: + """Lower an OpenAI chat-completions request into the internal request shape.""" + # Capability gates: these OpenAI request features have no gateway/session + # equivalent yet, so reject them before building the internal request. + n_value = payload.get("n", 1) + if n_value != 1: + raise MalformedRequestError(f"n={n_value} is not supported (only n=1)") + if payload.get("response_format") is not None: + raise MalformedRequestError("response_format is not supported") + + tool_choice_payload = payload.get("tool_choice", "auto") + if isinstance(tool_choice_payload, str): + tool_choice = tool_choice_payload.lower() + if tool_choice not in {"auto", "none"}: + raise MalformedRequestError( + f'tool_choice="{tool_choice_payload}" is not supported (only "auto" / "none" are supported)' + ) + elif isinstance(tool_choice_payload, dict): + raise MalformedRequestError( + 'tool_choice with a specific function is not supported (only "auto" / "none" are supported)' + ) + else: + raise MalformedRequestError("tool_choice must be a string or object") + + # Required payload fields and template kwargs. + messages = payload.get("messages") + if not isinstance(messages, list) or not messages: + raise MalformedRequestError("messages must be non-empty") + + chat_template_kwargs = payload.get("chat_template_kwargs") + if chat_template_kwargs is not None and not isinstance(chat_template_kwargs, dict): + raise MalformedRequestError("chat_template_kwargs must be an object") + + # Tool injection policy. + tools = payload.get("tools") + if tools is not None and not isinstance(tools, list): + raise MalformedRequestError("tools must be a list") + if tool_choice == "none": + tools = None + + # Sampling params are gateway-owned allowlist merges, not message canonicalization. + sampling_params = dict(base_sampling_params) + for key in allowed_sampling_keys: + if key in payload: + sampling_params[key] = payload[key] + + return { + "messages": [_normalize_message(message) for message in messages], + "tools": tools, + "chat_template_kwargs": dict(chat_template_kwargs) if chat_template_kwargs else {}, + "sampling_params": sampling_params, + } diff --git a/uni_agent/gateway/adapters/types.py b/uni_agent/gateway/adapters/types.py new file mode 100644 index 00000000..1ffa35d5 --- /dev/null +++ b/uni_agent/gateway/adapters/types.py @@ -0,0 +1,137 @@ +"""Adapter-side wire protocol typing and request-lowering errors. + +These types document the provider wire fragments the adapters lower. They are +intentionally partial: validation remains in each adapter and the internal +shape stays OpenAI-like for chat templates. +""" + +from __future__ import annotations + +from typing import Any, Literal, NotRequired, TypedDict + + +class MalformedRequestError(ValueError): + """Raised when adapter request lowering rejects a client payload.""" + + pass + + +class OpenAIChatCompletionFunction(TypedDict, total=False): + """``tool_calls[i].function`` object inside an OpenAI chat message.""" + + name: str + arguments: Any # OpenAI spec is JSON string; gateway also accepts dict (Qwen-style chat templates) + + +class OpenAIChatCompletionToolCall(TypedDict, total=False): + """One entry in an OpenAI assistant message's ``tool_calls`` array.""" + + id: str + type: Literal["function"] + function: OpenAIChatCompletionFunction + + +class OpenAIChatMessage(TypedDict, total=False): + """A single OpenAI chat-completion message. + + Includes OpenAI-compatible extension ``reasoning_content`` that the gateway + preserves on input. + """ + + role: str + content: Any + name: str + tool_calls: list[OpenAIChatCompletionToolCall] + tool_call_id: str + reasoning_content: str | None + + +class OpenAIChatCompletionTool(TypedDict, total=False): + """One entry in an OpenAI request ``tools`` array.""" + + type: Literal["function"] + function: dict[str, Any] + + +class OpenAIChatCompletionRequest(TypedDict, total=False): + """Incoming ``POST /v1/chat/completions`` request body shape. + + ``chat_template_kwargs`` is an OpenAI-compatible server extension used by + the gateway to forward per-request chat template overrides (e.g. + ``enable_thinking``) into ``MessageCodec``. + """ + + model: str + messages: list[OpenAIChatMessage] + tools: list[OpenAIChatCompletionTool] + tool_choice: Any + stream: bool + n: int + response_format: Any + chat_template_kwargs: dict[str, Any] + temperature: float + top_p: float + top_k: int + max_tokens: int + stop: str | list[str] + + +class OpenAIChatCompletionUsage(TypedDict): + """Token usage block inside an OpenAI response.""" + + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class OpenAIChatCompletionChoice(TypedDict): + """One entry in an OpenAI response ``choices`` array. + + ``finish_reason`` is the gateway's internal value; the OpenAI adapter passes + it through as the OpenAI-compatible response value. + """ + + index: int + message: OpenAIChatMessage + finish_reason: Literal["stop", "length", "tool_calls", "content_filter", "function_call"] + + +class OpenAIChatCompletionResponse(TypedDict): + """Outgoing ``POST /v1/chat/completions`` response body shape.""" + + id: str + object: Literal["chat.completion"] + created: int + model: str + choices: list[OpenAIChatCompletionChoice] + usage: OpenAIChatCompletionUsage + + +class AnthropicContentBlock(TypedDict, total=False): + type: str + text: str + source: dict[str, Any] + id: str + name: str + input: dict[str, Any] + tool_use_id: str + content: str | list[dict[str, Any]] + + +class AnthropicMessage(TypedDict): + role: str + content: str | list[AnthropicContentBlock] + + +class AnthropicRequest(TypedDict): + model: NotRequired[str] + messages: list[AnthropicMessage] + system: NotRequired[str | list[AnthropicContentBlock]] + tools: NotRequired[list[dict[str, Any]]] + tool_choice: NotRequired[str | dict[str, Any]] + stream: NotRequired[bool] + max_tokens: NotRequired[int] + stop_sequences: NotRequired[list[str]] + temperature: NotRequired[float] + top_p: NotRequired[float] + top_k: NotRequired[int] diff --git a/uni_agent/gateway/config.py b/uni_agent/gateway/config.py index ccf18f5d..7d3a323f 100644 --- a/uni_agent/gateway/config.py +++ b/uni_agent/gateway/config.py @@ -24,7 +24,7 @@ class GatewayActorConfig: apply_chat_template_kwargs: Default kwargs passed to chat-template rendering. base_sampling_params: Sampling params applied before per-request overrides. allowed_request_sampling_param_keys: Request sampling keys accepted by the - codec when merging payload sampling params. + provider adapters when merging payload sampling params. vision_info_extractor: Optional async extractor for image/video inputs. vision_info_extractor_kwargs: Static kwargs forwarded to the extractor. prompt_length: Optional prompt-token budget stored on gateway sessions. @@ -36,7 +36,7 @@ class GatewayActorConfig: tool_parser_name: str | None = None apply_chat_template_kwargs: dict[str, Any] | None = None base_sampling_params: dict[str, Any] | None = None - allowed_request_sampling_param_keys: frozenset[str] | None = None + allowed_request_sampling_param_keys: set[str] | frozenset[str] | None = None vision_info_extractor: Callable | None = None vision_info_extractor_kwargs: dict[str, Any] | None = None prompt_length: int | None = None diff --git a/uni_agent/gateway/gateway.py b/uni_agent/gateway/gateway.py index e2966e0b..0846771f 100644 --- a/uni_agent/gateway/gateway.py +++ b/uni_agent/gateway/gateway.py @@ -1,21 +1,35 @@ -"""Thin FastAPI/Ray actor layer for gateway routing and JSON serialization.""" +"""Thin FastAPI/Ray actor layer for routing and session ownership. + +The actor owns routing, capability gates, and per-session ``GatewaySession`` +instances. Provider adapters own wire-to-internal translation and the response +or SSE envelopes returned to clients. +""" from __future__ import annotations import asyncio -import time -from logging import getLogger +import logging from typing import Any -from uuid import uuid4 import ray from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse +from uni_agent.gateway.adapters.anthropic import ( + anthropic_build_response, + anthropic_error_body, + anthropic_stream_response, + anthropic_to_internal, +) +from uni_agent.gateway.adapters.openai import ( + openai_build_response, + openai_error_body, + openai_stream_response, + openai_to_internal, +) +from uni_agent.gateway.adapters.types import AnthropicRequest, MalformedRequestError, OpenAIChatCompletionRequest from uni_agent.gateway.config import GatewayActorConfig from uni_agent.gateway.session import ( - ChatCompletionRequest, - ChatCompletionResponse, GatewaySession, MessageCodec, SessionHandle, @@ -23,14 +37,18 @@ ) from verl.workers.rollout.utils import run_uvicorn +DEFAULT_ALLOWED_REQUEST_SAMPLING_KEYS = frozenset({"temperature", "top_p", "top_k", "max_tokens", "stop"}) + +logger = logging.getLogger("gateway") + class _GatewayActor: """Ray actor implementation exposed as ``GatewayActor = ray.remote(...)``. Runtime and manager callers invoke public methods with - ``actor.method.remote(...)``. The actor owns FastAPI routing, OpenAI - capability gates, JSON response envelopes, and per-session - ``GatewaySession`` instances. + ``actor.method.remote(...)``. The actor owns FastAPI routing, provider + capability gates, response envelopes, and per-session ``GatewaySession`` + instances. """ def __init__(self, config: GatewayActorConfig, backend): @@ -46,8 +64,12 @@ def __init__(self, config: GatewayActorConfig, backend): vision_info_extractor_kwargs=config.vision_info_extractor_kwargs, tool_parser_name=config.tool_parser_name, apply_chat_template_kwargs=config.apply_chat_template_kwargs, - base_sampling_params=config.base_sampling_params, - allowed_request_sampling_param_keys=config.allowed_request_sampling_param_keys, + ) + self._base_sampling_params = dict(config.base_sampling_params or {}) + self._allowed_request_sampling_param_keys = ( + DEFAULT_ALLOWED_REQUEST_SAMPLING_KEYS + if config.allowed_request_sampling_param_keys is None + else frozenset(config.allowed_request_sampling_param_keys) ) self._prompt_length = config.prompt_length self._response_length = config.response_length @@ -59,33 +81,53 @@ def __init__(self, config: GatewayActorConfig, backend): self._register_routes() def _register_routes(self) -> None: - """Register HTTP handlers for chat completions and reward metadata.""" + """Register provider HTTP handlers and reward metadata.""" + + def _error_body_for_path(path: str, status_code: int, message: str) -> dict[str, Any]: + if path.endswith("/v1/messages"): + return anthropic_error_body(status_code, message) + return openai_error_body(status_code, message) @self._app.exception_handler(HTTPException) - async def _http_exception_handler(_request: Request, exc: HTTPException): + async def _http_exception_handler(request: Request, exc: HTTPException): if isinstance(exc.detail, str): message = exc.detail elif isinstance(exc.detail, dict) and "message" in exc.detail: message = str(exc.detail["message"]) else: message = str(exc.detail) - error_type = "invalid_request_error" if 400 <= exc.status_code < 500 else "internal_server_error" return JSONResponse( status_code=exc.status_code, - content={ - "error": { - "message": message, - "type": error_type, - "code": None, - "param": None, - } - }, + content=_error_body_for_path(request.url.path, exc.status_code, message), + ) + + @self._app.exception_handler(Exception) + async def _unhandled_exception_handler(request: Request, exc: Exception): + # Any exception that escapes the route handlers (programmer bugs, + # ungraceful third-party failures, etc.) reaches here. We log the + # full traceback for diagnosis but return a provider-shaped error + # body so client SDKs can still parse it. + logger.exception("Unhandled exception in gateway route %s", request.url.path) + return JSONResponse( + status_code=500, + content=_error_body_for_path(request.url.path, 500, "Internal server error"), ) @self._app.post("/sessions/{session_id}/v1/chat/completions") - async def _chat_completions(session_id: str, request: Request): - payload = await request.json() - return await self._handle_chat_completions(session_id=session_id, payload=payload) + async def _openai_chat_completions(session_id: str, request: Request): + try: + payload = await request.json() + except ValueError as exc: + raise HTTPException(status_code=400, detail="Invalid JSON body") from exc + return await self._handle_openai_chat_completions(session_id=session_id, payload=payload) + + @self._app.post("/sessions/{session_id}/v1/messages") + async def _anthropic_messages(session_id: str, request: Request): + try: + payload = await request.json() + except ValueError as exc: + raise HTTPException(status_code=400, detail="Invalid JSON body") from exc + return await self._handle_anthropic_messages(session_id=session_id, payload=payload) @self._app.post("/sessions/{session_id}/reward_info") async def _reward_info(session_id: str, request: Request): @@ -109,58 +151,55 @@ def _get_session(self, session_id: str) -> GatewaySession: raise KeyError(f"Unknown session_id: {session_id}") return session - async def _handle_chat_completions( + async def _handle_openai_chat_completions( self, session_id: str, - payload: ChatCompletionRequest, - ) -> JSONResponse: - """Validate a chat-completion payload and serialize the session outcome.""" + payload: OpenAIChatCompletionRequest, + ) -> JSONResponse | StreamingResponse: + """Validate an OpenAI Chat Completions payload and serialize the session outcome.""" session = self._sessions.get(session_id) if session is None: raise HTTPException(status_code=404, detail=f"Unknown session_id: {session_id}") - if payload.get("stream") is True: - getLogger("gateway").warning( - "session=%s stream=true requested; gateway returns non-streaming response", - session_id, - ) - n_value = payload.get("n", 1) - if n_value != 1: - raise HTTPException(status_code=400, detail=f"n={n_value} is not supported (only n=1)") - if payload.get("response_format") is not None: - raise HTTPException(status_code=400, detail="response_format is not supported") - tool_choice_payload = payload.get("tool_choice") - if isinstance(tool_choice_payload, dict): - raise HTTPException( - status_code=400, - detail='tool_choice with a specific function is not supported (only "auto" / "none" are supported)', + try: + internal = openai_to_internal( + payload, + base_sampling_params=self._base_sampling_params, + allowed_sampling_keys=self._allowed_request_sampling_param_keys, ) - if isinstance(tool_choice_payload, str) and tool_choice_payload.lower() == "required": - raise HTTPException( - status_code=400, - detail='tool_choice="required" is not supported (only "auto" / "none" are supported)', + except MalformedRequestError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + outcome = await session.run_generation(internal, self._backend) + model = str(payload.get("model") or "unknown") + if payload.get("stream") is True: + return openai_stream_response(outcome, model=model) + return JSONResponse(openai_build_response(outcome, model=model)) + + async def _handle_anthropic_messages( + self, + session_id: str, + payload: AnthropicRequest, + ) -> JSONResponse | StreamingResponse: + """Validate an Anthropic Messages payload and serialize the session outcome.""" + session = self._sessions.get(session_id) + if session is None: + raise HTTPException(status_code=404, detail=f"Unknown session_id: {session_id}") + + try: + internal = anthropic_to_internal( + payload, + base_sampling_params=self._base_sampling_params, + allowed_sampling_keys=self._allowed_request_sampling_param_keys, ) + except MalformedRequestError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc - outcome = await session.run_generation(payload, self._backend) - response: ChatCompletionResponse = { - "id": f"chatcmpl-{uuid4().hex}", - "object": "chat.completion", - "created": int(time.time()), - "model": str(payload.get("model") or "unknown"), - "choices": [ - { - "index": 0, - "message": outcome.assistant_msg, - "finish_reason": outcome.finish_reason, - } - ], - "usage": { - "prompt_tokens": outcome.prompt_tokens, - "completion_tokens": outcome.completion_tokens, - "total_tokens": outcome.prompt_tokens + outcome.completion_tokens, - }, - } - return JSONResponse(response) + outcome = await session.run_generation(internal, self._backend) + model = str(payload.get("model") or "unknown") + if payload.get("stream") is True: + return anthropic_stream_response(outcome, model=model) + return JSONResponse(anthropic_build_response(outcome, model=model)) async def start(self) -> None: """Start the FastAPI server backing this gateway actor.""" @@ -183,7 +222,7 @@ async def shutdown(self) -> None: self._server_base_url = None async def create_session(self, session_id: str, metadata: dict[str, Any] | None = None) -> SessionHandle: - """Create an actor-owned session and return its OpenAI-compatible handle.""" + """Create an actor-owned session and return its provider-compatible handle.""" self._require_started() if session_id in self._sessions: raise RuntimeError(f"Session {session_id} already exists") diff --git a/uni_agent/gateway/session/__init__.py b/uni_agent/gateway/session/__init__.py index f9c9e2b0..eb39e722 100644 --- a/uni_agent/gateway/session/__init__.py +++ b/uni_agent/gateway/session/__init__.py @@ -1,21 +1,18 @@ -"""Session domain for the gateway: per-session state, codec, and wire protocol. +"""Session domain for the gateway: per-session state and model codec. The gateway is a thin HTTP layer; this package holds the session-side logic it -serves — trajectory buffering, message encoding/decoding, and the chat-completion -protocol types. ``SessionHandle`` and ``Trajectory`` are the cross-package public -surface (consumed by the framework runners). +serves: trajectory buffering and message encoding/decoding. +``SessionHandle`` / ``Trajectory`` are consumed by framework runners, while +``InternalGenerationRequest`` is the adapter-to-session request boundary. """ -from .codec import MalformedRequestError, MessageCodec -from .protocol import ChatCompletionRequest, ChatCompletionResponse +from .codec import MessageCodec from .session import GatewaySession, TrajectoryBuffer -from .types import SessionHandle, Trajectory +from .types import InternalGenerationRequest, SessionHandle, Trajectory __all__ = [ - "ChatCompletionRequest", - "ChatCompletionResponse", + "InternalGenerationRequest", "GatewaySession", - "MalformedRequestError", "MessageCodec", "SessionHandle", "Trajectory", diff --git a/uni_agent/gateway/session/codec.py b/uni_agent/gateway/session/codec.py index 747dbde8..648752ca 100644 --- a/uni_agent/gateway/session/codec.py +++ b/uni_agent/gateway/session/codec.py @@ -1,34 +1,24 @@ -"""Model-scoped codec for tokenizer, processor, tool-parser, and sampling boundaries.""" +"""Model-scoped codec for tokenizer, processor, tool-parser, and decode paths. + +This layer stays within the model boundary: it applies chat templates, handles +processor-backed multimodal inputs, parses tools, and decodes backend outputs. +""" from __future__ import annotations import json +import logging +from types import SimpleNamespace from typing import Any from uuid import uuid4 -from verl.experimental.agent_loop.tool_parser import ToolParser from verl.utils.chat_template import apply_chat_template as _apply_chat_template from verl.utils.chat_template import initialize_system_prompt from verl.utils.tokenizer import normalize_token_ids +logger = logging.getLogger("gateway") -class MalformedRequestError(ValueError): - """Raised when request payload normalization finds unsupported input.""" - - pass - - -_DEFAULT_ALLOWED_REQUEST_SAMPLING_PARAM_KEYS = frozenset( - { - "temperature", - "top_p", - "top_k", - "max_tokens", - } -) - - -# Map backend stop_reason values to OpenAI-spec finish_reason values. +# Map backend stop_reason values into the gateway's internal finish_reason vocabulary. _FINISH_REASON_MAP = { "completed": "stop", "stop": "stop", @@ -40,78 +30,15 @@ class MalformedRequestError(ValueError): "abort": "stop", } +_SGLANG_TOOL_PARSER_ALIASES = { + "qwen3_xml": "qwen3_coder", +} -def _normalize_message_content(content: Any) -> Any: - """Normalize message content: coerce None to empty string, validate type.""" - if isinstance(content, list | dict | str): - return content - if content is None: - return "" - raise MalformedRequestError(f"Unsupported content type: {type(content).__name__}") - - -def _normalize_tool_calls(tool_calls: Any) -> list[dict[str, Any]]: - """Validate tool_calls and parse JSON-string function arguments.""" - if not isinstance(tool_calls, list): - raise MalformedRequestError("tool_calls must be a list") - result = [] - for tool_call in tool_calls: - if not isinstance(tool_call, dict): - raise MalformedRequestError("tool_calls entries must be objects") - function = tool_call.get("function") - if not isinstance(function, dict): - raise MalformedRequestError("tool_call.function must be an object") - - normalized_tool_call = dict(tool_call) - normalized_function = dict(function) - arguments = normalized_function.get("arguments") - if isinstance(arguments, str): - try: - normalized_function["arguments"] = json.loads(arguments) - except (json.JSONDecodeError, TypeError): - pass - normalized_tool_call["function"] = normalized_function - result.append(normalized_tool_call) - return result - - -def _normalize_message(message: Any) -> dict[str, Any]: - """Normalize a single message and filter to known OpenAI chat fields.""" - if not isinstance(message, dict): - raise MalformedRequestError("messages entries must be objects") - - role = message.get("role") - if not isinstance(role, str) or not role: - raise MalformedRequestError("message.role must be a non-empty string") - - normalized: dict[str, Any] = { - "role": role, - "content": _normalize_message_content(message.get("content", "")), - } - if "name" in message: - name = message["name"] - if not isinstance(name, str): - raise MalformedRequestError("message.name must be a string") - normalized["name"] = name - if "tool_calls" in message: - normalized["tool_calls"] = _normalize_tool_calls(message["tool_calls"]) - if "tool_call_id" in message: - normalized["tool_call_id"] = str(message["tool_call_id"]) - if "reasoning_content" in message: - reasoning_content = message["reasoning_content"] - if reasoning_content is not None and not isinstance(reasoning_content, str): - raise MalformedRequestError("message.reasoning_content must be a string or null") - normalized["reasoning_content"] = reasoning_content - return normalized - - -def _validate_tools(tools: Any) -> list[Any] | None: - """Validate tools structure. Does not modify content.""" - if tools is None: - return None - if not isinstance(tools, list): - raise MalformedRequestError("tools must be a list") - return tools +_VLLM_TOOL_PARSER_ALIASES = { + "qwen": "qwen3_xml", + "qwen25": "qwen3_xml", + "qwen3": "qwen3_xml", +} def _canonicalize_tool_arguments_for_comparison(arguments: Any) -> tuple[str, Any]: @@ -125,13 +52,78 @@ def _canonicalize_tool_arguments_for_comparison(arguments: Any) -> tuple[str, An return ("raw", arguments) +def _process_tool_calls_sglang( + text: str, + tools: list[dict[str, Any]], + parser_name: str, +) -> tuple[str, list[Any]]: + from sglang.srt.entrypoints.openai.protocol import Function as SglFunction + from sglang.srt.entrypoints.openai.protocol import Tool as SglTool + from sglang.srt.function_call.function_call_parser import FunctionCallParser + + sglang_tools = [SglTool(type=tool["type"], function=SglFunction(**tool["function"])) for tool in tools] + parser = FunctionCallParser(sglang_tools, parser_name) + if not parser.has_tool_call(text): + return text, [] + + content, calls = parser.parse_non_stream(text) + return content, [SimpleNamespace(name=call.name, arguments=call.parameters) for call in calls] + + +def _process_tool_calls_vllm( + text: str, + tools: list[dict[str, Any]], + parser_name: str, + tokenizer, +) -> tuple[str, list[Any]]: + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionToolsParam + from vllm.tool_parsers import ToolParserManager + + parser_cls = ToolParserManager.get_tool_parser(parser_name) + parser = parser_cls(tokenizer) + request = SimpleNamespace( + tools=[ChatCompletionToolsParam(**tool) if isinstance(tool, dict) else tool for tool in tools], + tool_choice="auto", + skip_special_tokens=True, + ) + parsed = parser.extract_tool_calls(text, request) + if not parsed.tools_called: + return text, [] + return parsed.content or "", [tool_call.function for tool_call in parsed.tool_calls] + + +def _extract_tool_calls_with_sglang_or_vllm( + text: str, + tools: list[dict[str, Any]], + parser_name: str, + tokenizer, +) -> tuple[str, list[Any]]: + sglang_name = _SGLANG_TOOL_PARSER_ALIASES.get(parser_name, parser_name) + try: + return _process_tool_calls_sglang(text, tools, sglang_name) + except ModuleNotFoundError: + pass + except Exception: + logger.warning("SGLang tool-call parsing failed; trying vLLM", exc_info=True) + + vllm_name = _VLLM_TOOL_PARSER_ALIASES.get(parser_name, parser_name) + try: + return _process_tool_calls_vllm(text, tools, vllm_name, tokenizer) + except ModuleNotFoundError: + pass + except Exception: + logger.warning("vLLM tool-call parsing failed; returning raw text", exc_info=True) + + return text, [] + + class MessageCodec: - """Model-scoped request codec used by gateway sessions and actors. + """Model-scoped request codec used by gateway sessions. ``_GatewayActor`` owns one codec per actor and injects it into - ``GatewaySession`` instances. The codec normalizes OpenAI-shaped request - payloads, renders chat templates, handles multimodal processor inputs, and - decodes backend token outputs without reading session state. + ``GatewaySession`` instances. The codec renders chat templates, handles + multimodal processor inputs, and decodes backend token outputs without + reading session state. """ def __init__( @@ -143,25 +135,17 @@ def __init__( vision_info_extractor_kwargs: dict[str, Any] | None = None, tool_parser_name: str | None = None, apply_chat_template_kwargs: dict[str, Any] | None = None, - base_sampling_params: dict[str, Any] | None = None, - allowed_request_sampling_param_keys: set[str] | frozenset[str] | None = None, ): self._tokenizer = tokenizer self._processor = processor self._vision_info_extractor = vision_info_extractor or self._default_vision_info_extractor self._vision_info_extractor_kwargs = dict(vision_info_extractor_kwargs or {}) self._apply_chat_template_kwargs = apply_chat_template_kwargs or {} - self._base_sampling_params = dict(base_sampling_params or {}) - self._allowed_request_sampling_param_keys = ( - _DEFAULT_ALLOWED_REQUEST_SAMPLING_PARAM_KEYS - if allowed_request_sampling_param_keys is None - else frozenset(allowed_request_sampling_param_keys) - ) self._system_prompt = initialize_system_prompt( tokenizer, **self._apply_chat_template_kwargs, ) - self._tool_parser = ToolParser.get_tool_parser(tool_parser_name, tokenizer) if tool_parser_name else None + self._tool_parser_name = tool_parser_name async def _default_vision_info_extractor( self, @@ -303,15 +287,14 @@ async def decode_response( stop_reason: str | None = None, ) -> tuple[dict[str, Any], str]: """Decode model output tokens into an assistant message and finish reason.""" - if self._tool_parser is not None and tools: - parsed_tools = None - try: - from verl.tools.schemas import OpenAIFunctionToolSchema - - parsed_tools = [OpenAIFunctionToolSchema(**t) if isinstance(t, dict) else t for t in tools] - except Exception: - pass - content, function_calls = await self._tool_parser.extract_tool_calls(response_ids, parsed_tools) + if self._tool_parser_name and tools: + response_text = self._tokenizer.decode(response_ids, skip_special_tokens=False) + content, function_calls = _extract_tool_calls_with_sglang_or_vllm( + response_text, + tools, + self._tool_parser_name, + self._tokenizer, + ) if function_calls: tool_calls = [ { @@ -331,30 +314,13 @@ async def decode_response( finish_reason = _FINISH_REASON_MAP.get(stop_reason, stop_reason) if stop_reason else "stop" return {"role": "assistant", "content": response_text}, finish_reason - def normalize_request(self, payload: dict[str, Any]) -> dict[str, Any]: - """Normalize chat messages, tool schemas, and chat-template kwargs.""" - messages = payload.get("messages") - if not isinstance(messages, list) or not messages: - raise MalformedRequestError("messages must be non-empty") - chat_template_kwargs = payload.get("chat_template_kwargs") - if chat_template_kwargs is not None and not isinstance(chat_template_kwargs, dict): - raise MalformedRequestError("chat_template_kwargs must be an object") - - tool_choice_payload = payload.get("tool_choice") - tool_choice = tool_choice_payload.lower() if isinstance(tool_choice_payload, str) else "auto" - tools = _validate_tools(payload.get("tools")) - if tool_choice == "none": - tools = None - - return { - "messages": [_normalize_message(message) for message in messages], - "tools": tools, - "chat_template_kwargs": dict(chat_template_kwargs) if chat_template_kwargs else {}, - } - def canonicalize_message_for_prefix_comparison(self, message: dict[str, Any]) -> dict[str, Any]: """Canonicalize one message before session prefix comparison.""" normalized = dict(message) + # Tool-result correlation ids are wire noise; prefix comparison ignores + # them while preserving unexpected top-level fields on other roles. + if normalized.get("role") == "tool": + normalized.pop("tool_call_id", None) tool_calls = normalized.get("tool_calls") if not isinstance(tool_calls, list): return normalized @@ -362,6 +328,7 @@ def canonicalize_message_for_prefix_comparison(self, message: dict[str, Any]) -> normalized_tool_calls: list[dict[str, Any]] = [] for tool_call in tool_calls: normalized_tool_call = dict(tool_call) + normalized_tool_call.pop("id", None) function = normalized_tool_call.get("function") if isinstance(function, dict) and "arguments" in function: normalized_function = dict(function) @@ -370,11 +337,3 @@ def canonicalize_message_for_prefix_comparison(self, message: dict[str, Any]) -> normalized_tool_calls.append(normalized_tool_call) normalized["tool_calls"] = normalized_tool_calls return normalized - - def build_sampling_params(self, payload: dict[str, Any]) -> dict[str, Any]: - """Merge allowed per-request sampling params with codec defaults.""" - sampling_params = dict(self._base_sampling_params) - for key in self._allowed_request_sampling_param_keys: - if key in payload: - sampling_params[key] = payload[key] - return sampling_params diff --git a/uni_agent/gateway/session/protocol.py b/uni_agent/gateway/session/protocol.py deleted file mode 100644 index 2fd2b67a..00000000 --- a/uni_agent/gateway/session/protocol.py +++ /dev/null @@ -1,101 +0,0 @@ -"""OpenAI-compatible Chat Completions HTTP protocol types for the gateway. - -These TypedDicts give the actor's HTTP layer a single source of truth -for request / response shape so that ``_handle_chat_completions`` no -longer constructs anonymous dicts. They also serve as documentation for -the expected shape of these objects. -""" - -from __future__ import annotations - -from typing import Any, Literal, TypedDict - - -class ChatCompletionFunction(TypedDict, total=False): - """``tool_calls[i].function`` object inside a ChatMessage.""" - - name: str - arguments: Any # OpenAI spec is JSON string; gateway also accepts dict (Qwen-style chat templates) - - -class ChatCompletionToolCall(TypedDict, total=False): - """One entry in an assistant message's ``tool_calls`` array.""" - - id: str - type: Literal["function"] - function: ChatCompletionFunction - - -class ChatMessage(TypedDict, total=False): - """A single chat-completion message (system / user / assistant / tool). - - Includes OpenAI-compatible extension ``reasoning_content`` that the - gateway preserves on input. - """ - - role: str - content: Any - name: str - tool_calls: list[ChatCompletionToolCall] - tool_call_id: str - reasoning_content: str | None - - -class ChatCompletionTool(TypedDict, total=False): - """One entry in the request ``tools`` array.""" - - type: Literal["function"] - function: dict[str, Any] - - -class ChatCompletionRequest(TypedDict, total=False): - """Incoming ``POST /v1/chat/completions`` request body shape. - - ``chat_template_kwargs`` is an OpenAI-compatible server extension - used by the gateway to forward per-request chat template overrides - (e.g. ``enable_thinking``) into ``MessageCodec``. - """ - - model: str - messages: list[ChatMessage] - tools: list[ChatCompletionTool] - tool_choice: Any - stream: bool - n: int - response_format: Any - chat_template_kwargs: dict[str, Any] - temperature: float - top_p: float - top_k: int - max_tokens: int - - -class ChatCompletionUsage(TypedDict): - """Token usage block inside the response.""" - - prompt_tokens: int - completion_tokens: int - total_tokens: int - - -class ChatCompletionChoice(TypedDict): - """One entry in the response ``choices`` array. - - ``finish_reason`` is the OpenAI-spec value mapped from the backend's - raw stop reason by ``MessageCodec.decode_response``. - """ - - index: int - message: ChatMessage - finish_reason: Literal["stop", "length", "tool_calls", "content_filter", "function_call"] - - -class ChatCompletionResponse(TypedDict): - """Outgoing ``POST /v1/chat/completions`` response body shape.""" - - id: str - object: Literal["chat.completion"] - created: int - model: str - choices: list[ChatCompletionChoice] - usage: ChatCompletionUsage diff --git a/uni_agent/gateway/session/session.py b/uni_agent/gateway/session/session.py index c0a5f282..60ef01c9 100644 --- a/uni_agent/gateway/session/session.py +++ b/uni_agent/gateway/session/session.py @@ -10,8 +10,8 @@ from fastapi import HTTPException -from uni_agent.gateway.session.codec import MalformedRequestError, MessageCodec -from uni_agent.gateway.session.types import SessionHandle, Trajectory +from uni_agent.gateway.session.codec import MessageCodec +from uni_agent.gateway.session.types import InternalGenerationRequest, SessionHandle, Trajectory class SessionPhase(str, Enum): @@ -86,7 +86,7 @@ class GenerationOutcome: """Business result returned by ``GatewaySession.run_generation``. The session emits this instead of an HTTP response dict. ``_GatewayActor`` - converts it into the OpenAI chat-completion JSON envelope. + passes it to the provider adapter for wire response serialization. Attributes: assistant_msg: Decoded assistant message, or an empty assistant message @@ -107,8 +107,8 @@ class GatewaySession: ``_GatewayActor`` owns instances of this class, calls ``run_generation`` for chat requests, and delegates lifecycle operations here. The session owns the - conversation state and trajectory materialization, while the actor owns HTTP - routing and OpenAI response serialization. + conversation state and trajectory materialization, while the actor owns + HTTP routing and provider response serialization. """ def __init__( @@ -139,19 +139,16 @@ def __init__( # This is an implementation detail, not GatewaySession's public contract. self.generation_lock = asyncio.Lock() - async def run_generation(self, payload: dict[str, Any], backend) -> GenerationOutcome: - """Run one chat-completion request and return its business outcome. + async def run_generation(self, request: InternalGenerationRequest, backend) -> GenerationOutcome: + """Run one provider-normalized generation request and return its business outcome. The backend is passed in for this call only; the session does not own the - backend lifecycle. Protocol capability checks happen in the actor before - this method, while malformed payloads and backend errors are converted - into HTTP exceptions here. + backend lifecycle. The actor/provider adapter has already lowered the + wire payload to the internal canonical request; session never sees raw + wire payloads. Protocol capability checks happen in the actor before + this method, while backend errors are converted into HTTP exceptions + here. """ - try: - request_context = self._codec.normalize_request(payload) - except MalformedRequestError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc - async with self.generation_lock: async with self.request_lock: if self.phase != SessionPhase.ACTIVE: @@ -159,7 +156,7 @@ async def run_generation(self, payload: dict[str, Any], backend) -> GenerationOu status_code=409, detail=f"Session {self.handle.session_id} is {self.phase.value.lower()}", ) - encoded = await self._prepare_generation_inputs(payload, request_context) + encoded = await self._prepare_generation_inputs(request) if encoded.length_exhausted_trajectory is not None: empty_msg = {"role": "assistant", "content": ""} self.trajectories.append(encoded.length_exhausted_trajectory) @@ -222,10 +219,10 @@ async def run_generation(self, payload: dict[str, Any], backend) -> GenerationOu completion_tokens=len(response_ids), ) - async def _prepare_generation_inputs(self, payload: dict[str, Any], request_context: dict[str, Any]) -> EncodedData: - messages = request_context["messages"] - tools = request_context["tools"] - request_chat_template_kwargs = request_context["chat_template_kwargs"] + async def _prepare_generation_inputs(self, request: InternalGenerationRequest) -> EncodedData: + messages = request["messages"] + tools = request["tools"] + request_chat_template_kwargs = request["chat_template_kwargs"] materialized_trajectory = None image_data = None video_data = None @@ -297,7 +294,7 @@ async def _prepare_generation_inputs(self, payload: dict[str, Any], request_cont buffer = TrajectoryBuffer(prompt_ids=prompt_ids) context_ids = buffer.prompt_ids + buffer.response_ids - sampling_params = self._codec.build_sampling_params(payload) + sampling_params = dict(request["sampling_params"]) remaining_response_budget = ( self._response_length - len(buffer.response_mask) if self._response_length is not None else None ) diff --git a/uni_agent/gateway/session/types.py b/uni_agent/gateway/session/types.py index 9d43d7bc..ee430d51 100644 --- a/uni_agent/gateway/session/types.py +++ b/uni_agent/gateway/session/types.py @@ -1,26 +1,40 @@ -"""Shared gateway-owned dataclasses passed across session boundaries.""" +"""Shared gateway-owned types passed across session boundaries.""" from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypedDict if TYPE_CHECKING: import numpy as np import torch +class InternalGenerationRequest(TypedDict): + """Lowered request consumed by GatewaySession.run_generation. + + Provider adapters lower OpenAI / Anthropic wire requests into this + template-facing canonical before the session sees them. It is not a + provider-neutral block model. + """ + + messages: list[dict[str, Any]] + tools: list[dict[str, Any]] | None + chat_template_kwargs: dict[str, Any] + sampling_params: dict[str, Any] + + @dataclass class SessionHandle: """Address returned to agent runners for a newly created gateway session. Attributes: session_id: Stable session identifier assigned by the caller. - base_url: Per-session OpenAI-compatible ``/v1`` API root, or ``None`` + base_url: Per-session provider-compatible ``/v1`` API root, or ``None`` when the handle only needs to identify the session. reward_info_url: Per-session endpoint used by runners to attach reward - metadata; this is a sibling of the OpenAI-compatible ``/v1`` root - rather than part of that API. + metadata; this is a sibling of the provider ``/v1`` root rather + than part of that API. """ session_id: str