diff --git a/docs/content/docs/framework/agent/model.en.mdx b/docs/content/docs/framework/agent/model.en.mdx index 71fa2a115..0652b4a67 100644 --- a/docs/content/docs/framework/agent/model.en.mdx +++ b/docs/content/docs/framework/agent/model.en.mdx @@ -15,6 +15,7 @@ VeADK builds the model client on [LiteLLM](https://github.com/BerriAI/litellm) b | `enable_responses` | `bool` | `False` | Use the Ark Responses API; see [Responses API](/en/docs/framework/agent/responses-api). | | `enable_responses_cache` | `bool` | `True` | Reuse `previous_response_id` for multi-turn continuation and caching; only when `enable_responses=True`. | | `context_cache_config` | `ContextCacheConfig` | `None` | Context-cache config to reduce repeated computation on long contexts. | +| `query_routing` | `bool \| dict \| QueryRoutingConfig` | `None` | Query routing between cheap and strong models. `None` auto-loads the `query_routing` block from `config.yaml`; `False` disables it. See [Query routing](#query-routing). | ## Use a different model @@ -36,6 +37,123 @@ Pass a list to `model_name`: the first model is tried first; on failure the next agent = Agent(model_name=["doubao-seed-1-8-251228", "deepseek-r1-250528"]) ``` +## Query routing + +SystemOne selects a model once at the start of each user turn, using the effective history, instructions, tool declarations and candidate capabilities. The selection remains fixed through tool loops, nested AgentTool calls and sub-agent transfers. The next user turn is classified again. Routing is disabled unless explicitly configured; an unconfigured project retains its original execution behavior. + +Add the block to `config.yaml` in the project root: + +```yaml title="config.yaml" +decision: + model: doubao-seed-2-1-lite-260915 + +query_routing: + simple_models: + - doubao-seed-2-1-lite-260915 + # strong_model: kimi-k3-1-pro-261024 # optional; defaults to the Agent's model_name + mode: shadow # observe first; set enforce to execute routing + decision_timeout: 2.0 + questions: + - question: Does the answer need up-to-date, real-time information? + threshold: 0.2 +``` + +The Agent discovers `config.yaml` in the current working directory and its parents at startup, so no extra code is needed: + +```python title="agent.py" +agent = Agent(model_name="kimi-k3-1-pro-261024") +``` + +You can also pass a mapping directly with `Agent(query_routing={...})`, or disable routing explicitly with `query_routing=False`. + +| Field | Type | Default | Description | +| :--- | :--- | :--- | :--- | +| `simple_models` | `list[str \| dict]` | required | Low-cost models, at least one. The first model that supports the current request is selected in list order. A string is shorthand for a mapping with only `name`. | +| `strong_model` | `str` | Agent's `model_name` | Strong model name. When it differs from the Agent's primary model, it is built separately with the same connection settings. | +| `mode` | `"enforce" \| "shadow"` | `"shadow"` | `shadow` only records routing decisions and always uses the strong model, so you can observe results before switching. | +| `decision_timeout` | `float` | `2.0` | Per-classification timeout in seconds, in `(0, 30]`. A timeout routes to the strong model. | +| `builtin_thresholds` | `dict[str, float]` | `{}` | Optional overrides for built-in question thresholds, each in `[0, 1]`. Calibrate against representative queries before enforcing. | +| `questions` | `list[str \| dict]` | `[]` | Custom questions appended to the four built-in questions. A string is shorthand for `{question: ...}`. | + + +The root `decision` block is shared by query routing and content safety. It enables neither feature. `Agent(decision=DecisionConfig(...))` supplies explicit immutable connection values; each consumer owns its client and questions. Configure connections only here, not inside either feature. + +| Field | Required | Default / rule | +| :--- | :--- | :--- | +| `decision.api_url` | No | Derived from the Agent's Ark endpoint as `/api/v1/systemone`; required for a custom gateway. | +| `decision.api_key` | No | Reuses the Agent's Ark key; required for a non-Ark decision endpoint. | +| `decision.model` | No | `doubao-seed-2-1-lite-260915`; the model used for SystemOne decisions. | + +`DECISION_API_URL`, `DECISION_API_KEY`, and `DECISION_MODEL` override YAML. Overriding a configured destination requires its own `DECISION_API_KEY`. Explicit `DecisionConfig` values bypass file/environment loading. Questions, thresholds, and `decision_timeout` belong to each feature; changing routing policy does not change content-safety policy. + +```python +from veadk import Agent +from veadk.decision import DecisionConfig + +agent = Agent( + name="router", + decision=DecisionConfig(model="doubao-seed-2-1-lite-260915"), + query_routing={"simple_models": ["doubao-seed-2-1-lite-260915"]}, +) +``` + +Each `simple_models` entry can declare capabilities; anything not declared is treated as unsupported: + +| Capability field | Type | Default | Description | +| :--- | :--- | :--- | :--- | +| `name` | `str` | required | Model name; `model` and `model_name` are accepted aliases. | +| `provider` | `str` | Agent provider | LiteLLM provider prefix for unqualified names. | +| `api` | `str` | Agent protocol | `chat_completions` or `responses` (Ark Responses adapter). | +| `api_base` | `str` | Agent endpoint | Independent endpoint. A different endpoint requires an explicit credential. | +| `api_key_env` / `api_key` | `str` | Agent credential | Environment variable name or secret value; mutually exclusive. Prefer `api_key_env`. | +| `extra_config` | `dict` | `{}` | Provider options. Cannot override routing, history or connection fields. Independent endpoints do not inherit the Agent’s provider options. | +| `supports_tools` | `bool` | `False` | Whether the model supports tool calls. | +| `capability_description` | `str \| None` | `None` | Optional, 1–2000 characters. Describe validated task suitability, limitations, and tool dependencies to SystemOne. This does not override instructions, capability flags, or thresholds, and does not guarantee quality. Do not include credentials. | +| `supports_media` | `bool` | `False` | Whether the model supports image, audio, video, or file input. | +| `supports_structured_output` | `bool` | `False` | Whether the model supports structured output. | +| `max_context_chars` | `int` | `32000` | Character budget including instructions, tools, schemas and history; media uses an estimate. This is not a tokenizer-based context-window guarantee. | + +In enforce mode, tool arguments are validated against the tool's declared schema before execution. Unambiguous decimal integer strings can be converted losslessly when the schema explicitly requires an integer; fractions, booleans, and ambiguous formats are not coerced to integers. Invalid arguments return a structured error without original argument values for correction by the model already selected for this turn. At most two correction responses are allowed across the root turn, including nested agents; further invalid calls fail explicitly. Tools that have already executed are not automatically retried. Tools without parameter schemas retain their existing behavior; schema references must be local. Responses replay places explanatory text before the structured tool calls from the same model message, preserving all text, call IDs, arguments and tool results without inserting an assistant message between calls and results. Disabled routing and shadow mode preserve the original tool argument behavior. + +Control questions always use the `noul` type. Four questions are built in: + +| Question ID | Default threshold | What it checks | +| :--- | :--- | :--- | +| `is_complex` | `0.5` | Whether reasoning, planning, code or math exceeds the candidate’s declared capabilities. | +| `needs_tools` | `0.5` | Whether tool orchestration or recovery exceeds the candidate’s declared capabilities. Supported simple tool calls are not automatically escalated. | +| `needs_long_context` | `0.5` | Whether the current task requires global analysis across a long history or multiple documents. Extracting stated facts from a few recent messages is not long-context analysis. | +| `is_high_stakes` | `0.05` | Whether the current task asks for high-risk medical, legal, financial, or safety conclusions, or irreversible actions. A past high-risk topic does not automatically escalate an ordinary factual extraction; follow-up requests for related advice remain high-risk. | + +The default thresholds are `0.5` for complex reasoning, tool orchestration and long context, and `0.05` for high stakes. Explicit `builtin_thresholds` values take precedence; custom questions still default to `0.15` when their threshold is omitted. A threshold bounds the SystemOne question probability, not the task error rate. These defaults apply only when routing is enabled; `shadow` mode still executes with the strong model. + +A custom question only needs `question` and `threshold`; its ID is generated as `custom_1`, `custom_2`, and so on. Built-in IDs are reserved and cannot be overridden. Every question returns a probability from `0` to `1`; **the simple model is used only when all probabilities are at or below their thresholds**. Any probability above its threshold, a missing or invalid answer, a Decision timeout, or a Decision error routes the request to the strong model. Incomplete decision input, unsupported media or unavailable dynamic tool declarations also select the strong model. Effective history is never silently truncated to enable a downgrade. After selection, a provider failure propagates; it never triggers a model switch. Late capability growth fails explicitly. Underlying cross-model fallbacks must be removed when routing is enabled. + +Responses continuation is reused only when the target, endpoint, identity, request configuration and effective history match. Model changes, including A → B → A, rebuild the chain from full business history; expired IDs retry with that history on the same model. Provider-specific reasoning and signatures are not replayed to a different target. `enable_responses_cache=False` always sends full history. Server-side caching still requires the provider’s cache-service entitlement; Responses continuation can also use `store=True` with `extra_body={"caching": {"type": "disabled"}}`. + +Nested AgentTool calls retain separate sessions and inherit one root decision. Native ADK LlmAgent and standard sub-agent workflows are adapted on owned copies; custom native LlmAgent subclasses must implement turn binding explicitly. Resume uses the persisted session decision and requires unchanged routing configuration. Live bidirectional execution and external runtimes are outside this ADK text-turn contract. Concurrent execution of the same session still requires application-level serialization; cross-process simultaneous classification is not a distributed lock. + +Model events expose `veadk_query_routing` metadata with root turn ID, route, model and target digest. The decision event also includes reason, latency and reported input tokens. Credentials and raw query text are excluded. + +```python title="agent.py" +agent = Agent( + model_name="strong-model", + query_routing={ + "mode": "enforce", + "simple_models": [{ + "name": "small-model", + "provider": "openai", + "api": "chat_completions", + "api_base": "https://models.example.com/v1", + "api_key_env": "SMALL_MODEL_API_KEY", + "supports_tools": True, + }], + }, +) +``` + +Real protocol tests are opt-in: set `VEADK_QUERY_ROUTING_LIVE=1`, `MODEL_AGENT_API_KEY` and `MODEL_AGENT_API_BASE`, then run `uv run --extra dev pytest tests/e2e/test_query_routing_live.py`. Model names can be set with `VEADK_QUERY_ROUTING_SIMPLE_MODEL` and `VEADK_QUERY_ROUTING_STRONG_MODEL`. These tests validate transport and turn contracts; production routing quality and cost savings require a representative task benchmark. For separate-process HTTP validation, run `uv run --extra dev pytest tests/e2e/test_query_routing_http_live.py`; it starts an ADK API Server and calls `/run` and `/run_sse` over real TCP. Routing does not guarantee tool-protocol conformance for every model: an undeclared tool name raises an ADK error, and an aborted call without a result can cause the provider to reject later history replay. Applications must handle failed sessions; routing does not silently delete tool history or switch models to conceal errors. + + ## Customize the model client For full control, pass a LiteLLM client directly: diff --git a/docs/content/docs/framework/agent/model.mdx b/docs/content/docs/framework/agent/model.mdx index 7064405b3..7d3a12fc0 100644 --- a/docs/content/docs/framework/agent/model.mdx +++ b/docs/content/docs/framework/agent/model.mdx @@ -15,6 +15,7 @@ VeADK 默认基于 [LiteLLM](https://github.com/BerriAI/litellm) 构建模型客 | `enable_responses` | `bool` | `False` | 启用方舟 Responses API,详见 [Responses API 支持](/cn/docs/framework/agent/responses-api)。 | | `enable_responses_cache` | `bool` | `True` | 在 Responses API 下复用 `previous_response_id` 做多轮续接与缓存,仅在 `enable_responses=True` 时生效。 | | `context_cache_config` | `ContextCacheConfig` | `None` | 上下文缓存配置,用于降低长上下文的重复计算成本。 | +| `query_routing` | `bool \| dict \| QueryRoutingConfig` | `None` | 查询分级与模型路由。`None` 自动读取 `config.yaml` 的 `query_routing` 配置,`False` 显式关闭。详见[查询分级与模型路由](#查询分级与模型路由)。 | ## 使用不同模型 @@ -36,6 +37,123 @@ agent = Agent( agent = Agent(model_name=["doubao-seed-1-8-251228", "deepseek-r1-250528"]) ``` +## 查询分级与模型路由 + +SystemOne 在每个用户 turn 开始时,根据有效历史、指令、工具声明及候选模型能力选择一次模型。工具循环、嵌套 AgentTool 和子 Agent 转交全程保持该选择;下一个用户 turn 重新判断。该能力默认关闭,只有显式配置才生效;不配置时保留原有执行行为。 + +在项目根目录的 `config.yaml` 中新增配置块: + +```yaml title="config.yaml" +decision: + model: doubao-seed-2-1-lite-260915 + +query_routing: + simple_models: + - doubao-seed-2-1-lite-260915 + # strong_model: kimi-k3-1-pro-261024 # 可选;默认使用 Agent 的 model_name + mode: shadow # 先观察;设为 enforce 才实际执行路由 + decision_timeout: 2.0 + questions: + - question: 回答是否需要最新的实时信息? + threshold: 0.2 +``` + +Agent 启动时会自动发现当前工作目录及上级目录中的 `config.yaml`,无需额外代码: + +```python title="agent.py" +agent = Agent(model_name="kimi-k3-1-pro-261024") +``` + +也可以直接传入字典:`Agent(query_routing={...})`;显式关闭则传 `query_routing=False`。 + +| 字段 | 类型 | 默认值 | 说明 | +| :--- | :--- | :--- | :--- | +| `simple_models` | `list[str \| dict]` | 必填 | 低成本模型列表,至少一个。按列出顺序选择第一个满足当前请求能力要求的模型。字符串等价于只声明 `name` 的字典。 | +| `strong_model` | `str` | Agent 的 `model_name` | 强模型名称。与 Agent 的主模型不同时,用同一套连接信息单独构建。 | +| `mode` | `"enforce" \| "shadow"` | `"shadow"` | `shadow` 只记录分级结果,所有请求仍走强模型,便于观察效果后再切换。 | +| `decision_timeout` | `float` | `2.0` | 单次分级超时(秒),范围 `(0, 30]`。超时按强模型处理。 | +| `builtin_thresholds` | `dict[str, float]` | `{}` | 可覆盖内置问题阈值,范围 `[0, 1]`;正式执行前应使用代表性查询校准。 | +| `questions` | `list[str \| dict]` | `[]` | 追加到四个内置问题之后的自定义问题。字符串等价于 `{question: ...}`。 | + + +根级 `decision` 由查询路由和内容安全共用,配置它不会启用任何功能。代码中通过 `Agent(decision=DecisionConfig(...))` 显式传入不可变连接配置;各功能独立创建客户端、管理问题。连接只在公共位置配置,不放在功能内部。 + +| 配置项 | 必填 | 默认值 / 规则 | +| :--- | :--- | :--- | +| `decision.api_url` | 否 | 从 Agent 的方舟地址推导 `/api/v1/systemone`;使用自定义网关时必填。 | +| `decision.api_key` | 否 | 复用 Agent 的方舟密钥;非方舟决策地址必须显式填写。 | +| `decision.model` | 否 | `doubao-seed-2-1-lite-260915`,用于 SystemOne 决策的模型。 | + +`DECISION_API_URL`、`DECISION_API_KEY`、`DECISION_MODEL` 可覆盖 YAML。覆盖已配置的目标地址时,必须同时提供对应的 `DECISION_API_KEY`。显式构造 `DecisionConfig` 不读取文件或环境变量。问题、阈值及 `decision_timeout` 归各功能管理;修改路由规则不会修改内容安全规则。 + +```python +from veadk import Agent +from veadk.decision import DecisionConfig + +agent = Agent( + name="router", + decision=DecisionConfig(model="doubao-seed-2-1-lite-260915"), + query_routing={"simple_models": ["doubao-seed-2-1-lite-260915"]}, +) +``` + +每个 `simple_models` 条目可声明能力;未声明的能力一律视为不支持: + +| 能力字段 | 类型 | 默认值 | 说明 | +| :--- | :--- | :--- | :--- | +| `name` | `str` | 必填 | 模型名称,也可写成 `model` 或 `model_name`。 | +| `provider` | `str` | Agent provider | 无前缀名称使用的 LiteLLM provider。 | +| `api` | `str` | Agent 协议 | `chat_completions` 或 `responses`(方舟 Responses 适配器)。 | +| `api_base` | `str` | Agent endpoint | 独立 endpoint;不同 endpoint 必须显式配置凭据。 | +| `api_key_env` / `api_key` | `str` | Agent 凭据 | 环境变量名或密钥值,二者互斥;推荐 `api_key_env`。 | +| `extra_config` | `dict` | `{}` | Provider 选项,不得覆盖路由、历史及连接字段。独立 endpoint 不继承 Agent 的 provider 选项。 | +| `supports_tools` | `bool` | `False` | 是否支持工具调用。 | +| `capability_description` | `str \| None` | `None` | 可选,1–2000 字符。向 SystemOne 说明候选模型经验证的适用任务、限制及工具依赖;不是提示词覆盖,不替代能力标记或阈值,也不自动保证质量。不得包含凭据。 | +| `supports_media` | `bool` | `False` | 是否支持图片、音频、视频或文件输入。 | +| `supports_structured_output` | `bool` | `False` | 是否支持结构化输出。 | +| `max_context_chars` | `int` | `32000` | 包含指令、工具、schema 和历史的字符预算;媒体按固定成本估算,不等同于 tokenizer 精确的上下文窗口保证。 | + +在 enforce 模式下,工具执行前会按工具声明的 schema 校验参数。明确声明为整数的十进制整数字符串可无损转换;小数、布尔值及歧义格式不会被强制转成整数。不合法的参数返回不包含原始参数值的结构化错误,由本 turn 已选择的模型纠正,整个根 turn(含嵌套 Agent)最多返回两次纠正反馈,随后明确报错。已经执行的工具不会因此被自动重试。工具没有参数 schema 时保持原有行为;schema 的引用必须为本地引用。Responses 重放会将同一条模型消息中的说明文本置于结构化工具调用之前,保留全部文本、调用 ID、参数和工具结果,避免在工具调用与结果之间插入模型消息。未启用路由以及 shadow 模式均保留原有工具参数行为。 + +内部控制问题固定为 `noul` 类型,内置四个: + +| 问题 ID | 默认阈值 | 判断内容 | +| :--- | :--- | :--- | +| `is_complex` | `0.5` | 推理、规划、代码或数学是否超出候选模型已声明的能力。 | +| `needs_tools` | `0.5` | 工具编排或恢复是否超出候选模型已声明的能力;已支持的简单工具调用不会因此自动升级。 | +| `needs_long_context` | `0.5` | 当前任务是否需要跨较长历史或多份材料做全局分析;从少量最近消息提取已给出的事实不属于长上下文分析。 | +| `is_high_stakes` | `0.05` | 当前任务是否要求医疗、法律、金融、安全等高风险结论或不可撤销操作;历史出现高风险主题不会自动将当前的普通事实提取升级,但继续追问相关建议仍属于高风险。 | + +默认阈值为:复杂推理、工具编排、长上下文 `0.5`,高风险 `0.05`。`builtin_thresholds` 中显式填写的值优先;自定义问题省略阈值时仍使用 `0.15`。阈值是 SystemOne 问题概率的上限,不是任务错误率。未启用路由时不生效,`shadow` 模式仍由强模型执行。 + +自定义问题只需提供 `question` 与 `threshold`,ID 自动生成为 `custom_1`、`custom_2`……;内置 ID 保留,不能覆盖。所有问题都返回 `0` 到 `1` 的概率,**只有全部不超过各自阈值时才走简单模型**;任一问题超过阈值、返回值缺失或非法、Decision 超时或报错时,都走强模型。决策输入不完整、媒体无法表示或动态工具声明不可用时也选择强模型,不会静默截断有效历史以促成降级。选择确定后,provider 失败直接传播,不切换模型;后续请求超出候选能力时显式报错。开启路由时必须移除底层跨模型 fallback。 + +只有目标、endpoint、身份、请求配置及有效历史一致时才复用 Responses 续链。切换模型(包括 A → B → A)会使用完整业务历史重建;ID 过期时在同一模型上携带完整历史重试。不同目标间不重放 provider 专属的推理内容及签名。`enable_responses_cache=False` 始终发送完整历史。服务端缓存仍需对应模型的缓存服务权限;Responses 续接也可使用 `store=True` 与 `extra_body={"caching": {"type": "disabled"}}`。 + +嵌套 AgentTool 保持独立 session,并继承同一个根决策。原生 ADK LlmAgent 及标准子 Agent 工作流通过独立副本适配;自定义原生 LlmAgent 子类需显式实现 turn 绑定。恢复执行使用 session 中持久化的决策,要求路由配置不变。双向 Live 执行及外部 runtime 不在此 ADK 文本 turn 契约内。同一 session 的并发执行仍需应用侧串行化;跨进程同时分级不提供分布式锁。 + +模型事件提供 `veadk_query_routing` 元数据,包含根 turn ID、路由、模型及目标摘要。决策事件还包含原因、耗时及服务上报的输入 token 数,不包含凭据或原始 query。 + +```python title="agent.py" +agent = Agent( + model_name="strong-model", + query_routing={ + "mode": "enforce", + "simple_models": [{ + "name": "small-model", + "provider": "openai", + "api": "chat_completions", + "api_base": "https://models.example.com/v1", + "api_key_env": "SMALL_MODEL_API_KEY", + "supports_tools": True, + }], + }, +) +``` + +真实协议测试需显式开启:设置 `VEADK_QUERY_ROUTING_LIVE=1`、`MODEL_AGENT_API_KEY` 和 `MODEL_AGENT_API_BASE`,执行 `uv run --extra dev pytest tests/e2e/test_query_routing_live.py`。可通过 `VEADK_QUERY_ROUTING_SIMPLE_MODEL` 和 `VEADK_QUERY_ROUTING_STRONG_MODEL` 指定模型。测试验证协议与 turn 契约;生产路由质量及成本收益仍需代表性任务集评测。独立进程 HTTP 验收可执行 `uv run --extra dev pytest tests/e2e/test_query_routing_http_live.py`,测试会启动 ADK API Server 并通过真实 TCP 调用 `/run`、`/run_sse`。开启路由不保证任意模型都满足工具协议;例如输出不存在的工具名会按 ADK 契约报错,中止后未获得结果的工具调用还可能使后续历史重放被 provider 拒绝。应用需处理失败会话,路由不会静默删除工具历史或切换模型来掩盖错误。 + + ## 自定义模型客户端 如需高度自定义,直接传入一个 LiteLLM 客户端: diff --git a/docs/content/docs/framework/tools/guardrail.en.mdx b/docs/content/docs/framework/tools/guardrail.en.mdx index 9ad443157..a76d9b073 100644 --- a/docs/content/docs/framework/tools/guardrail.en.mdx +++ b/docs/content/docs/framework/tools/guardrail.en.mdx @@ -92,18 +92,19 @@ print(response) Configure it in `config.yaml`; agents run through `Runner` or deployed through AgentKit mount the plugin automatically: ```yaml title="config.yaml" +decision: + # Optional; derived from the current Ark API base when omitted + api_url: https://ark.cn-beijing.volces.com/api/v1/systemone + # Optional; reuses the Ark API key when omitted + api_key: + # Optional; defaults to doubao-seed-2-1-lite-260915 + model: doubao-seed-2-1-lite-260915 + tool: content_safety: provider: decision on_error: allow # allow (default, marked degraded) or block block_message: "Your request was blocked by the content safety policy." - decision: - # Optional; derived from the current Ark API base when omitted - api_url: https://ark.cn-beijing.volces.com/api/v1/systemone - # Optional; reuses the Ark API key when omitted - api_key: - # Optional; defaults to doubao-seed-2-1-lite-260915 - model: doubao-seed-2-1-lite-260915 use_default_questions: true questions: # Override a built-in question by name, or append a new one @@ -112,7 +113,9 @@ tool: threshold: 0.5 ``` -The connection fields can also be overridden with `TOOL_CONTENT_SAFETY_DECISION_API_URL`, `TOOL_CONTENT_SAFETY_DECISION_API_KEY`, and `TOOL_CONTENT_SAFETY_DECISION_MODEL`. +The root `decision` block is shared connection configuration; it does not enable either feature or share mutable clients. Content-safety questions and thresholds remain local. Import `DecisionConfig` from `veadk.decision` and pass it through `Agent(decision=...)` or `build_content_safety(config, decision=...)`. See [model configuration](../agent/model) for all connection fields. + +The connection fields can also be overridden with `DECISION_API_URL`, `DECISION_API_KEY`, and `DECISION_MODEL`. When a custom `decision.api_url` is configured, `decision.api_key` must also be supplied explicitly; the SDK never reuses Ark credentials for a third-party endpoint. You can configure it in code instead: diff --git a/docs/content/docs/framework/tools/guardrail.mdx b/docs/content/docs/framework/tools/guardrail.mdx index ade631f49..1de0fe7f0 100644 --- a/docs/content/docs/framework/tools/guardrail.mdx +++ b/docs/content/docs/framework/tools/guardrail.mdx @@ -92,18 +92,19 @@ print(response) 在 `config.yaml` 中配置后,通过 `Runner` 运行或部署 Agent 时会自动挂载插件: ```yaml title="config.yaml" +decision: + # 可选;省略时由当前 Ark API Base 推导 /api/v1/systemone + api_url: https://ark.cn-beijing.volces.com/api/v1/systemone + # 可选;省略时复用 Ark API Key + api_key: + # 可选;默认 doubao-seed-2-1-lite-260915 + model: doubao-seed-2-1-lite-260915 + tool: content_safety: provider: decision on_error: allow # allow(默认,标记 degraded)或 block block_message: "Your request was blocked by the content safety policy." - decision: - # 可选;省略时由当前 Ark API Base 推导 /api/v1/systemone - api_url: https://ark.cn-beijing.volces.com/api/v1/systemone - # 可选;省略时复用 Ark API Key - api_key: - # 可选;默认 doubao-seed-2-1-lite-260915 - model: doubao-seed-2-1-lite-260915 use_default_questions: true questions: # 覆盖同名内置问题,或追加新问题 @@ -112,7 +113,9 @@ tool: threshold: 0.5 ``` -连接字段也可以用环境变量覆盖 YAML:`TOOL_CONTENT_SAFETY_DECISION_API_URL`、`TOOL_CONTENT_SAFETY_DECISION_API_KEY`、`TOOL_CONTENT_SAFETY_DECISION_MODEL`。 +根级 `decision` 是公共连接配置,本身不启用功能,也不共享可变客户端;内容安全问题和阈值仍独立管理。从 `veadk.decision` 导入 `DecisionConfig`,通过 `Agent(decision=...)` 或 `build_content_safety(config, decision=...)` 传入。完整连接项见[模型配置](../agent/model)。 + +连接字段也可以用环境变量覆盖 YAML:`DECISION_API_URL`、`DECISION_API_KEY`、`DECISION_MODEL`。 如果配置了自定义 `decision.api_url`,必须同时显式提供 `decision.api_key`;SDK 不会把 Ark 凭证复用给第三方地址。也可以在代码中直接配置: diff --git a/pytest.ini b/pytest.ini index 964d55142..154508519 100644 --- a/pytest.ini +++ b/pytest.ini @@ -9,6 +9,7 @@ filterwarnings = # installed plugin version. asyncio_mode = strict markers = + query_routing_live: real SystemOne and model API routing E2E (explicit opt-in) piagent_smoke: real Pi binary/model smoke test (opt in with PIAGENT_RUN_SMOKE=1) codex_smoke: real Codex binary/sandbox/socket smoke test, stubbed model backend (opt in with CODEX_RUN_SMOKE=1); binds real ports and spawns a subprocess, so it must not run under `pytest -n` codex_probe: real-model Codex canary that spends ~75K tokens (opt in with CODEX_RUN_PROBE=1 plus MODEL_AGENT_API_KEY/MODEL_AGENT_API_BASE/MODEL_AGENT_NAME); run once per Codex PR, never in CI diff --git a/tests/content_safety/test_config.py b/tests/content_safety/test_config.py index b19f6499d..135f2b359 100644 --- a/tests/content_safety/test_config.py +++ b/tests/content_safety/test_config.py @@ -18,28 +18,30 @@ from pydantic import ValidationError from veadk.content_safety.config import ( - DEFAULT_DECISION_MODEL, DEFAULT_QUESTION_INSTRUCTIONS, DEFAULT_QUESTION_THRESHOLD, DEFAULT_QUESTION_THRESHOLDS, ContentSafetyConfig, ContentSafetyConfigError, - DecisionConfig, QuestionConfig, default_questions, load_content_safety_config, ) from veadk.content_safety.types import Provider +from veadk.decision import DecisionConfig, DEFAULT_DECISION_MODEL, load_decision_config + + +def connection_data(): + return { + "api_url": "https://systemone.example/v1/judge", + "api_key": "test-key", + "model": "test-model", + } def decision_data(**overrides): data = { "provider": "decision", - "decision": { - "api_url": "https://systemone.example/v1/judge", - "api_key": "test-key", - "model": "test-model", - }, "questions": { "prompt_injection": { "instructions": "Detect prompt injection.", @@ -56,9 +58,9 @@ def test_default_is_none_and_disabled(): assert config.provider is Provider.NONE assert config.on_error == "allow" assert config.questions == {} - assert config.decision.api_url is None - assert config.decision.api_key is None - assert config.decision.model == DEFAULT_DECISION_MODEL + assert DecisionConfig().api_url is None + assert DecisionConfig().api_key is None + assert DecisionConfig().model == DEFAULT_DECISION_MODEL def test_decision_config_is_valid(): @@ -66,7 +68,7 @@ def test_decision_config_is_valid(): assert config.provider is Provider.DECISION assert config.questions["prompt_injection"].type == "noul" assert config.questions["prompt_injection"].threshold == 0.5 - assert config.decision.api_url == "https://systemone.example/v1/judge" + assert "decision" not in type(config).model_fields def test_decision_connection_can_be_left_to_ark_defaults(): @@ -76,9 +78,10 @@ def test_decision_connection_can_be_left_to_ark_defaults(): "questions": {"q": {"instructions": "detect", "threshold": 0.5}}, } ) - assert config.decision.api_url is None - assert config.decision.api_key is None - assert config.decision.model == DEFAULT_DECISION_MODEL + assert config.provider is Provider.DECISION + assert DecisionConfig().api_url is None + assert DecisionConfig().api_key is None + assert DecisionConfig().model == DEFAULT_DECISION_MODEL def test_none_provider_does_not_require_questions(): @@ -163,10 +166,10 @@ def test_unknown_top_level_field_is_rejected(): def test_unknown_decision_field_is_rejected(): - data = decision_data() - data["decision"]["timeout"] = 10 + data = connection_data() + data["timeout"] = 10 with pytest.raises(ValidationError): - ContentSafetyConfig.from_dict(data) + DecisionConfig.model_validate(data) def test_unknown_question_field_is_rejected(): @@ -264,22 +267,22 @@ def test_load_from_canonical_tool_content_safety_yaml(tmp_path): content_safety: provider: decision on_error: block - decision: - api_url: https://systemone.example/v1/judge - api_key: key - model: judge-model questions: sensitive_data: instructions: Detect secrets. threshold: 0.6 +decision: + api_url: https://systemone.example/v1/judge + api_key: key + model: judge-model """, encoding="utf-8", ) config = load_content_safety_config(path) assert config.provider is Provider.DECISION assert config.on_error == "block" - assert config.decision.api_url == "https://systemone.example/v1/judge" - assert config.decision.model == "judge-model" + assert load_decision_config(path).api_url == "https://systemone.example/v1/judge" + assert load_decision_config(path).model == "judge-model" assert config.questions["sensitive_data"].threshold == 0.6 @@ -302,26 +305,24 @@ def test_env_overrides_decision_fields(tmp_path, monkeypatch): tool: content_safety: provider: decision - decision: - api_url: https://yaml.example/v1/judge - api_key: yaml-key - model: yaml-model questions: q: instructions: detect threshold: 0.5 +decision: + api_url: https://yaml.example/v1/judge + api_key: yaml-key + model: yaml-model """, encoding="utf-8", ) - monkeypatch.setenv( - "TOOL_CONTENT_SAFETY_DECISION_API_URL", "https://env.example/v1/judge" - ) - monkeypatch.setenv("TOOL_CONTENT_SAFETY_DECISION_API_KEY", "env-key") - monkeypatch.setenv("TOOL_CONTENT_SAFETY_DECISION_MODEL", "env-model") - config = load_content_safety_config(path) - assert config.decision.api_url == "https://env.example/v1/judge" - assert config.decision.api_key_value == "env-key" - assert config.decision.model == "env-model" + monkeypatch.setenv("DECISION_API_URL", "https://env.example/v1/judge") + monkeypatch.setenv("DECISION_API_KEY", "env-key") + monkeypatch.setenv("DECISION_MODEL", "env-model") + config = load_decision_config(path) + assert config.api_url == "https://env.example/v1/judge" + assert config.api_key_value == "env-key" + assert config.model == "env-model" def test_partial_env_override_keeps_yaml_values(tmp_path, monkeypatch): @@ -331,32 +332,30 @@ def test_partial_env_override_keeps_yaml_values(tmp_path, monkeypatch): tool: content_safety: provider: decision - decision: - api_url: https://yaml.example/v1/judge - api_key: yaml-key - model: yaml-model questions: q: instructions: detect threshold: 0.5 +decision: + api_url: https://yaml.example/v1/judge + api_key: yaml-key + model: yaml-model """, encoding="utf-8", ) - monkeypatch.setenv("TOOL_CONTENT_SAFETY_DECISION_MODEL", "env-model") - monkeypatch.delenv("TOOL_CONTENT_SAFETY_DECISION_API_URL", raising=False) - monkeypatch.delenv("TOOL_CONTENT_SAFETY_DECISION_API_KEY", raising=False) - config = load_content_safety_config(path) - assert config.decision.api_url == "https://yaml.example/v1/judge" - assert config.decision.api_key_value == "yaml-key" - assert config.decision.model == "env-model" + monkeypatch.setenv("DECISION_MODEL", "env-model") + monkeypatch.delenv("DECISION_API_URL", raising=False) + monkeypatch.delenv("DECISION_API_KEY", raising=False) + config = load_decision_config(path) + assert config.api_url == "https://yaml.example/v1/judge" + assert config.api_key_value == "yaml-key" + assert config.model == "env-model" def test_from_dict_does_not_apply_env(monkeypatch): - monkeypatch.setenv( - "TOOL_CONTENT_SAFETY_DECISION_API_URL", "https://env.example/v1/judge" - ) - config = ContentSafetyConfig.from_dict(decision_data()) - assert config.decision.api_url == "https://systemone.example/v1/judge" + monkeypatch.setenv("DECISION_API_URL", "https://env.example/v1/judge") + config = DecisionConfig.model_validate(connection_data()) + assert config.api_url == "https://systemone.example/v1/judge" def test_env_override_without_yaml_decision_block(tmp_path, monkeypatch): @@ -373,10 +372,10 @@ def test_env_override_without_yaml_decision_block(tmp_path, monkeypatch): """, encoding="utf-8", ) - monkeypatch.setenv("TOOL_CONTENT_SAFETY_DECISION_API_KEY", "env-key") - config = load_content_safety_config(path) - assert config.decision.api_key_value == "env-key" - assert config.decision.api_url is None + monkeypatch.setenv("DECISION_API_KEY", "env-key") + config = load_decision_config(path) + assert config.api_key_value == "env-key" + assert config.api_url is None def test_load_rejects_duplicate_yaml_keys(tmp_path): @@ -441,7 +440,7 @@ def test_env_override_ignores_non_mapping_decision(tmp_path, monkeypatch): "tool:\n content_safety:\n provider: none\n decision: nope\n", encoding="utf-8", ) - monkeypatch.setenv("TOOL_CONTENT_SAFETY_DECISION_API_KEY", "env-key") + monkeypatch.setenv("DECISION_API_KEY", "env-key") with pytest.raises(ContentSafetyConfigError, match="invalid `content_safety`"): load_content_safety_config(path) @@ -588,18 +587,18 @@ def test_threshold_must_be_numeric_not_string(): def test_api_key_is_not_exposed_in_repr_or_dump(): - config = ContentSafetyConfig.from_dict(decision_data()) - assert config.decision.api_key_value == "test-key" + config = DecisionConfig.model_validate(connection_data()) + assert config.api_key_value == "test-key" assert "test-key" not in repr(config) - assert "test-key" not in repr(config.decision) + assert "test-key" not in config.model_dump_json() assert "test-key" not in repr(config.model_dump()) def test_http_api_url_is_rejected_outside_loopback(): - data = decision_data() - data["decision"]["api_url"] = "http://systemone.example/v1/judge" + data = connection_data() + data["api_url"] = "http://systemone.example/v1/judge" with pytest.raises(ValidationError, match="https"): - ContentSafetyConfig.from_dict(data) + DecisionConfig.model_validate(data) def test_http_api_url_is_allowed_for_loopback(): @@ -608,9 +607,9 @@ def test_http_api_url_is_allowed_for_loopback(): "http://localhost:8080/v1/judge", "http://[::1]:8080/v1/judge", ): - data = decision_data() - data["decision"]["api_url"] = url - assert ContentSafetyConfig.from_dict(data).decision.api_url == url + data = connection_data() + data["api_url"] = url + assert DecisionConfig.model_validate(data).api_url == url def test_load_applies_dotenv_decision_fields(tmp_path, monkeypatch): @@ -622,34 +621,32 @@ def test_load_applies_dotenv_decision_fields(tmp_path, monkeypatch): tool: content_safety: provider: decision - decision: - api_url: https://yaml.example/v1/judge - api_key: yaml-key questions: q: instructions: detect threshold: 0.5 +decision: + api_url: https://yaml.example/v1/judge + api_key: yaml-key """, encoding="utf-8", ) - (tmp_path / ".env").write_text( - "TOOL_CONTENT_SAFETY_DECISION_API_KEY=env-key\n", encoding="utf-8" - ) + (tmp_path / ".env").write_text("DECISION_API_KEY=env-key\n", encoding="utf-8") monkeypatch.chdir(tmp_path) - monkeypatch.delenv("TOOL_CONTENT_SAFETY_DECISION_API_KEY", raising=False) + monkeypatch.delenv("DECISION_API_KEY", raising=False) try: - config = load_content_safety_config(path) - assert config.decision.api_key_value == "env-key" + config = load_decision_config(path) + assert config.api_key_value == "env-key" finally: - os.environ.pop("TOOL_CONTENT_SAFETY_DECISION_API_KEY", None) + os.environ.pop("DECISION_API_KEY", None) @pytest.mark.parametrize("url", ["https://[::1", "https://"]) def test_api_url_must_be_parseable_with_host(url): - data = decision_data() - data["decision"]["api_url"] = url + data = connection_data() + data["api_url"] = url with pytest.raises(ValidationError, match="decision.api_url"): - ContentSafetyConfig.from_dict(data) + DecisionConfig.model_validate(data) def test_llm_shield_app_id_warns_for_non_llm_shield_provider(caplog): diff --git a/tests/content_safety/test_content_safety_plugin.py b/tests/content_safety/test_content_safety_plugin.py index e1af29bbb..67a546935 100644 --- a/tests/content_safety/test_content_safety_plugin.py +++ b/tests/content_safety/test_content_safety_plugin.py @@ -83,10 +83,6 @@ def config(*, on_error="allow", questions=None): "threshold": 0.5, } }, - "decision": { - "api_url": "https://systemone.example", - "api_key": "k", - }, } ) diff --git a/tests/content_safety/test_factory.py b/tests/content_safety/test_factory.py index 65a2bd4e2..e9d32c745 100644 --- a/tests/content_safety/test_factory.py +++ b/tests/content_safety/test_factory.py @@ -21,13 +21,16 @@ import pytest +from veadk.decision import ( + DecisionConfig, + DecisionConfigError, + resolve_decision_connection, +) from veadk.content_safety import ( ContentSafetyConfig, - ContentSafetyConfigError, build_content_safety, ) -from veadk.content_safety.factory import ( - _ark_defaults, +from veadk.decision.connection import ( _is_ark_api_base, _systemone_url, ) @@ -68,11 +71,6 @@ def decision_config(): { "provider": "decision", "questions": {"q": {"instructions": "detect", "threshold": 0.5}}, - "decision": { - "api_url": "https://systemone.example/v1/judge", - "api_key": "key", - "model": "judge-model", - }, } ) @@ -89,7 +87,14 @@ def test_decision_provider_uses_injected_client(): def test_decision_provider_builds_default_client(): - plugin = build_content_safety(decision_config()) + plugin = build_content_safety( + decision_config(), + decision=DecisionConfig( + api_url="https://systemone.example/v1/judge", + api_key="key", + model="judge-model", + ), + ) assert isinstance(plugin, DecisionPlugin) assert isinstance(plugin._client, DecisionClient) assert plugin._client.endpoint == "https://systemone.example/v1/judge" @@ -259,14 +264,13 @@ def test_systemone_url_derivation(base, expected): assert _systemone_url(base) == expected -def test_ark_defaults_uses_explicit_values(): - assert _ark_defaults("https://ark.example", "key") == ( - "https://ark.example", - "key", - ) +def test_connection_uses_explicit_values(): + assert resolve_decision_connection( + DecisionConfig(), "https://ark.cn-beijing.volces.com/api/v3", "key" + ) == ("https://ark.cn-beijing.volces.com/api/v1/systemone", "key") -def test_ark_defaults_reads_veadk_settings(monkeypatch): +def test_connection_reads_only_missing_veadk_settings(monkeypatch): import veadk.config as veadk_config monkeypatch.setattr( @@ -274,22 +278,20 @@ def test_ark_defaults_reads_veadk_settings(monkeypatch): "settings", SimpleNamespace( model=SimpleNamespace( - api_base="https://ark.example/api/v3", api_key="ark-key" + api_base="https://ark.cn-beijing.volces.com/api/v3", api_key="ark-key" ) ), ) - assert _ark_defaults(None, None) == ( - "https://ark.example/api/v3", + assert resolve_decision_connection(DecisionConfig()) == ( + "https://ark.cn-beijing.volces.com/api/v1/systemone", "ark-key", ) - assert _ark_defaults("https://override", None) == ( - "https://override", - "ark-key", - ) - assert _ark_defaults(None, "override-key") == ( - "https://ark.example/api/v3", - "override-key", - ) + assert resolve_decision_connection( + DecisionConfig(), "https://ark.ap-southeast.bytepluses.com/api/v3" + ) == ("https://ark.ap-southeast.bytepluses.com/api/v1/systemone", "ark-key") + assert resolve_decision_connection( + DecisionConfig(), ark_api_key="override-key" + ) == ("https://ark.cn-beijing.volces.com/api/v1/systemone", "override-key") def test_is_ark_api_base_accepts_ark_hosts_and_rejects_unsafe_urls(): @@ -352,7 +354,7 @@ def test_decision_provider_rejects_non_ark_fallback_base(): "questions": {"q": {"instructions": "detect", "threshold": 0.5}}, } ) - with pytest.raises(ContentSafetyConfigError, match="decision.api_url"): + with pytest.raises(DecisionConfigError, match="decision.api_url"): build_content_safety( config, ark_api_base="https://api.openai.com/v1", @@ -376,13 +378,15 @@ def test_decision_provider_uses_ark_api_url_but_resolves_missing_key(monkeypatch config = ContentSafetyConfig.from_dict( { "provider": "decision", - "decision": { - "api_url": "https://ark.cn-beijing.volces.com/api/v1/systemone" - }, "questions": {"q": {"instructions": "detect", "threshold": 0.5}}, } ) - plugin = build_content_safety(config) + plugin = build_content_safety( + config, + decision=DecisionConfig.model_validate( + {"api_url": "https://ark.cn-beijing.volces.com/api/v1/systemone"} + ), + ) assert ( plugin._client.endpoint == "https://ark.cn-beijing.volces.com/api/v1/systemone" ) @@ -405,29 +409,33 @@ def test_decision_provider_rejects_custom_api_url_without_key(monkeypatch): config = ContentSafetyConfig.from_dict( { "provider": "decision", - "decision": {"api_url": "https://custom.example/v1/systemone"}, "questions": {"q": {"instructions": "detect", "threshold": 0.5}}, } ) - with pytest.raises(ContentSafetyConfigError, match="decision.api_key"): - build_content_safety(config) + with pytest.raises(DecisionConfigError, match="decision.api_key"): + build_content_safety( + config, + decision=DecisionConfig.model_validate( + {"api_url": "https://custom.example/v1/systemone"} + ), + ) def test_decision_provider_rejects_ark_url_with_non_ark_fallback(): config = ContentSafetyConfig.from_dict( { "provider": "decision", - "decision": { - "api_url": ("https://ark.cn-beijing.volces.com/api/v1/systemone") - }, "questions": {"q": {"instructions": "detect", "threshold": 0.5}}, } ) - with pytest.raises(ContentSafetyConfigError, match="non-Ark"): + with pytest.raises(DecisionConfigError, match="not an Ark"): build_content_safety( config, ark_api_base="https://api.openai.com/v1", ark_api_key="openai-key", + decision=DecisionConfig.model_validate( + {"api_url": ("https://ark.cn-beijing.volces.com/api/v1/systemone")} + ), ) diff --git a/tests/content_safety/test_integration.py b/tests/content_safety/test_integration.py index 4f1dd9f9f..748244508 100644 --- a/tests/content_safety/test_integration.py +++ b/tests/content_safety/test_integration.py @@ -19,6 +19,7 @@ from google.adk.plugins import BasePlugin from veadk import Agent, Runner +from veadk.decision import DecisionConfig, load_decision_config from veadk.content_safety import ( DEFAULT_QUESTION_INSTRUCTIONS, ContentSafetyConfig, @@ -33,13 +34,13 @@ tool: content_safety: provider: decision - decision: - api_url: https://systemone.example/v1/judge - api_key: test-key questions: prompt_injection: instructions: test instructions threshold: 0.5 +decision: + api_url: https://systemone.example/v1/judge + api_key: test-key """ @@ -82,7 +83,7 @@ def test_loads_configured_block(self, decision_cwd): config = load_content_safety_config_if_present() assert config is not None assert config.provider is Provider.DECISION - assert config.decision.api_url == "https://systemone.example/v1/judge" + assert load_decision_config().api_url == "https://systemone.example/v1/judge" def test_omitted_questions_use_defaults(self, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) @@ -92,9 +93,9 @@ def test_omitted_questions_use_defaults(self, tmp_path, monkeypatch): tool: content_safety: provider: decision - decision: - api_url: https://systemone.example/v1/judge - api_key: test-key +decision: + api_url: https://systemone.example/v1/judge + api_key: test-key """, ) plugin = resolve_content_safety_plugin() @@ -160,10 +161,6 @@ def test_agent_mapping_overrides_yaml(self, tmp_path, monkeypatch): config = ContentSafetyConfig.from_dict( { "provider": "decision", - "decision": { - "api_url": "https://inline.example/v1/judge", - "api_key": "inline-key", - }, "questions": { "prompt_injection": { "instructions": "inline instructions", @@ -172,7 +169,15 @@ def test_agent_mapping_overrides_yaml(self, tmp_path, monkeypatch): }, } ) - agent = make_agent(content_safety=config) + agent = make_agent( + content_safety=config, + decision=DecisionConfig.model_validate( + { + "api_url": "https://inline.example/v1/judge", + "api_key": "inline-key", + } + ), + ) plugin = resolve_content_safety_plugin(agent) assert plugin is not None assert plugin.name == "content_safety_decision" diff --git a/tests/decision/test_shared_config.py b/tests/decision/test_shared_config.py new file mode 100644 index 000000000..946f3e695 --- /dev/null +++ b/tests/decision/test_shared_config.py @@ -0,0 +1,220 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Common connection configuration stays independent of consumer policy.""" + +import os + +import pytest +import yaml +from pydantic import ValidationError + +from veadk.configs.routing_configs import ( + QueryRoutingConfig, + load_query_routing_config_if_present, +) +from veadk.content_safety import ( + ContentSafetyConfig, + load_content_safety_config_if_present, +) +from veadk.decision import DecisionConfig, DecisionConfigError, load_decision_config + + +def config_file(tmp_path, **overrides): + document = { + "decision": { + "model": "shared-judge", + "api_url": "https://shared.example/systemone", + "api_key": "test-placeholder", + }, + **overrides, + } + path = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(document)) + return path + + +def test_only_common_import_is_public(): + import veadk.content_safety as safety + import veadk.content_safety.config as config + + assert not hasattr(safety, "DecisionConfig") + assert not hasattr(config, "DecisionConfig") + + +@pytest.mark.parametrize( + "consumer,kwargs", + [ + (QueryRoutingConfig, {"simple_models": ["candidate"]}), + (ContentSafetyConfig, {"provider": "decision"}), + ], +) +def test_old_nested_connection_is_rejected(consumer, kwargs): + with pytest.raises(ValidationError, match="decision"): + consumer(**kwargs, decision={"model": "old"}) + + +def test_root_connection_is_immutable_and_enables_nothing(tmp_path): + path = config_file(tmp_path) + first = load_decision_config(path) + second = load_decision_config(path) + assert first.model == second.model == "shared-judge" + assert first is not second + with pytest.raises(ValidationError): + first.model = "changed" + assert load_query_routing_config_if_present(path) is None + assert load_content_safety_config_if_present(path) is None + + +def test_explicit_code_ignores_file_and_environment(tmp_path, monkeypatch): + config_file(tmp_path) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DECISION_MODEL", "env") + assert DecisionConfig(model="code").model == "code" + + +def test_common_environment_overrides_file(tmp_path, monkeypatch): + path = config_file(tmp_path) + monkeypatch.setenv("DECISION_MODEL", "environment-judge") + monkeypatch.setenv("TOOL_CONTENT_SAFETY_DECISION_MODEL", "old-ignored") + assert load_decision_config(path).model == "environment-judge" + + +@pytest.mark.parametrize("value", [[], "invalid", {"unsupported": True}]) +def test_invalid_root_config_fails(tmp_path, value): + with pytest.raises(DecisionConfigError, match="decision"): + load_decision_config(config_file(tmp_path, decision=value)) + + +def test_environment_endpoint_cannot_inherit_file_key(tmp_path, monkeypatch): + path = config_file(tmp_path) + monkeypatch.setenv("DECISION_API_URL", "https://other.example/systemone") + with pytest.raises(DecisionConfigError, match="DECISION_API_KEY"): + load_decision_config(path) + monkeypatch.setenv("DECISION_API_KEY", "other-placeholder") + assert load_decision_config(path).api_key_value == "other-placeholder" + + +def test_structured_config_does_not_become_process_environment(tmp_path, monkeypatch): + from veadk.utils.misc import set_envs + + path = config_file( + tmp_path, + query_routing={"simple_models": ["candidate"]}, + tool={"content_safety": {"provider": "decision"}}, + ) + monkeypatch.setattr(os, "environ", os.environ.copy()) + before = dict(os.environ) + document, exported = set_envs(str(path)) + assert document["decision"]["model"] == "shared-judge" + assert not exported + assert dict(os.environ) == before + + +def test_agent_shares_values_but_never_client_instances(tmp_path, monkeypatch): + from veadk import Agent + from veadk.content_safety import resolve_content_safety_plugin + + config_file(tmp_path) + monkeypatch.chdir(tmp_path) + shared = DecisionConfig( + api_url="https://inline.example/systemone", + api_key="inline-placeholder", + model="inline-judge", + ) + agent = Agent( + name="shared", + model_name="strong", + model_api_key="model-placeholder", + decision=shared, + query_routing={"simple_models": ["candidate"], "decision_timeout": 3}, + content_safety={"provider": "decision"}, + ) + safety = resolve_content_safety_plugin(agent) + routing = agent.model.classifier._judge + assert routing is not safety._client + assert routing._model == safety._client._model == "inline-judge" + assert ( + routing.endpoint + == safety._client.endpoint + == "https://inline.example/systemone" + ) + assert agent.decision == shared + assert routing._timeout != safety._client._timeout + + +@pytest.mark.parametrize( + "text", + [ + "[]", + "decision: []", + "decision: {model: a, model: b}", + "decision: [", + "decision: {1: bad}", + ], +) +def test_malformed_common_yaml_fails(tmp_path, text): + path = tmp_path / "config.yaml" + path.write_text(text) + with pytest.raises(DecisionConfigError): + load_decision_config(path) + + +def test_missing_explicit_file_fails_but_missing_discovered_file_defaults( + tmp_path, monkeypatch +): + monkeypatch.chdir(tmp_path) + assert load_decision_config() == DecisionConfig() + with pytest.raises(DecisionConfigError): + load_decision_config(tmp_path / "absent.yaml") + + +def test_root_absent_environment_can_supply_connection(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DECISION_API_KEY", "environment-placeholder") + assert load_decision_config().api_key_value == "environment-placeholder" + + +def test_explicit_decision_key_does_not_resolve_unused_global_credentials(monkeypatch): + import veadk.config as config + from types import SimpleNamespace + from veadk.decision import resolve_decision_connection + + class Model: + api_base = "https://ark.cn-beijing.volces.com/api/v3" + + @property + def api_key(self): + raise AssertionError( + "unused global key resolution may call a cloud service" + ) + + monkeypatch.setattr(config, "settings", SimpleNamespace(model=Model())) + assert resolve_decision_connection( + DecisionConfig(api_key="explicit-placeholder") + ) == ("https://ark.cn-beijing.volces.com/api/v1/systemone", "explicit-placeholder") + with pytest.raises(DecisionConfigError, match="not an Ark"): + resolve_decision_connection(DecisionConfig(), "https://unrelated.example/v1") + + +@pytest.mark.parametrize( + "consumer,kwargs", + [ + (QueryRoutingConfig, {"simple_models": ["candidate"]}), + (ContentSafetyConfig, {"provider": "decision"}), + ], +) +def test_rejecting_removed_connection_never_prints_its_key(consumer, kwargs): + with pytest.raises(ValidationError) as error: + consumer(**kwargs, decision={"api_key": "obsolete-key-placeholder"}) + assert "obsolete-key-placeholder" not in str(error.value) diff --git a/tests/e2e/query_routing_http_agent.py b/tests/e2e/query_routing_http_agent.py new file mode 100644 index 000000000..2b7bd1b1e --- /dev/null +++ b/tests/e2e/query_routing_http_agent.py @@ -0,0 +1,177 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Agents loaded by the real ADK API server for opt-in HTTP acceptance tests.""" + +import asyncio +import json +import os +from pathlib import Path + +from google.adk.tools import AgentTool, ToolContext +from google.genai import types +from pydantic import BaseModel + +from veadk import Agent +from veadk.decision import DecisionClient +from veadk.routing import factory + +SIMPLE = os.getenv("VEADK_QUERY_ROUTING_SIMPLE_MODEL", "doubao-seed-2-1-lite-260915") +STRONG = os.getenv("VEADK_QUERY_ROUTING_STRONG_MODEL", "doubao-seed-2-1-pro-260915") + + +def audit(**record): + # Only synthetic test data and allowlisted metadata are written here. + with Path(os.environ["VEADK_QUERY_ROUTING_AUDIT"]).open("a") as output: + output.write(json.dumps(record, ensure_ascii=False) + "\n") + + +class ObservedDecisionClient(DecisionClient): + app_name = "" + + async def judge(self, **kwargs): + audit(kind="decision_start", app=self.app_name) + return await super().judge(**kwargs) + + +factory.DecisionClient = ObservedDecisionClient + + +async def observe_model(callback_context, llm_request): + audit( + kind="model", + agent=callback_context.agent_name, + invocation=callback_context.invocation_id, + model=llm_request.model, + ) + + +def lookup_order(order_id: str, tool_context: ToolContext) -> dict: + """Look up a synthetic order; retry once when the result is temporarily unavailable.""" + count = tool_context.state.get("lookup_count", 0) + 1 + tool_context.state["lookup_count"] = count + if count == 1: + return {"error": "temporarily_unavailable", "retryable": True} + orders = {f"ORD-{41 + i:03d}": (7 + i, 19 + i) for i in range(12)} + if order_id not in orders: + return {"error": "unknown_order"} + quantity, price = orders[order_id] + return {"order_id": order_id, "quantity": quantity, "unit_price": price} + + +def calculate_total(quantity: int, unit_price: int) -> dict: + """Calculate an order total from the values returned by lookup_order.""" + audit(kind="calculate_total", quantity=quantity, unit_price=unit_price) + return {"total": quantity * unit_price} + + +async def wait_for_inventory(tool_context: ToolContext) -> dict: + """Wait for an inventory lookup requested by the user.""" + tool_context.state["inventory_started"] = True + audit(kind="inventory_started") + cancelled = False + try: + await asyncio.sleep(20) + return {"inventory": 12} + except asyncio.CancelledError: + cancelled = True + raise + finally: + audit(kind="inventory_finished", cancelled=cancelled) + + +class CodeResult(BaseModel): + code: str + quantity: int + + +def build_agent(kind: str) -> Agent: + responses = kind != "chat" + routing = { + "mode": "enforce", + "simple_models": [ + { + "name": SIMPLE, + "api": "chat_completions" if kind in ("chat", "mixed") else "responses", + "supports_tools": True, + "supports_media": True, + "supports_structured_output": True, + } + ], + } + if kind == "shadow": + routing.pop("mode") + elif kind == "timeout": + routing["decision_timeout"] = 0.001 + + options = { + "name": kind, + "model_name": STRONG, + "model_api_base": os.environ["MODEL_AGENT_API_BASE"], + "model_api_key": os.environ["MODEL_AGENT_API_KEY"], + "enable_responses": responses, + "model_extra_config": { + "extra_body": {"thinking": {"type": "disabled"}}, + }, + "generate_content_config": types.GenerateContentConfig( + temperature=0, max_output_tokens=1024 + ), + "instruction": "Follow the user. Keep answers brief and remember facts across turns.", + "before_model_callback": observe_model, + } + if responses: + options["model_extra_config"]["store"] = True + options["model_extra_config"]["extra_body"]["caching"] = {"type": "disabled"} + if kind != "default_off": + options["query_routing"] = False if kind == "disabled" else routing + if kind == "tools": + options["tools"] = [lookup_order, calculate_total] + options["instruction"] = ( + "For order requests call lookup_order; retry once if it reports a retryable " + "error. Then call calculate_total with the returned quantity and unit_price. " + "Do not guess tool results. Return only the total." + ) + if kind == "cancel": + options["tools"] = [wait_for_inventory] + options["instruction"] = ( + "Call wait_for_inventory only when the user asks to look up inventory. " + "Otherwise answer directly." + ) + if kind == "schema": + options["output_schema"] = CodeResult + if kind == "nested": + leaf = Agent( + **{ + **options, + "name": "leaf", + "query_routing": False, + "instruction": "Return exactly NESTED_HTTP_OK.", + } + ) + child = Agent( + **{ + **options, + "name": "child", + "query_routing": False, + "tools": [AgentTool(leaf)], + "instruction": "Pass the user's full request to leaf once and return its answer exactly.", + } + ) + options["tools"] = [AgentTool(child)] + options["instruction"] = ( + "Pass the user's full request to child once and return its answer exactly." + ) + root = Agent(**options) + if kind not in ("disabled", "default_off"): + root.model.classifier._judge.app_name = kind + return root diff --git a/tests/e2e/test_content_safety_e2e.py b/tests/e2e/test_content_safety_e2e.py index 576f7e424..875737b54 100644 --- a/tests/e2e/test_content_safety_e2e.py +++ b/tests/e2e/test_content_safety_e2e.py @@ -37,6 +37,7 @@ from veadk import Agent as VeADKAgent from veadk import Runner as VeADKRunner +from veadk.decision import DecisionConfig from veadk.content_safety import ( Checkpoint, ContentSafetyConfig, @@ -155,6 +156,14 @@ def dangerous_tool(query: str) -> dict: return {"ok": True, "query": query} +def connection(server): + return DecisionConfig( + api_url=f"http://127.0.0.1:{server.server_address[1]}/api/v1/systemone", + api_key="e2e-key", + model="e2e-judge", + ) + + def make_config(server, *, on_error="allow"): return ContentSafetyConfig.from_dict( { @@ -169,11 +178,6 @@ def make_config(server, *, on_error="allow"): } for name in ALL_QUESTIONS }, - "decision": { - "api_url": f"http://127.0.0.1:{server.server_address[1]}/v1/judge", - "api_key": "e2e-key", - "model": "e2e-judge", - }, } ) @@ -223,7 +227,7 @@ def event_text(event): def test_e2e_allow_full_adk_turn(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[[text_response("hello from the model")]], @@ -255,7 +259,7 @@ def test_e2e_allow_full_adk_turn(): def test_e2e_blocks_user_input_before_model_call(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[[text_response("should never be called")]], @@ -278,7 +282,7 @@ def test_e2e_blocks_user_input_before_model_call(): def test_e2e_blocks_model_output_after_model_call(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[[text_response("SECRET value")]], @@ -301,7 +305,7 @@ def test_e2e_blocks_model_output_after_model_call(): def test_e2e_blocks_streamed_output_without_leaking_partials(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[ @@ -330,7 +334,7 @@ def test_e2e_blocks_streamed_output_without_leaking_partials(): def test_e2e_blocks_tool_before_execution(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[ @@ -362,7 +366,7 @@ def test_e2e_blocks_tool_before_execution(): def test_e2e_tool_runs_when_content_is_allowed(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[ @@ -390,7 +394,7 @@ def test_e2e_tool_runs_when_content_is_allowed(): def test_e2e_tool_argument_metadata_key_is_audited(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[ @@ -418,7 +422,9 @@ def test_e2e_tool_argument_metadata_key_is_audited(): def test_e2e_on_error_allow_degrades_open(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server, on_error="allow")) + plugin = build_content_safety( + make_config(server, on_error="allow"), decision=connection(server) + ) llm, events = asyncio.run( run_agent( turns=[[text_response("still answered")]], @@ -439,7 +445,9 @@ def test_e2e_on_error_allow_degrades_open(): def test_e2e_on_error_block_fails_closed(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server, on_error="block")) + plugin = build_content_safety( + make_config(server, on_error="block"), decision=connection(server) + ) llm, events = asyncio.run( run_agent( turns=[[text_response("should not appear")]], @@ -480,10 +488,6 @@ def test_e2e_config_yaml_auto_mounts_in_veadk_runner(tmp_path, monkeypatch): provider: decision use_default_questions: false block_message: {BLOCK_MESSAGE!r} - decision: - api_url: http://127.0.0.1:{server.server_address[1]}/v1/judge - api_key: e2e-key - model: e2e-judge questions: prompt_injection: instructions: Detect prompt_injection. @@ -494,6 +498,10 @@ def test_e2e_config_yaml_auto_mounts_in_veadk_runner(tmp_path, monkeypatch): high_risk_operation: instructions: Detect high_risk_operation. threshold: 0.5 +decision: + api_url: http://127.0.0.1:{server.server_address[1]}/v1/judge + api_key: e2e-key + model: e2e-judge """, encoding="utf-8", ) @@ -524,7 +532,7 @@ def test_e2e_config_yaml_auto_mounts_in_veadk_runner(tmp_path, monkeypatch): def test_e2e_multimodal_image_reaches_decision_api(): TOOL_CALLS.clear() with systemone_server() as server: - plugin = build_content_safety(make_config(server)) + plugin = build_content_safety(make_config(server), decision=connection(server)) llm, events = asyncio.run( run_agent( turns=[[text_response("image noted")]], diff --git a/tests/e2e/test_content_safety_live.py b/tests/e2e/test_content_safety_live.py index a898e9f0d..c08a0fb21 100644 --- a/tests/e2e/test_content_safety_live.py +++ b/tests/e2e/test_content_safety_live.py @@ -17,7 +17,7 @@ service from normal test runs:: VEADK_CONTENT_SAFETY_LIVE=1 \ - TOOL_CONTENT_SAFETY_DECISION_API_KEY=... \ + DECISION_API_KEY=... \ pytest tests/e2e/test_content_safety_live.py -q The API key is read from the environment and is never logged or asserted. @@ -35,6 +35,7 @@ import pytest from google.genai import types +from veadk.decision import DecisionConfig from veadk.content_safety import ( AuditAction, AuditStatus, @@ -45,9 +46,7 @@ ) LIVE_ENABLED = os.getenv("VEADK_CONTENT_SAFETY_LIVE") == "1" -LIVE_API_KEY = os.getenv("TOOL_CONTENT_SAFETY_DECISION_API_KEY") or os.getenv( - "ARK_API_KEY" -) +LIVE_API_KEY = os.getenv("DECISION_API_KEY") or os.getenv("ARK_API_KEY") DEFAULT_ENDPOINT = "https://ark.cn-beijing.volces.com/api/v1/systemone" RED_PNG = base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAIAAAD8GO2jAAAAK0lEQVR4nO3NQQEAAATAQGTQ" @@ -90,8 +89,7 @@ pytestmark = pytest.mark.skipif( not (LIVE_ENABLED and LIVE_API_KEY), reason=( - "set VEADK_CONTENT_SAFETY_LIVE=1 and " - "TOOL_CONTENT_SAFETY_DECISION_API_KEY (or ARK_API_KEY) to run" + "set VEADK_CONTENT_SAFETY_LIVE=1 and DECISION_API_KEY (or ARK_API_KEY) to run" ), ) @@ -107,20 +105,22 @@ def _plugin(): config = ContentSafetyConfig.from_dict( { "provider": "decision", - "decision": { - "api_url": os.getenv( - "TOOL_CONTENT_SAFETY_DECISION_API_URL", DEFAULT_ENDPOINT - ), + "questions": questions, + } + ) + return build_content_safety( + config, + decision=DecisionConfig.model_validate( + { + "api_url": os.getenv("DECISION_API_URL", DEFAULT_ENDPOINT), "api_key": LIVE_API_KEY, "model": os.getenv( - "TOOL_CONTENT_SAFETY_DECISION_MODEL", + "DECISION_MODEL", "doubao-seed-2-1-lite-260915", ), - }, - "questions": questions, - } + } + ), ) - return build_content_safety(config) def _request(text: str): @@ -180,19 +180,21 @@ def _plugin_with_question(instructions): "threshold": 0.5, } }, - "decision": { - "api_url": os.getenv( - "TOOL_CONTENT_SAFETY_DECISION_API_URL", DEFAULT_ENDPOINT - ), + } + ) + return build_content_safety( + config, + decision=DecisionConfig.model_validate( + { + "api_url": os.getenv("DECISION_API_URL", DEFAULT_ENDPOINT), "api_key": LIVE_API_KEY, "model": os.getenv( - "TOOL_CONTENT_SAFETY_DECISION_MODEL", + "DECISION_MODEL", "doubao-seed-2-1-lite-260915", ), - }, - } + } + ), ) - return build_content_safety(config) def _media_request(parts): diff --git a/tests/e2e/test_query_routing_http_live.py b/tests/e2e/test_query_routing_http_live.py new file mode 100644 index 000000000..b87622b6c --- /dev/null +++ b/tests/e2e/test_query_routing_http_live.py @@ -0,0 +1,528 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Live HTTP acceptance: separate ADK server process, real SystemOne and models. + +Run with VEADK_QUERY_ROUTING_LIVE=1 and MODEL_AGENT_API_BASE/API_KEY set. +No Runner, ASGI test transport or fabricated provider answers are used by the client. +Default questions, thresholds and the two-second budget are preserved except +in the explicitly named deadline-failure case. +""" + +import asyncio +import base64 +import json +import os +import socket +import struct +import subprocess +import sys +import time +import uuid +import zlib +from pathlib import Path + +import httpx +import pytest +from google.adk.events.event import Event + +pytestmark = [ + pytest.mark.query_routing_live, + pytest.mark.skipif( + os.getenv("VEADK_QUERY_ROUTING_LIVE") != "1" + or not os.getenv("MODEL_AGENT_API_KEY") + or not os.getenv("MODEL_AGENT_API_BASE"), + reason="requires explicit live enablement and model credentials", + ), +] + +APPS = ( + "responses", + "chat", + "mixed", + "tools", + "nested", + "schema", + "shadow", + "disabled", + "default_off", + "timeout", + "cancel", +) + + +@pytest.fixture(scope="module") +def server(tmp_path_factory): + root = tmp_path_factory.mktemp("query-routing-http") + agents = root / "agents" + agents.mkdir() + for name in APPS: + package = agents / name + package.mkdir() + (package / "__init__.py").write_text("from . import agent\n") + (package / "agent.py").write_text( + "from query_routing_http_agent import build_agent\n" + f"root_agent = build_agent({name!r})\n" + ) + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join( + [str(Path(__file__).parent), str(Path(__file__).resolve().parents[2])] + ) + env["VEADK_QUERY_ROUTING_AUDIT"] = str(root / "audit.jsonl") + for key in ("NO_PROXY", "no_proxy"): + env[key] = env.get(key, "") + ",127.0.0.1,localhost,::1" + command = [ + str(Path(sys.executable).with_name("adk")), + "api_server", + str(agents), + "--host", + "127.0.0.1", + "--port", + str(port), + "--no-reload", + "--session_service_uri", + "memory://", + "--artifact_service_uri", + "memory://", + ] + log = (root / "server.log").open("w") + process = subprocess.Popen(command, cwd=root, env=env, stdout=log, stderr=log) + (root / "server-command.json").write_text( + json.dumps({"command": command, "pid": process.pid}) + ) + url = f"http://127.0.0.1:{port}" + try: + with httpx.Client(base_url=url, trust_env=False, timeout=1) as client: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + assert process.poll() is None, ( + f"server exited; see {root / 'server.log'}" + ) + try: + response = client.get("/list-apps") + if response.status_code == 200: + assert set(APPS).issubset(response.json()) + break + except httpx.TransportError: + pass + time.sleep(0.1) + else: + pytest.fail("ADK API server did not start within 30 seconds") + yield {"url": url, "root": root, "command": command} + finally: + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + log.close() + + +def decisions(server, app): + path = server["root"] / "audit.jsonl" + records = ( + [json.loads(line) for line in path.read_text().splitlines()] + if path.exists() + else [] + ) + return [r for r in records if r["kind"] == "decision_start" and r["app"] == app] + + +async def session(client, app): + user = "http_acceptance" + response = await client.post(f"/apps/{app}/users/{user}/sessions", json={}) + response.raise_for_status() + return {"appName": app, "userId": user, "sessionId": response.json()["id"]} + + +async def turn(client, identity, text, *, stream=False, parts=None): + payload = { + **identity, + "newMessage": {"role": "user", "parts": parts or [{"text": text}]}, + } + if stream: + events = [] + async with client.stream( + "POST", "/run_sse", json={**payload, "streaming": True} + ) as response: + response.raise_for_status() + assert response.headers["content-type"].startswith("text/event-stream") + async for line in response.aiter_lines(): + if line.startswith("data: "): + event = json.loads(line[6:]) + events.append(event) + else: + response = await client.post("/run", json=payload) + response.raise_for_status() + events = response.json() + assert events, "Agent API returned no events" + answer = "" + for value in reversed(events): + event = Event.model_validate(value) + if event.is_final_response() and event.content: + answer = "".join( + p.text for p in event.content.parts or [] if p.text and not p.thought + ).strip() + if answer: + break + routes = [e.get("customMetadata", {}).get("veadk_query_routing") for e in events] + routes = [r for r in routes if r] + result = { + "events": events, + "answer": answer.strip(), + "routes": routes, + "errors": [ + { + "code": e.get("errorCode"), + "message": e.get("errorMessage") or e.get("error"), + } + for e in events + if e.get("errorCode") or e.get("error") + ], + } + assert not result["errors"], result["errors"] + return result + + +def save(server, case, result): + with (server["root"] / "results.jsonl").open("a") as output: + output.write(json.dumps({"case": case, **result}, ensure_ascii=False) + "\n") + + +def fixed_route(result, expected=None): + assert result["answer"], result + routes = result["routes"] + assert routes, result + assert len({r["model"] for r in routes}) == 1, routes + assert len({r["root_turn_id"] for r in routes}) == 1, routes + if expected: + assert routes[0]["route"] == expected, routes[0] + return routes[0] + + +NATURAL_CASES = [ + ( + "extract", + "从这段文字中提取工单编号,只输出编号:工单编号 TK-482,状态处理中。", + "simple", + "TK-482", + ), + ( + "medical", + "我服用华法林后出现黑便和头晕,今晚应该如何调整剂量,是否需要立即就医?", + "strong", + None, + ), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("app", ["responses", "chat"]) +@pytest.mark.parametrize( + "case,prompt,expected,answer", NATURAL_CASES, ids=[c[0] for c in NATURAL_CASES] +) +async def test_http_default_natural_queries( + server, app, case, prompt, expected, answer +): + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + identity = await session(client, app) + before = len(decisions(server, app)) + result = await turn(client, identity, prompt, stream=app == "responses") + save(server, f"natural/{app}/{case}", result) + route = fixed_route(result, expected) + assert route["reason"] not in ("decision_timeout", "decision_error") + assert len(decisions(server, app)) == before + 1 + if answer: + assert result["answer"].rstrip("。.") == answer + + +@pytest.mark.asyncio +@pytest.mark.parametrize("app", ["responses", "chat", "mixed"]) +async def test_http_natural_multiturn_history(server, app): + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + identity = await session(client, app) + before = len(decisions(server, app)) + for index, (prompt, expected, values) in enumerate( + [ + ("记住项目代号 ORCHID47。只回复该代号。", "simple", ["ORCHID47"]), + ("项目代号是什么?只回复代号。", "simple", ["ORCHID47"]), + ( + "记住里程碑代号 MAPLE83。另外,我服用华法林后出现黑便、头晕,是否应立即就医?", + "strong", + [], + ), + ( + "现在只提取之前给出的项目代号和里程碑代号,不讨论其他事情。", + "simple", + ["ORCHID47", "MAPLE83"], + ), + ] + ): + result = await turn(client, identity, prompt, stream=index % 2 == 1) + save(server, f"history/{app}/{index}", result) + fixed_route(result, expected) + assert all(value in result["answer"] for value in values) + assert len(decisions(server, app)) == before + 4 + fetched = await client.get( + f"/apps/{app}/users/{identity['userId']}/sessions/{identity['sessionId']}" + ) + fetched.raise_for_status() + assert ( + len( + { + e["invocationId"] + for e in fetched.json()["events"] + if e.get("customMetadata", {}).get("veadk_query_routing") + } + ) + == 4 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("app", ["disabled", "default_off", "shadow", "timeout"]) +async def test_http_opt_in_shadow_and_deadline(server, app): + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + identity = await session(client, app) + before = len(decisions(server, app)) + result = await turn(client, identity, "计算 17+25,只输出数字。") + save(server, f"config/{app}", result) + assert result["answer"] == "42" + if app in ("disabled", "default_off"): + assert not result["routes"] + assert len(decisions(server, app)) == before + else: + route = fixed_route(result, "strong") + assert route["reason"] == ( + "shadow" if app == "shadow" else "decision_timeout" + ) + assert len(decisions(server, app)) == before + 1 + + +@pytest.mark.asyncio +async def test_http_tool_error_recovery_keeps_one_decision(server): + app = "tools" + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + identity = await session(client, app) + before = len(decisions(server, app)) + result = await turn( + client, + identity, + "查询订单 ORD-041 并计算总价。查询暂时失败可以重试一次。", + stream=True, + ) + save(server, "tools/retry", result) + fixed_route(result) + assert result["answer"].rstrip(".。") == "133" + calls = [ + p["functionCall"] + for e in result["events"] + for p in (e.get("content") or {}).get("parts", []) + if p.get("functionCall") and not e.get("partial") + ] + assert [c["name"] for c in calls] == [ + "lookup_order", + "lookup_order", + "calculate_total", + ] + # Events preserve provider arguments; validate values at execution, + # after the schema guard normalizes unambiguous integer strings. + executed = [ + record + for line in (server["root"] / "audit.jsonl").read_text().splitlines() + if (record := json.loads(line))["kind"] == "calculate_total" + ] + assert executed == [ + {"kind": "calculate_total", "quantity": 7, "unit_price": 19} + ] + assert len(decisions(server, app)) == before + 1 + + +@pytest.mark.asyncio +async def test_http_nested_agents_and_schema(server): + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + for app, prompt in [ + ("nested", "请通过 child 获取 leaf 的结果。"), + ("schema", "项目编号 ORCHID47,数量 7。提取为要求的结构。"), + ]: + identity = await session(client, app) + before = len(decisions(server, app)) + result = await turn(client, identity, prompt) + save(server, app, result) + fixed_route(result) + assert len(decisions(server, app)) == before + 1 + if app == "nested": + assert "NESTED_HTTP_OK" in result["answer"] + records = [ + json.loads(line) + for line in (server["root"] / "audit.jsonl") + .read_text() + .splitlines() + ] + models = [ + r + for r in records + if r["kind"] == "model" + and r["agent"] in ("nested", "child", "leaf") + ] + assert {r["agent"] for r in models} == {"nested", "child", "leaf"} + assert len({r["model"] for r in models}) == 1 + else: + assert json.loads(result["answer"]) == { + "code": "ORCHID47", + "quantity": 7, + } + + +@pytest.mark.asyncio +async def test_http_concurrent_sessions_do_not_leak(server): + app = "mixed" + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + identities = [await session(client, app) for _ in range(4)] + codes = ["SESSION_" + uuid.uuid4().hex[:12] for _ in identities] + before = len(decisions(server, app)) + + async def conversation(index): + own = codes[index] + prompt = f"记住我的项目代号 {own}。只回复代号。" + if index % 2: + prompt += "另外,我服用华法林后出现黑便和头晕,是否应立即就医?" + first = await turn(client, identities[index], prompt, stream=True) + save(server, f"concurrent/{index}/first", first) + fixed_route(first, "strong" if index % 2 else "simple") + last = await turn( + client, identities[index], "只回复我之前给出的项目代号。", stream=True + ) + save(server, f"concurrent/{index}/last", last) + fixed_route(last, "simple") + assert own in last["answer"] + assert all( + code not in json.dumps(last["events"]) for code in codes if code != own + ) + + await asyncio.gather(*(conversation(i) for i in range(4))) + assert len(decisions(server, app)) == before + 8 + + +@pytest.mark.asyncio +async def test_http_image_with_default_decision(server): + def chunk(kind, data): + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data)) + ) + + png = ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", 64, 64, 8, 2, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress((b"\x00" + b"\xff\x00\x00" * 64) * 64)) + + chunk(b"IEND", b"") + ) + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + identity = await session(client, "responses") + result = await turn( + client, + identity, + "", + stream=True, + parts=[ + { + "text": "What is the dominant image color? Reply with one English color word." + }, + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(png).decode(), + } + }, + ], + ) + save(server, "image", result) + fixed_route(result, "simple") + assert result["answer"].lower().rstrip(".") == "red" + + +@pytest.mark.asyncio +async def test_http_sse_disconnect_cancels_tool_and_next_turn_recovers(server): + app = "cancel" + + def audit_records(): + return [ + json.loads(line) + for line in (server["root"] / "audit.jsonl").read_text().splitlines() + ] + + async with httpx.AsyncClient( + base_url=server["url"], trust_env=False, timeout=60 + ) as client: + identity = await session(client, app) + before = len(decisions(server, app)) + payload = { + **identity, + "newMessage": {"role": "user", "parts": [{"text": "请调用工具查询库存。"}]}, + "streaming": True, + } + async with client.stream("POST", "/run_sse", json=payload) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if line.startswith("data: "): + event = json.loads(line[6:]) + if any( + p.get("functionCall", {}).get("name") == "wait_for_inventory" + for p in (event.get("content") or {}).get("parts", []) + ): + break + deadline = time.monotonic() + 3 + while not any(r["kind"] == "inventory_started" for r in audit_records()): + assert time.monotonic() < deadline, "tool never started" + await asyncio.sleep(0.05) + deadline = time.monotonic() + 5 + while not any(r["kind"] == "inventory_finished" for r in audit_records()): + assert time.monotonic() < deadline, ( + "client disconnect did not stop the in-flight tool" + ) + await asyncio.sleep(0.05) + assert any( + r["kind"] == "inventory_finished" and r["cancelled"] + for r in audit_records() + ) + result = await turn( + client, + identity, + "取消刚才的查询。现在不要调用工具,只计算 17+25,输出数字。", + stream=True, + ) + save(server, "disconnect/recovery", result) + fixed_route(result) + assert result["answer"] == "42" + assert len(decisions(server, app)) == before + 2 diff --git a/tests/e2e/test_query_routing_live.py b/tests/e2e/test_query_routing_live.py new file mode 100644 index 000000000..1fbf2a5c9 --- /dev/null +++ b/tests/e2e/test_query_routing_live.py @@ -0,0 +1,445 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Real SystemOne + Ark model E2E, explicitly enabled by environment. + +Run with VEADK_QUERY_ROUTING_LIVE=1, MODEL_AGENT_API_KEY and +MODEL_AGENT_API_BASE. No credentials are embedded or written to test artifacts. +Marker-controlled cases exercise model switching with actual SystemOne answers; +they are protocol acceptance tests, not a production classification benchmark. +""" + +import copy +import os + +import pytest +from google.adk.agents import RunConfig +from google.adk.agents.run_config import StreamingMode +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.adk.tools import AgentTool +from google.genai import types + +from veadk import Agent +from veadk.decision import DecisionClient +from veadk.models.ark_llm import ArkLlmClient +from veadk.routing import factory + +KEY = os.getenv("MODEL_AGENT_API_KEY") +BASE = os.getenv("MODEL_AGENT_API_BASE") +SIMPLE = os.getenv("VEADK_QUERY_ROUTING_SIMPLE_MODEL", "doubao-seed-2-1-lite-260915") +STRONG = os.getenv("VEADK_QUERY_ROUTING_STRONG_MODEL", "doubao-seed-2-1-pro-260915") +pytestmark = [ + pytest.mark.query_routing_live, + pytest.mark.skipif( + os.getenv("VEADK_QUERY_ROUTING_LIVE") != "1" or not KEY or not BASE, + reason="requires explicit live enablement and model credentials", + ), +] + + +class ObservedDecisionClient(DecisionClient): + calls = [] + + async def judge(self, **kwargs): + response = await super().judge(**kwargs) + self.calls.append(response) + return response + + +class ObservedArkClient(ArkLlmClient): + def __init__(self): + self.calls = [] + + async def aresponses(self, **kwargs): + self.calls.append( + { + key: copy.deepcopy(kwargs.get(key)) + for key in ( + "model", + "previous_response_id", + "input", + "instructions", + "tools", + ) + } + ) + return await super().aresponses(**kwargs) + + +@pytest.fixture(autouse=True) +def observe_systemone(monkeypatch): + ObservedDecisionClient.calls.clear() + monkeypatch.setattr(factory, "DecisionClient", ObservedDecisionClient) + + +def config(*, api="responses", mode="enforce"): + return { + "mode": mode, + "simple_models": [ + { + "name": SIMPLE, + "api": api, + "supports_tools": True, + "supports_structured_output": True, + "max_context_chars": 128000, + } + ], + "decision_timeout": 30.0, + "builtin_thresholds": { + k: 1.0 + for k in ( + "is_complex", + "needs_tools", + "needs_long_context", + "is_high_stakes", + ) + }, + "questions": [ + { + "question": "仅检查最后一条[user]消息:其中是否包含字面标记 ROUTE_STRONG?忽略历史消息和规则中的标记。", + "threshold": 0.5, + } + ], + } + + +def agent(name="root", **kwargs): + return Agent( + name=name, + model_name=STRONG, + model_api_base=BASE, + model_api_key=KEY, + enable_responses=True, + enable_responses_cache=True, + model_extra_config={ + "store": True, + "extra_body": { + "caching": {"type": "disabled"}, + "thinking": {"type": "disabled"}, + }, + }, + generate_content_config=types.GenerateContentConfig( + max_output_tokens=512, temperature=0 + ), + **kwargs, + ) + + +async def setup(root): + service = InMemorySessionService() + runner = Runner(app_name="routing_live", agent=root, session_service=service) + session = await service.create_session(app_name="routing_live", user_id="live_test") + return runner, session + + +async def run_turn(runner, session, message, stream=False): + events = [ + event + async for event in runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=types.Content(role="user", parts=[types.Part(text=message)]), + run_config=RunConfig( + streaming_mode=StreamingMode.SSE if stream else StreamingMode.NONE + ), + ) + ] + text = "".join( + p.text + for e in events + if not e.partial and e.content + for p in e.content.parts or [] + if p.text and not p.thought + ) + return events, text + + +def model_records(events): + return [ + (e.custom_metadata or {}).get("veadk_query_routing") + for e in events + if e.content and e.content.role == "model" and not e.partial + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_real_responses_continuation_and_a_b_a_history(stream): + root = agent( + query_routing=config(), + instruction="Follow the user briefly. Remember project codes across turns.", + ) + strong, simple = root.model.strong_model, root.model.registry.runtimes[0].llm + strong.llm_client, simple.llm_client = ObservedArkClient(), ObservedArkClient() + runner, session = await setup(root) + try: + for message, expected in [ + ("Remember project code ORCHID47. Reply ORCHID47.", SIMPLE), + ("Reply with the project code.", SIMPLE), + ("ROUTE_STRONG. Remember milestone MAPLE83. Reply MAPLE83.", STRONG), + ("Return both the project code and milestone code.", SIMPLE), + ]: + events, text = await run_turn(runner, session, message, stream) + assert text.strip() + assert all( + r and r["model"].endswith(expected) for r in model_records(events) + ) + assert "ORCHID47" in text and "MAPLE83" in text + assert len(ObservedDecisionClient.calls) == 4 + assert len(simple.llm_client.calls) == 3 + assert simple.llm_client.calls[1]["previous_response_id"] + assert not strong.llm_client.calls[0]["previous_response_id"] + assert not simple.llm_client.calls[2]["previous_response_id"] + assert "ORCHID47" in str(simple.llm_client.calls[2]["input"]) + assert "MAPLE83" in str(simple.llm_client.calls[2]["input"]) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_real_chat_adapter_with_two_nested_agent_tools(): + seen = [] + + async def observe(callback_context, llm_request): + seen.append( + ( + callback_context.agent_name, + callback_context.invocation_id, + llm_request.model, + ) + ) + + leaf = agent( + "leaf", + query_routing=False, + instruction="Reply exactly NESTED_OK.", + before_model_callback=observe, + ) + child = agent( + "child", + query_routing=False, + tools=[AgentTool(leaf)], + instruction="Call leaf exactly once, then return its answer.", + before_model_callback=observe, + ) + root = agent( + query_routing=config(api="chat_completions"), + tools=[AgentTool(child)], + instruction="Call child exactly once, then return its answer.", + before_model_callback=observe, + ) + runner, session = await setup(root) + try: + events, text = await run_turn( + runner, session, "Use child to obtain the leaf response." + ) + assert "NESTED_OK" in text, (text, seen) + assert len(ObservedDecisionClient.calls) == 1 + assert {name for name, _, _ in seen} == {"root", "child", "leaf"} + assert len({inv for _, inv, _ in seen}) == 3 + assert all(model.endswith(SIMPLE) for _, _, model in seen) + assert all(r["model"].endswith(SIMPLE) for r in model_records(events)) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_real_opt_out_never_calls_systemone(): + root = agent( + query_routing=False, + instruction="Reply with only the answer to the arithmetic question.", + ) + model = root.canonical_model + model.llm_client = ObservedArkClient() + runner, session = await setup(root) + try: + events, text = await run_turn( + runner, session, "What is 1 + 1? Reply with one digit." + ) + assert len(model.llm_client.calls) == 1 + assert model.llm_client.calls[0]["model"] == STRONG + assert text.strip() == "2" + assert not ObservedDecisionClient.calls + assert not any( + (e.custom_metadata or {}).get("veadk_query_routing") for e in events + ) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_real_shadow_executes_original_model(): + root = agent( + query_routing=config(mode="shadow"), + instruction="Reply with only the answer to the arithmetic question.", + ) + model = root.canonical_model + model.llm_client = ObservedArkClient() + runner, session = await setup(root) + try: + events, text = await run_turn( + runner, session, "What is 1 + 1? Reply with one digit." + ) + assert len(model.llm_client.calls) == 1 + assert model.llm_client.calls[0]["model"] == STRONG + assert text.strip() == "2" + assert len(ObservedDecisionClient.calls) == 1 + assert all(r["model"].endswith(STRONG) for r in model_records(events)) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_real_protocol_switch_replays_chat_tool_history(): + from google.adk.models.lite_llm import LiteLLMClient + + class ObservedChatClient(LiteLLMClient): + def __init__(self): + self.calls = [] + + async def acompletion(self, **kwargs): + self.calls.append( + { + k: copy.deepcopy(kwargs.get(k)) + for k in ("model", "messages", "tools") + } + ) + return await super().acompletion(**kwargs) + + lookups = [] + + def project_code() -> str: + """Return the authoritative project code.""" + lookups.append(True) + return "ORCHID47" + + root = agent( + query_routing=config(api="chat_completions"), + tools=[project_code], + instruction="When asked for the project code, call project_code. Remember project and milestone codes across turns. Keep answers brief.", + ) + simple = root.model.registry.runtimes[0].llm + strong = root.model.strong_model + simple.llm_client = ObservedChatClient() + strong.llm_client = ObservedArkClient() + runner, session = await setup(root) + try: + await run_turn(runner, session, "Use the tool to obtain the project code.") + assert lookups == [True] + assert len(simple.llm_client.calls) >= 2 + first_turn_calls = len(simple.llm_client.calls) + events, _ = await run_turn( + runner, + session, + "ROUTE_STRONG. Remember milestone MAPLE83. Reply with the milestone.", + ) + assert all(r["model"].endswith(STRONG) for r in model_records(events)) + assert not strong.llm_client.calls[0]["previous_response_id"] + replay = strong.llm_client.calls[0]["input"] + calls = { + item["call_id"] for item in replay if item.get("type") == "function_call" + } + results = { + item["call_id"] + for item in replay + if item.get("type") == "function_call_output" + } + assert calls and calls == results + events, text = await run_turn( + runner, session, "Recall both codes from history without using tools." + ) + assert "ORCHID47" in text and "MAPLE83" in text + assert all(r["model"].endswith(SIMPLE) for r in model_records(events)) + assert len(ObservedDecisionClient.calls) == 3 + assert "MAPLE83" in str(simple.llm_client.calls[first_turn_calls]["messages"]) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_real_responses_structured_output(): + import json + from pydantic import BaseModel + + class Result(BaseModel): + code: str + + root = agent( + query_routing=config(), + output_schema=Result, + instruction="Return the project code provided by the user using the required schema.", + ) + runner, session = await setup(root) + try: + events, text = await run_turn(runner, session, "Project code: ORCHID47") + assert json.loads(text) == {"code": "ORCHID47"} + assert len(ObservedDecisionClient.calls) == 1 + assert all(r["model"].endswith(SIMPLE) for r in model_records(events)) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_real_media_reaches_systemone_and_selected_responses_model(): + import struct + import zlib + + def chunk(kind, data): + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data)) + ) + + png = ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", 64, 64, 8, 2, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress((b"\x00" + b"\xff\x00\x00" * 64) * 64)) + + chunk(b"IEND", b"") + ) + routing = config() + routing["simple_models"][0]["supports_media"] = True + root = agent( + query_routing=routing, + instruction="Identify the dominant image color. Reply with one English color word.", + ) + runner, session = await setup(root) + try: + events = [ + event + async for event in runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=types.Content( + role="user", + parts=[ + types.Part(text="What color fills this image?"), + types.Part.from_bytes(data=png, mime_type="image/png"), + ], + ), + ) + ] + text = " ".join( + p.text + for e in events + if e.content + for p in e.content.parts or [] + if p.text and not p.thought + ) + assert "red" in text.lower() + assert len(ObservedDecisionClient.calls) == 1 + assert all(r["model"].endswith(SIMPLE) for r in model_records(events)) + finally: + await runner.close() diff --git a/tests/routing/__init__.py b/tests/routing/__init__.py new file mode 100644 index 000000000..7f463206f --- /dev/null +++ b/tests/routing/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/routing/conftest.py b/tests/routing/conftest.py new file mode 100644 index 000000000..d2761ed49 --- /dev/null +++ b/tests/routing/conftest.py @@ -0,0 +1,24 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Keep loopback integration traffic local, even on proxy-configured hosts.""" + +import os + +import pytest + + +@pytest.fixture(autouse=True) +def local_http_bypass(monkeypatch): + for name in ("NO_PROXY", "no_proxy"): + monkeypatch.setenv(name, os.getenv(name, "") + ",127.0.0.1,localhost,::1") diff --git a/tests/routing/test_agent_e2e.py b/tests/routing/test_agent_e2e.py new file mode 100644 index 000000000..154efbca2 --- /dev/null +++ b/tests/routing/test_agent_e2e.py @@ -0,0 +1,331 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end routing through a real ADK Runner and a real HTTP Decision service. + +The model itself is deterministic, but everything above it is the production +path: Agent construction, config.yaml auto-discovery, the ADK callback, the +routing proxy, httpx, TCP, the local System One-shaped server, and sessions. +""" + +import json +import threading +from collections.abc import AsyncGenerator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest +from google.adk.apps import App +from google.adk.models.base_llm import BaseLlm +from google.adk.models.llm_response import LlmResponse +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.genai import types + +from veadk import Agent +from veadk.configs.routing_configs import QueryRoutingConfig +from veadk.routing.llm import RoutingLlm +from veadk.routing.store import get_route_store + + +class _Handler(BaseHTTPRequestHandler): + def do_POST(self): # noqa: N802 - stdlib naming + length = int(self.headers.get("Content-Length", "0")) + body = json.loads(self.rfile.read(length)) + state = "\n".join( + part.get("text", "") + for part in body.get("state", []) + if part.get("type") == "text" + ) + self.server.requests.append(body) + + if "FAIL" in state: + self.send_response(500) + self.end_headers() + return + + values = { + "is_complex": 0.99 if "COMPLEX" in state else 0.01, + "needs_tools": 0.01, + "needs_long_context": 0.01, + "is_high_stakes": 0.01, + } + payload = { + "answers": { + name: {"type": "noul", "noul": value} for name, value in values.items() + } + } + encoded = json.dumps(payload).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *args): + pass + + +class _Server(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.requests: list[dict] = [] + + +@contextmanager +def decision_server(): + server = _Server(("127.0.0.1", 0), _Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class NamedLlm(BaseLlm): + model: str = "named" + reply: str = "" + calls: int = 0 + + async def generate_content_async( + self, llm_request, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + self.calls += 1 + yield LlmResponse( + content=types.Content(role="model", parts=[types.Part(text=self.reply)]) + ) + + +def write_config(tmp_path, server, *, mode="enforce"): + (tmp_path / "config.yaml").write_text( + f""" +query_routing: + mode: {mode} + simple_models: + - cheap-model +decision: + api_url: http://127.0.0.1:{server.server_address[1]}/api/v1/systemone + api_key: e2e-key + model: e2e-judge +""", + encoding="utf-8", + ) + + +def make_agent(): + agent = Agent( + name="e2e_agent", + model_name="strong-model", + model_api_base="https://ark.cn-beijing.volces.com/api/v3/", + model_api_key="e2e-model-key", + ) + assert isinstance(agent.model, RoutingLlm) + agent.model.strong_model = NamedLlm(model="strong", reply="STRONG") + agent.model.registry.runtimes[0].llm = NamedLlm(model="cheap", reply="CHEAP") + return agent + + +async def run_turn(agent, *, session_id, message): + app = App(name="e2e-app", root_agent=agent) + service = InMemorySessionService() + await service.create_session( + app_name="e2e-app", user_id="user", session_id=session_id + ) + runner = Runner(app=app, session_service=service) + events = [ + event + async for event in runner.run_async( + user_id="user", + session_id=session_id, + new_message=types.Content(role="user", parts=[types.Part(text=message)]), + ) + ] + return "".join( + part.text + for event in events + if event.content + for part in event.content.parts + if isinstance(part.text, str) + ) + + +@pytest.fixture(autouse=True) +def _clear_store(): + get_route_store().clear() + yield + get_route_store().clear() + + +@pytest.mark.asyncio +async def test_simple_query_uses_cheap_model_and_one_decision_call( + tmp_path, monkeypatch +): + with decision_server() as server: + write_config(tmp_path, server) + monkeypatch.chdir(tmp_path) + agent = make_agent() + cheap = agent.model.registry.runtimes[0].llm + + text = await run_turn( + agent, session_id="simple", message="你好,介绍一下你自己" + ) + + assert text == "CHEAP" + assert cheap.calls == 1 + assert len(server.requests) == 1 + assert set(server.requests[0]["questions"]) == { + "is_complex", + "needs_tools", + "needs_long_context", + "is_high_stakes", + } + assert "你好" in json.dumps(server.requests[0]["state"], ensure_ascii=False) + + +@pytest.mark.asyncio +async def test_next_invocation_is_classified_independently(tmp_path, monkeypatch): + with decision_server() as server: + write_config(tmp_path, server) + monkeypatch.chdir(tmp_path) + agent = make_agent() + strong = agent.model.strong_model + cheap = agent.model.registry.runtimes[0].llm + + first = await run_turn( + agent, session_id="complex", message="COMPLEX: 设计一个系统" + ) + second = await run_turn(agent, session_id="complex", message="再解释一下") + + assert first == "STRONG" + assert second == "CHEAP" + assert cheap.calls == 1 + assert strong.calls == 1 + # A new user invocation is classified again. + assert len(server.requests) == 2 + + +@pytest.mark.asyncio +async def test_decision_failure_falls_back_to_strong(tmp_path, monkeypatch): + with decision_server() as server: + write_config(tmp_path, server) + monkeypatch.chdir(tmp_path) + agent = make_agent() + + text = await run_turn(agent, session_id="failure", message="FAIL please") + + assert text == "STRONG" + assert len(server.requests) == 1 + + +@pytest.mark.asyncio +async def test_shadow_mode_never_switches(tmp_path, monkeypatch): + with decision_server() as server: + write_config(tmp_path, server, mode="shadow") + monkeypatch.chdir(tmp_path) + agent = make_agent() + cheap = agent.model.registry.runtimes[0].llm + + text = await run_turn(agent, session_id="shadow", message="你好") + + assert text == "STRONG" + assert cheap.calls == 0 + assert len(server.requests) == 1 + + +@pytest.mark.asyncio +async def test_config_without_routing_block_is_untouched(tmp_path, monkeypatch): + (tmp_path / "config.yaml").write_text("model:\n agent:\n name: x\n") + monkeypatch.chdir(tmp_path) + + agent = Agent( + name="plain_agent", + model_name="strong-model", + model_api_base="https://ark.cn-beijing.volces.com/api/v3/", + model_api_key="e2e-model-key", + ) + + assert not isinstance(agent.model, RoutingLlm) + assert agent.before_model_callback is None + + +def make_router_agent(name, **kwargs): + return Agent( + name=name, + model_name="strong-model", + model_api_base="https://ark.cn-beijing.volces.com/api/v3/", + model_api_key="e2e-model-key", + query_routing={"simple_models": ["cheap-model"]}, + **kwargs, + ) + + +async def noop_callback(callback_context, llm_request): + return None + + +def test_existing_callable_callback_is_preserved(): + agent = make_router_agent("callback_agent", before_model_callback=noop_callback) + assert isinstance(agent.model, RoutingLlm) + assert agent.before_model_callback is noop_callback + + +def test_existing_callback_list_is_extended(): + agent = make_router_agent( + "callback_list_agent", before_model_callback=[noop_callback] + ) + assert agent.before_model_callback[0] is noop_callback + assert len(agent.before_model_callback) == 1 + + +def test_routing_on_a_non_adk_runtime_is_rejected(): + agent = Agent( + name="runtime_agent", + model_name="strong-model", + model_api_base="https://ark.cn-beijing.volces.com/api/v3/", + model_api_key="e2e-model-key", + query_routing=False, + ) + agent.runtime = "codex" + agent.query_routing = QueryRoutingConfig(simple_models=["cheap-model"]) + + with pytest.raises(ValueError, match="only supported by the ADK runtime"): + agent._setup_query_routing() + + +def test_update_model_is_rejected_while_routing_is_enabled(): + agent = make_router_agent("update_agent") + + with pytest.raises(ValueError, match="update_model is not supported"): + agent.update_model("another-model") + + +@pytest.mark.asyncio +async def test_explicit_false_disables_yaml_routing(tmp_path, monkeypatch): + with decision_server() as server: + write_config(tmp_path, server) + monkeypatch.chdir(tmp_path) + + agent = Agent( + name="disabled_agent", + model_name="strong-model", + model_api_base="https://ark.cn-beijing.volces.com/api/v3/", + model_api_key="e2e-model-key", + query_routing=False, + ) + + assert not isinstance(agent.model, RoutingLlm) diff --git a/tests/routing/test_callback.py b/tests/routing/test_callback.py new file mode 100644 index 000000000..2bc6345e8 --- /dev/null +++ b/tests/routing/test_callback.py @@ -0,0 +1,30 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest +from google.adk.models.llm_request import LlmRequest + +from veadk.routing.callback import build_before_model_callback +from veadk.routing.context import get_turn_context + + +@pytest.mark.asyncio +async def test_callback_does_not_classify_outside_routed_execution(): + callback = build_before_model_callback(None, None, None) + request = LlmRequest(model="original") + assert await callback(SimpleNamespace(), request) is None + assert request.model == "original" + assert get_turn_context() is None diff --git a/tests/routing/test_capability.py b/tests/routing/test_capability.py new file mode 100644 index 000000000..5eaf6bbd6 --- /dev/null +++ b/tests/routing/test_capability.py @@ -0,0 +1,125 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from google.adk.models.llm_request import LlmRequest +from google.genai import types + +from veadk.configs.routing_configs import SimpleModelConfig +from veadk.routing.capability import CapabilityGuard +from veadk.routing.registry import ModelRegistry, SimpleModelRuntime +from veadk.routing.state import request_has_media, request_text_chars + +PNG_BYTES = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def request_with(parts, *, config=None): + kwargs = {"config": config} if config is not None else {} + return LlmRequest( + model="provider/model", + contents=[types.Content(role="user", parts=parts)], + **kwargs, + ) + + +def text_request(text="hello"): + return request_with([types.Part(text=text)]) + + +def test_text_only_is_allowed_by_default(): + check = CapabilityGuard().check(text_request(), SimpleModelConfig(name="cheap")) + assert check.allowed is True + assert check.reason == "" + + +def test_registry_check_matches_the_guard(): + runtime = SimpleModelRuntime(config=SimpleModelConfig(name="cheap"), llm=object()) + registry = ModelRegistry(runtimes=[runtime]) + check = registry.check(text_request(), runtime) + assert check.allowed is True + + +def test_media_is_rejected_by_default(): + request = request_with( + [types.Part(inline_data=types.Blob(mime_type="image/png", data=PNG_BYTES))] + ) + assert request_has_media(request) is True + check = CapabilityGuard().check(request, SimpleModelConfig(name="cheap")) + assert check.allowed is False + assert check.reason == "media_not_supported" + + +def test_media_is_allowed_when_declared(): + request = request_with( + [types.Part(inline_data=types.Blob(mime_type="image/png", data=PNG_BYTES))] + ) + model = SimpleModelConfig(name="cheap", supports_media=True) + assert CapabilityGuard().check(request, model).allowed is True + + +def test_tools_are_rejected_by_default(): + request = text_request() + request.tools_dict = {"lookup": object()} + check = CapabilityGuard().check(request, SimpleModelConfig(name="cheap")) + assert check.allowed is False + assert check.reason == "tools_not_supported" + + +def test_tools_are_allowed_when_declared(): + request = text_request() + request.tools_dict = {"lookup": object()} + model = SimpleModelConfig(name="cheap", supports_tools=True) + assert CapabilityGuard().check(request, model).allowed is True + + +def test_structured_output_is_rejected_by_default(): + config = types.GenerateContentConfig(response_schema={"type": "object"}) + request = request_with([types.Part(text="hello")], config=config) + check = CapabilityGuard().check(request, SimpleModelConfig(name="cheap")) + assert check.allowed is False + assert check.reason == "structured_output_not_supported" + + +def test_structured_output_is_allowed_when_declared(): + config = types.GenerateContentConfig(response_schema={"type": "object"}) + request = request_with([types.Part(text="hello")], config=config) + model = SimpleModelConfig(name="cheap", supports_structured_output=True) + assert CapabilityGuard().check(request, model).allowed is True + + +def test_context_limit_uses_request_characters(): + request = text_request("x" * 120) + assert request_text_chars(request) == 120 + model = SimpleModelConfig(name="cheap", max_context_chars=100) + check = CapabilityGuard().check(request, model) + assert check.allowed is False + assert check.reason == "context_too_large" + + +def test_context_limit_is_inclusive(): + model = SimpleModelConfig(name="cheap", max_context_chars=100) + exact = CapabilityGuard().check(text_request("x" * 100), model) + over = CapabilityGuard().check(text_request("x" * 101), model) + assert exact.allowed is True + assert over.allowed is False + assert over.reason == "context_too_large" + + +def test_media_counts_towards_context_budget(): + request = request_with( + [types.Part(inline_data=types.Blob(mime_type="image/png", data=PNG_BYTES))] + ) + assert request_text_chars(request) == 4_000 diff --git a/tests/routing/test_classifier.py b/tests/routing/test_classifier.py new file mode 100644 index 000000000..8ccaa6da7 --- /dev/null +++ b/tests/routing/test_classifier.py @@ -0,0 +1,204 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from types import SimpleNamespace + +import pytest +from google.adk.models.llm_request import LlmRequest +from google.genai import types + +from veadk.configs.routing_configs import QueryRoutingConfig +from veadk.decision import DecisionClientError +from veadk.routing.classifier import QueryClassifier +from veadk.routing.registry import ModelRegistry, SimpleModelRuntime + + +class FakeOuterLlm: + model = "simple" + + +def answers(**overrides): + values = { + "is_complex": 0.01, + "needs_tools": 0.01, + "needs_long_context": 0.01, + "is_high_stakes": 0.01, + } + values.update(overrides) + return {name: {"type": "noul", "noul": value} for name, value in values.items()} + + +class Judge: + def __init__(self, result=None, error=None, delay=0.0): + self.result = result or answers() + self.error = error + self.delay = delay + self.calls = [] + + async def judge(self, *, state, questions): + self.calls.append((state, questions)) + if self.delay: + await asyncio.sleep(self.delay) + if self.error is not None: + raise self.error + return SimpleNamespace(answers=self.result) + + +def request(text="hello"): + return LlmRequest( + model="strong", + contents=[types.Content(role="user", parts=[types.Part(text=text)])], + ) + + +def classifier(judge, *, mode="enforce", supports_tools=False): + config = QueryRoutingConfig.model_validate( + { + "simple_models": [{"name": "simple", "supports_tools": supports_tools}], + "mode": mode, + "decision_timeout": 0.05, + } + ) + runtime = SimpleModelRuntime(config=config.simple_models[0], llm=FakeOuterLlm()) + return QueryClassifier(config, judge, ModelRegistry(runtimes=[runtime])) + + +@pytest.mark.asyncio +async def test_simple_route_returns_selected_model(): + judge = Judge() + selection = await classifier(judge).classify( + request(), invocation_id="inv", session_id="s" + ) + assert selection.route == "simple" + assert selection.model_name == "simple" + assert selection.reason == "all_thresholds_passed" + assert judge.calls + + +def test_registry_property_exposes_the_model_pool(): + router = classifier(Judge()) + assert router.registry is router._registry + + +@pytest.mark.asyncio +async def test_optional_capability_description_is_sent_without_connection_secrets(): + description = ( + "Validated for extraction and one-step tools; not exact long-table sums." + ) + config = QueryRoutingConfig( + simple_models=[ + { + "name": "simple", + "capability_description": description, + "api_key_env": "TEST_CREDENTIAL_REFERENCE", + } + ] + ) + runtime = SimpleModelRuntime(config=config.simple_models[0], llm=FakeOuterLlm()) + judge = Judge() + classifier = QueryClassifier(config, judge, ModelRegistry(runtimes=[runtime])) + await classifier.classify(request(), invocation_id="inv") + state = str(judge.calls[0][0]) + assert description in state + assert "TEST_CREDENTIAL_REFERENCE" not in state + + +@pytest.mark.asyncio +async def test_no_capable_model_skips_decision_call(): + judge = Judge() + req = request() + req.tools_dict = {"tool": object()} + selection = await classifier(judge).classify(req, invocation_id="inv") + assert selection.route == "strong" + assert selection.reason == "no_capable_simple_model" + assert judge.calls == [] + + +@pytest.mark.asyncio +async def test_empty_state_routes_strong(): + judge = Judge() + empty = LlmRequest(model="strong", contents=[]) + selection = await classifier(judge).classify(empty, invocation_id="inv") + assert selection.route == "strong" + assert selection.reason == "empty_state" + assert judge.calls == [] + + +@pytest.mark.asyncio +async def test_timeout_routes_strong(): + judge = Judge(delay=0.2) + selection = await classifier(judge).classify(request(), invocation_id="inv") + assert selection.route == "strong" + assert selection.reason == "decision_timeout" + + +@pytest.mark.asyncio +async def test_client_error_routes_strong(): + judge = Judge(error=DecisionClientError("timeout", "slow")) + selection = await classifier(judge).classify(request(), invocation_id="inv") + assert selection.route == "strong" + assert selection.reason == "decision_error" + + +@pytest.mark.asyncio +async def test_unexpected_error_routes_strong(): + judge = Judge(error=RuntimeError("boom")) + selection = await classifier(judge).classify(request(), invocation_id="inv") + assert selection.route == "strong" + assert selection.reason == "decision_error" + + +@pytest.mark.asyncio +async def test_policy_rejection_routes_strong_and_upgrades_session(): + judge = Judge(result=answers(is_complex=0.9)) + selection = await classifier(judge).classify(request(), invocation_id="inv") + assert selection.route == "strong" + assert selection.reason == "is_complex_above_threshold" + + +@pytest.mark.asyncio +async def test_shadow_mode_records_simple_but_routes_strong(): + judge = Judge() + selection = await classifier(judge, mode="shadow").classify( + request(), invocation_id="inv" + ) + assert selection.route == "strong" + assert selection.reason == "shadow" + assert selection.model_name == "simple" + assert selection.shadow is True + + +@pytest.mark.asyncio +async def test_judge_receives_builtin_and_custom_questions(): + judge = Judge() + config = QueryRoutingConfig.model_validate( + { + "simple_models": ["simple"], + "mode": "enforce", + "questions": [{"question": "是否需要最新信息?", "threshold": 0.2}], + } + ) + runtime = SimpleModelRuntime(config=config.simple_models[0], llm=FakeOuterLlm()) + router = QueryClassifier(config, judge, ModelRegistry(runtimes=[runtime])) + await router.classify(request(), invocation_id="inv") + _, questions = judge.calls[0] + assert list(questions) == [ + "is_complex", + "needs_tools", + "needs_long_context", + "is_high_stakes", + "custom_1", + ] + assert all(question.type == "noul" for question in questions.values()) diff --git a/tests/routing/test_config.py b/tests/routing/test_config.py new file mode 100644 index 000000000..f1d94d8b6 --- /dev/null +++ b/tests/routing/test_config.py @@ -0,0 +1,357 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from pydantic import ValidationError + +from veadk.configs.routing_configs import ( + DEFAULT_QUERY_ROUTING_THRESHOLD, + QueryRoutingConfig, + QueryRoutingConfigError, + RoutingQuestion, + load_query_routing_config, + load_query_routing_config_if_present, +) + + +def write_config(tmp_path, text): + path = tmp_path / "config.yaml" + path.write_text(text, encoding="utf-8") + return path + + +def test_builtin_questions_are_always_present_and_ordered(): + config = QueryRoutingConfig(simple_models=["cheap-model"]) + assert [question.id for question in config.questions] == [ + "is_complex", + "needs_tools", + "needs_long_context", + "is_high_stakes", + ] + thresholds = {question.id: question.threshold for question in config.questions} + assert thresholds["is_high_stakes"] == 0.05 + assert all( + thresholds[name] == 0.5 + for name in ("is_complex", "needs_tools", "needs_long_context") + ) + assert all(question.type == "noul" for question in config.questions) + + +def test_simple_model_accepts_string_and_mapping(): + config = QueryRoutingConfig( + simple_models=[ + "cheap-a", + {"name": "cheap-b", "supports_tools": True, "max_context_chars": 1000}, + ] + ) + assert [model.name for model in config.simple_models] == ["cheap-a", "cheap-b"] + assert config.simple_models[0].supports_tools is False + assert config.simple_models[0].max_context_chars == 32_000 + assert config.simple_models[1].supports_tools is True + assert config.simple_models[1].max_context_chars == 1000 + + +def test_custom_question_only_needs_text_and_threshold(): + config = QueryRoutingConfig( + simple_models=["cheap"], + questions=[{"question": "是否需要最新信息?", "threshold": 0.1}], + ) + custom = config.questions[-1] + assert custom.id == "custom_1" + assert custom.instructions == "是否需要最新信息?" + assert custom.threshold == 0.1 + assert custom.type == "noul" + + +@pytest.mark.parametrize("description", ["", "x" * 2001, 1]) +def test_candidate_description_has_bounded_text_contract(description): + with pytest.raises(ValidationError): + QueryRoutingConfig( + simple_models=[{"name": "cheap", "capability_description": description}] + ) + + +def test_custom_question_accepts_leading_string(): + config = QueryRoutingConfig( + simple_models=["cheap"], questions=["是否需要实时数据?"] + ) + assert config.questions[-1].instructions == "是否需要实时数据?" + assert config.questions[-1].threshold == DEFAULT_QUERY_ROUTING_THRESHOLD + + +def test_reserved_question_id_is_rejected(): + with pytest.raises(ValidationError, match="reserved or duplicated"): + QueryRoutingConfig( + simple_models=["cheap"], + questions=[{"id": "is_complex", "question": "override", "threshold": 0}], + ) + + +def test_duplicate_custom_question_id_is_rejected(): + with pytest.raises(ValidationError, match="reserved or duplicated"): + QueryRoutingConfig( + simple_models=["cheap"], + questions=[ + {"id": "custom_x", "question": "a"}, + {"id": "custom_x", "question": "b"}, + ], + ) + + +def test_non_list_simple_models_is_left_for_validation(): + with pytest.raises(ValidationError): + QueryRoutingConfig(simple_models="cheap") + + +@pytest.mark.parametrize("value", ["raw", ["list"]]) +def test_non_mapping_input_is_left_for_validation(value): + with pytest.raises(ValidationError): + QueryRoutingConfig.model_validate(value) + + +def test_non_list_questions_is_left_for_validation(): + with pytest.raises(ValidationError): + QueryRoutingConfig(simple_models=["cheap"], questions="question") + + +def test_custom_question_accepts_routing_question_instance(): + config = QueryRoutingConfig( + simple_models=["cheap"], + questions=[ + RoutingQuestion( + id="custom_x", instructions="是否需要外部知识?", threshold=0.2 + ) + ], + ) + assert config.questions[-1].id == "custom_x" + assert config.questions[-1].threshold == 0.2 + + +def test_question_item_must_be_string_mapping_or_question(): + with pytest.raises(ValidationError, match="questions items must be"): + QueryRoutingConfig(simple_models=["cheap"], questions=[123]) + + +def test_valid_strong_model_is_kept(): + config = QueryRoutingConfig(simple_models=["cheap"], strong_model="strong") + assert config.strong_model == "strong" + + +def test_from_dict_builds_configuration(): + config = QueryRoutingConfig.from_dict({"simple_models": ["cheap"]}) + assert [model.name for model in config.simple_models] == ["cheap"] + + +@pytest.mark.parametrize("threshold", [-0.1, 1.1, float("nan"), float("inf")]) +def test_threshold_must_be_within_unit_interval(threshold): + with pytest.raises(ValidationError): + QueryRoutingConfig( + simple_models=["cheap"], + questions=[{"question": "q", "threshold": threshold}], + ) + + +@pytest.mark.parametrize("threshold", [0.0, 1.0]) +def test_threshold_accepts_unit_interval_boundaries(threshold): + config = QueryRoutingConfig( + simple_models=["cheap"], + questions=[{"question": "q", "threshold": threshold}], + ) + assert config.questions[-1].threshold == threshold + + +def test_duplicate_simple_models_are_rejected(): + with pytest.raises(ValidationError, match="duplicate"): + QueryRoutingConfig(simple_models=["cheap", "cheap"]) + + +def test_empty_simple_models_are_rejected(): + with pytest.raises(ValidationError): + QueryRoutingConfig(simple_models=[]) + + +def test_strong_model_must_not_be_blank(): + with pytest.raises(ValidationError, match="non-empty"): + QueryRoutingConfig(simple_models=["cheap"], strong_model=" ") + + +def test_shadow_mode_and_decision_timeout_bounds(): + config = QueryRoutingConfig(simple_models=["cheap"], mode="shadow") + assert config.mode == "shadow" + with pytest.raises(ValidationError): + QueryRoutingConfig(simple_models=["cheap"], decision_timeout=0) + + +def test_questions_hash_is_stable_and_ignores_model_names(): + first = QueryRoutingConfig(simple_models=["a"]) + second = QueryRoutingConfig(simple_models=["b"]) + assert first.questions_hash() == second.questions_hash() + + third = QueryRoutingConfig( + simple_models=["a"], questions=[{"question": "new gate"}] + ) + assert third.questions_hash() != first.questions_hash() + + +def test_load_from_explicit_config_yaml(tmp_path): + path = write_config( + tmp_path, + """ +query_routing: + simple_models: + - cheap-a + questions: + - question: 是否需要最新信息? + threshold: 0.2 +""", + ) + config = load_query_routing_config(path) + assert [model.name for model in config.simple_models] == ["cheap-a"] + assert config.questions[-1].id == "custom_1" + assert config.questions[-1].threshold == 0.2 + + +def test_load_if_present_returns_none_without_block(tmp_path): + path = write_config(tmp_path, "model:\n agent:\n name: x\n") + assert load_query_routing_config_if_present(path) is None + + +def test_auto_discovery_optional_returns_none_without_config_file( + tmp_path, monkeypatch +): + monkeypatch.chdir(tmp_path) + assert load_query_routing_config_if_present() is None + + +def test_auto_discovery_required_raises_without_config_file(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + with pytest.raises(QueryRoutingConfigError, match="no config.yaml"): + load_query_routing_config() + + +def test_empty_config_document_is_treated_as_empty_mapping(tmp_path): + path = write_config(tmp_path, "") + assert load_query_routing_config_if_present(path) is None + + +def test_load_if_present_rejects_missing_explicit_path(tmp_path): + with pytest.raises(QueryRoutingConfigError, match="cannot read"): + load_query_routing_config_if_present(tmp_path / "missing.yaml") + + +def test_load_requires_block(tmp_path): + path = write_config(tmp_path, "model: {}\n") + with pytest.raises(QueryRoutingConfigError, match="query_routing"): + load_query_routing_config(path) + + +def test_invalid_block_raises_with_path(tmp_path): + path = write_config(tmp_path, "query_routing:\n simple_models: []\n") + with pytest.raises(QueryRoutingConfigError, match="invalid `query_routing`"): + load_query_routing_config(path) + + +def test_non_mapping_block_raises(tmp_path): + path = write_config(tmp_path, "query_routing: []\n") + with pytest.raises(QueryRoutingConfigError, match="mapping"): + load_query_routing_config(path) + + +def test_non_mapping_block_raises_when_optional(tmp_path): + path = write_config(tmp_path, "query_routing: []\n") + with pytest.raises(QueryRoutingConfigError, match="mapping"): + load_query_routing_config_if_present(path) + + +def test_invalid_yaml_raises(tmp_path): + path = write_config(tmp_path, "query_routing: [\n") + with pytest.raises(QueryRoutingConfigError, match="invalid YAML"): + load_query_routing_config(path) + + +def test_non_mapping_root_raises(tmp_path): + path = write_config(tmp_path, "- just\n- a list\n") + with pytest.raises(QueryRoutingConfigError, match="root must be a mapping"): + load_query_routing_config_if_present(path) + + +def test_unreadable_path_raises(tmp_path): + with pytest.raises(QueryRoutingConfigError, match="cannot read"): + load_query_routing_config(tmp_path / "nope" / "config.yaml") + + +def test_decision_block_defaults(tmp_path): + path = write_config(tmp_path, "query_routing:\n simple_models: [cheap]\n") + config = load_query_routing_config(path) + assert "decision" not in type(config).model_fields + assert config.decision_timeout == 2.0 + + +def test_independent_credentials_use_environment_and_remain_redacted(monkeypatch): + from veadk.configs.routing_configs import SimpleModelConfig, QueryRoutingConfigError + + monkeypatch.setenv("ROUTING_TEST_KEY", "test-only-value") + profile = SimpleModelConfig(name="cheap", api_key_env="ROUTING_TEST_KEY") + assert profile.api_key_value == "test-only-value" + monkeypatch.delenv("ROUTING_TEST_KEY") + with pytest.raises(QueryRoutingConfigError, match="missing credential"): + profile.api_key_value + with pytest.raises(ValueError, match="only one"): + SimpleModelConfig( + name="cheap", api_key="test-only", api_key_env="ROUTING_TEST_KEY" + ) + hidden = SimpleModelConfig(name="cheap", api_key="test-only") + assert ( + "test-only" not in repr(hidden) and "test-only" not in hidden.model_dump_json() + ) + + +@pytest.mark.parametrize( + "field", + [ + "model", + "api_key", + "api_base", + "fallbacks", + "previous_response_id", + "input", + "instructions", + ], +) +def test_profile_cannot_override_routing_contract_through_extra_config(field): + from veadk.configs.routing_configs import SimpleModelConfig + + with pytest.raises(ValueError, match="not allowed"): + SimpleModelConfig(name="cheap", extra_config={field: "override"}) + + +def test_builtin_threshold_calibration_validates_ids_and_bounds(): + cfg = QueryRoutingConfig( + simple_models=["cheap"], builtin_thresholds={"needs_tools": 0.4} + ) + assert {q.id: q.threshold for q in cfg.questions} == { + "is_complex": 0.5, + "needs_tools": 0.4, + "needs_long_context": 0.5, + "is_high_stakes": 0.05, + } + fresh = QueryRoutingConfig(simple_models=["cheap"]) + assert next(q.threshold for q in fresh.questions if q.id == "needs_tools") == 0.5 + with pytest.raises(ValueError, match="unknown"): + QueryRoutingConfig(simple_models=["cheap"], builtin_thresholds={"unknown": 0.4}) + for threshold in (-0.1, 1.1, float("nan")): + with pytest.raises(ValueError): + QueryRoutingConfig( + simple_models=["cheap"], builtin_thresholds={"needs_tools": threshold} + ) diff --git a/tests/routing/test_execution_contract.py b/tests/routing/test_execution_contract.py new file mode 100644 index 000000000..ff2b629b5 --- /dev/null +++ b/tests/routing/test_execution_contract.py @@ -0,0 +1,494 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Failure, replay and lifetime contracts at the real ADK boundary.""" + +import asyncio +import copy +from types import SimpleNamespace + +import httpx +import pytest +from google.adk.agents import LlmAgent, SequentialAgent, ParallelAgent +from google.adk.events import Event +from google.adk.models.llm_response import LlmResponse +from google.adk.models.llm_request import LlmRequest +from google.adk.tools import AgentTool +from google.genai import types +from pydantic import BaseModel +from volcenginesdkarkruntime._exceptions import ArkBadRequestError + +from tests.routing.test_turn_contract import ( + Judge, + RecordingArk, + RecordingLlm, + routed, + setup, + turn, +) +from veadk.routing.context import get_turn_context +from veadk.routing.execution import _RECORD, _REQUEST, _select, _portable_contents + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["off", "shadow", "enforce"]) +async def test_responses_tool_order_is_strictly_opt_in(mode): + from veadk.routing.context import _TURN + + model = RecordingArk(model="openai/model") + content = types.Content( + role="model", + parts=[ + types.Part(text="Preparing lookup"), + types.Part( + function_call=types.FunctionCall(name="lookup", id="call-1", args={}) + ), + types.Part( + function_call=types.FunctionCall(name="lookup", id="call-2", args={}) + ), + ], + ) + original = content.model_copy(deep=True) + binding = ( + None + if mode == "off" + else SimpleNamespace(selection=SimpleNamespace(shadow=mode == "shadow")) + ) + token = _TURN.set(binding) + try: + results = types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + name="lookup", id=identifier, response={"value": "ORCHID"} + ) + ) + for identifier in ("call-1", "call-2") + ], + ) + async for _ in model.generate_content_async( + LlmRequest(contents=[content, results]) + ): + pass + finally: + _TURN.reset(token) + items = model.calls[0]["input"] + expected = ( + ["message", "function_call", "function_call"] + if mode == "enforce" + else ["function_call", "function_call", "message"] + ) + expected += ["function_call_output", "function_call_output"] + assert [item["type"] for item in items] == expected + assert [item["call_id"] for item in items if item["type"] == "function_call"] == [ + "call-1", + "call-2", + ] + assert content == original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cache", [False, True]) +async def test_responses_replay_keeps_explanation_before_tool_calls(monkeypatch, cache): + from veadk.models.ark_llm import request_reorganization_by_ark + + class MixedArk(RecordingArk): + async def generate_content_via_responses(self, responses_args, stream=False): + if self.calls: + async for response in super().generate_content_via_responses( + responses_args, stream + ): + yield response + return + self.calls.append( + request_reorganization_by_ark( + copy.deepcopy(responses_args), self.enable_responses_cache + ) + ) + yield LlmResponse( + content=types.Content( + role="model", + parts=[ + types.Part(text="I will retrieve the code."), + types.Part( + function_call=types.FunctionCall( + name="lookup", id="call-original", args={} + ) + ), + ], + ), + interaction_id="mixed-first", + ) + + def lookup() -> dict: + return {"code": "ORCHID"} + + cheap = MixedArk(model="openai/cheap", enable_responses_cache=cache) + strong = RecordingArk(model="openai/strong", enable_responses_cache=cache) + judge = Judge((0.01, 0.9, 0.01)) + root = routed( + monkeypatch, judge, strong, cheap, tools=[lookup], enable_responses=True + ) + runner, session = await setup(root) + try: + await turn(runner, session, "Retrieve the code") + await turn(runner, session, "Recall the code") + await turn(runner, session, "Recall it again") + assert len(judge.calls) == 3 + assert len(cheap.calls) == 3 and len(strong.calls) == 1 + assert cheap.calls[1].get("previous_response_id") == ( + "mixed-first" if cache else None + ) + for request in (strong.calls[0], cheap.calls[2]): + assert not request.get("previous_response_id") + items = request["input"] + index = next( + i for i, item in enumerate(items) if item.get("type") == "function_call" + ) + assert items[index - 1]["role"] == "assistant" + assert "retrieve the code" in str(items[index - 1]) + assert items[index + 1]["type"] == "function_call_output" + assert ( + items[index]["call_id"] + == items[index + 1]["call_id"] + == "call-original" + ) + assert "ORCHID" in items[index + 1]["output"] + finally: + await runner.close() + + +class ExpiringArk(RecordingArk): + async def generate_content_via_responses(self, responses_args, stream=False): + if responses_args.get("previous_response_id"): + self.calls.append(copy.deepcopy(responses_args)) + raise ArkBadRequestError( + "expired", + response=httpx.Response( + 400, request=httpx.Request("POST", "https://example.test/responses") + ), + body={"code": "InvalidParameter.PreviousResponseNotFound"}, + request_id="test-request", + ) + async for response in super().generate_content_via_responses( + responses_args, stream + ): + yield response + + +@pytest.mark.asyncio +async def test_expired_responses_id_retries_same_model_with_full_history(monkeypatch): + cheap = ExpiringArk(model="openai/cheap") + strong = RecordingLlm(model="strong") + judge = Judge() + root = routed(monkeypatch, judge, strong, cheap) + runner, session = await setup(root) + try: + await turn(runner, session, "ORCHID") + await turn(runner, session, "MAPLE") + assert len(cheap.calls) == 3 + assert cheap.calls[1]["previous_response_id"] + assert not cheap.calls[2].get("previous_response_id") + assert "ORCHID" in str(cheap.calls[2]["input"]) + assert "MAPLE" in str(cheap.calls[2]["input"]) + assert not strong.calls + assert len(judge.calls) == 2 + finally: + await runner.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["instruction", "history"]) +async def test_changed_effective_request_rebuilds_responses_chain(monkeypatch, change): + cheap = RecordingArk(model="openai/cheap") + calls = 0 + + async def callback(callback_context, llm_request): + nonlocal calls + calls += 1 + if calls == 2 and change == "history": + llm_request.contents[0] = types.Content( + role="user", parts=[types.Part(text="REVISED")] + ) + llm_request.model = "callback-must-not-switch-model" + + root = routed( + monkeypatch, + Judge(), + RecordingLlm(model="strong"), + cheap, + instruction="Remember codes", + before_model_callback=callback, + ) + runner, session = await setup(root) + try: + await turn(runner, session, "ORCHID") + if change == "instruction": + root.instruction = "Remember codes and reply briefly" + await turn(runner, session, "MAPLE") + assert not cheap.calls[1].get("previous_response_id") + assert cheap.calls[1]["model"] == "cheap" + assert calls == 2 + if change == "history": + assert "REVISED" in str(cheap.calls[1]["input"]) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_resume_reuses_persisted_selection_and_rejects_configuration_change( + monkeypatch, +): + judge = Judge() + cheap = RecordingLlm(model="cheap") + root = routed(monkeypatch, judge, RecordingLlm(model="strong"), cheap) + runner, session = await setup(root) + try: + await turn(runner, session, "ORCHID") + saved = await runner.session_service.get_session( + app_name=session.app_name, user_id=session.user_id, session_id=session.id + ) + record = saved.state[_RECORD] + ctx = SimpleNamespace(session=saved, invocation_id=record["invocation_id"]) + binding = await _select(root, ctx, root.model) + assert binding.model is cheap + assert len(judge.calls) == 1 + cheap.model = "changed" + with pytest.raises(ValueError, match="configuration changed"): + await _select(root, ctx, root.model) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_cancellation_closes_classifier_and_clears_turn_scope(monkeypatch): + started = asyncio.Event() + closed = [] + + class WaitingJudge(Judge): + async def judge(self, **kwargs): + started.set() + await asyncio.Event().wait() + + async def aclose(self): + closed.append(True) + + root = routed( + monkeypatch, + WaitingJudge(), + RecordingLlm(model="strong"), + RecordingLlm(model="cheap"), + ) + runner, session = await setup(root) + try: + task = asyncio.create_task(turn(runner, session, "ORCHID")) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert root.model.active_runs == 0 + assert closed == [True] + assert get_turn_context() is None and _REQUEST.get() is None + assert not root.model.store._inflight + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_late_capability_growth_fails_without_switching(monkeypatch): + async def enlarge(callback_context, llm_request): + llm_request.contents.append( + types.Content(role="user", parts=[types.Part(text="X" * 33000)]) + ) + + cheap, strong = RecordingLlm(model="cheap"), RecordingLlm(model="strong") + judge = Judge() + root = routed(monkeypatch, judge, strong, cheap, before_model_callback=enlarge) + runner, session = await setup(root) + try: + with pytest.raises(ValueError, match="context_too_large"): + await turn(runner, session, "small query") + assert len(judge.calls) == 1 + assert not cheap.calls and not strong.calls + assert get_turn_context() is None + finally: + await runner.close() + + +def test_foreign_reasoning_removed_without_losing_tool_pair_or_mutating_session(): + items = [ + types.Content( + role="model", + parts=[ + types.Part( + text="opaque reasoning", + thought=True, + thought_signature=b"signature", + ), + types.Part( + function_call=types.FunctionCall( + name="lookup", id="call-1", args={} + ), + thought=True, + thought_signature=b"tool-signature", + ), + types.Part(text="business answer"), + ], + ), + types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + name="lookup", id="call-1", response={"result": "ORCHID"} + ) + ) + ], + ), + ] + event = Event( + author="root", + content=items[0], + custom_metadata={"veadk_query_routing": {"target": "a"}}, + ) + ctx = SimpleNamespace(session=SimpleNamespace(events=[event])) + same = _portable_contents(ctx, items, "a") + foreign = _portable_contents(ctx, items, "b") + assert same[0].parts[0].thought_signature == b"signature" + assert "opaque reasoning" not in str(foreign) + assert ( + foreign[0].parts[0].function_call.id + == foreign[1].parts[0].function_response.id + == "call-1" + ) + assert all(p.thought_signature is None for c in foreign for p in c.parts) + assert "business answer" in str(foreign) and "ORCHID" in str(foreign) + assert items[0].parts[0].thought_signature == b"signature" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("workflow_type", [SequentialAgent, ParallelAgent]) +async def test_agent_tool_workflow_inherits_root_model(monkeypatch, workflow_type): + from tests.routing.test_turn_contract import DelegatingLlm + + judge = Judge() + unused, cheap = RecordingLlm(model="unused"), DelegatingLlm(model="cheap") + workflow = workflow_type( + name="workflow", + sub_agents=[ + LlmAgent(name="first", model=unused), + LlmAgent(name="second", model=unused), + ], + ) + root = routed(monkeypatch, judge, unused, cheap, tools=[AgentTool(workflow)]) + runner, session = await setup(root) + try: + await turn(runner, session, "ORCHID") + assert len(cheap.calls) == 4 + assert len(judge.calls) == 1 and not unused.calls + assert len(set(cheap.scopes)) == 1 + finally: + await runner.close() + + +def test_custom_native_subclass_rejected_instead_of_losing_its_behavior(monkeypatch): + class Custom(LlmAgent): + pass + + with pytest.raises(ValueError, match="custom ADK agents"): + routed( + monkeypatch, + Judge(), + RecordingLlm(model="strong"), + RecordingLlm(model="cheap"), + sub_agents=[Custom(name="custom", model="unused")], + ) + + +@pytest.mark.asyncio +async def test_output_schema_and_tool_description_are_classified(monkeypatch): + class Output(BaseModel): + code: str + + def lookup_project(code: str) -> str: + """Look up archived project codes.""" + return code + + cheap, strong = RecordingLlm(model="cheap"), RecordingLlm(model="strong") + judge = Judge() + root = routed( + monkeypatch, + judge, + strong, + cheap, + tools=[lookup_project], + global_instruction="Archive policy applies", + ) + runner, session = await setup(root) + try: + await turn(runner, session, "ORCHID") + assert "archived project" in str(judge.calls) + assert "Archive policy" in str(judge.calls) + root.output_schema = Output + await turn(runner, session, "format code") + assert len(judge.calls) == 1 + assert len(strong.calls) == 1 + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_callback_cannot_mutate_bound_model_configuration(monkeypatch): + cheap = RecordingLlm(model="cheap") + + async def mutate(callback_context, llm_request): + cheap.model = "other" + + root = routed( + monkeypatch, + Judge(), + RecordingLlm(model="strong"), + cheap, + before_model_callback=mutate, + ) + runner, session = await setup(root) + try: + with pytest.raises(ValueError, match="configuration changed during the turn"): + await turn(runner, session, "hello") + assert not cheap.calls + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_unresolved_dynamic_tools_choose_strong(monkeypatch): + from google.adk.tools.base_toolset import BaseToolset + + class Toolset(BaseToolset): + async def get_tools(self, readonly_context=None): + return [] + + async def close(self): + pass + + judge = Judge() + strong, cheap = RecordingLlm(model="strong"), RecordingLlm(model="cheap") + root = routed(monkeypatch, judge, strong, cheap, tools=[Toolset()]) + runner, session = await setup(root) + try: + await turn(runner, session, "hello") + assert len(strong.calls) == 1 and not cheap.calls + assert not judge.calls + finally: + await runner.close() diff --git a/tests/routing/test_factory.py b/tests/routing/test_factory.py new file mode 100644 index 000000000..4173b5289 --- /dev/null +++ b/tests/routing/test_factory.py @@ -0,0 +1,326 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import AsyncGenerator +from types import SimpleNamespace + +import pytest +from google.adk.models.base_llm import BaseLlm +from google.adk.models.llm_response import LlmResponse +from google.genai import types + +from veadk.decision import ( + DecisionConfig, + DecisionConfigError, + resolve_decision_connection, +) +from veadk.configs.routing_configs import ( + QueryRoutingConfig, + QueryRoutingConfigError, +) +from veadk.models.ark_llm import ArkLlm +from veadk.models.retrying_lite_llm import RetryingLiteLlm +from veadk.routing import factory +from veadk.routing.llm import RoutingLlm + +ARK_BASE = "https://ark.cn-beijing.volces.com/api/v3" +SYSTEMONE = "https://ark.cn-beijing.volces.com/api/v1/systemone" + + +class FakeStrongLlm(BaseLlm): + model: str = "strong" + + async def generate_content_async( + self, llm_request, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + yield LlmResponse( + content=types.Content(role="model", parts=[types.Part(text="strong")]) + ) + + +class FakeJudge: + async def judge(self, *, state, questions): + return SimpleNamespace(answers={}) + + +def fake_agent(**overrides): + values = { + "model_name": "strong", + "model": FakeStrongLlm(), + "model_provider": "ark", + "model_api_base": ARK_BASE, + "model_api_key": "agent-key", + "model_extra_config": {}, + "enable_responses": False, + "enable_responses_cache": True, + "runtime": "adk", + } + values.update(overrides) + return SimpleNamespace(**values) + + +def config(**overrides): + values = {"simple_models": ["cheap"]} + values.update(overrides) + return QueryRoutingConfig.model_validate(values) + + +def test_false_disables_routing(monkeypatch): + def unexpected(): + raise AssertionError("must not load configuration") + + monkeypatch.setattr(factory, "load_query_routing_config_if_present", unexpected) + monkeypatch.setattr(factory, "load_query_routing_config", unexpected) + assert factory.resolve_query_routing_config(False) is None + + +def test_none_auto_detects(monkeypatch): + sentinel = config() + monkeypatch.setattr( + factory, "load_query_routing_config_if_present", lambda: sentinel + ) + assert factory.resolve_query_routing_config(None) is sentinel + + +def test_true_requires_configuration(monkeypatch): + sentinel = config() + monkeypatch.setattr(factory, "load_query_routing_config", lambda: sentinel) + assert factory.resolve_query_routing_config(True) is sentinel + + +def test_config_instance_is_returned_unchanged(): + value = config() + assert factory.resolve_query_routing_config(value) is value + + +def test_mapping_is_validated(): + value = factory.resolve_query_routing_config( + {"simple_models": ["cheap"], "strong_model": "strong"} + ) + assert value.simple_models[0].name == "cheap" + assert value.strong_model == "strong" + + +def test_invalid_value_is_rejected(): + with pytest.raises(QueryRoutingConfigError, match="must be None"): + factory.resolve_query_routing_config("config.yaml") + + +def test_responses_cache_is_supported(): + agent = fake_agent(enable_responses=True, enable_responses_cache=True) + router, _ = factory.build_query_routing(config(), agent, judge=FakeJudge()) + assert router.runtimes[0].llm.enable_responses_cache is True + + +def test_build_uses_the_agent_model_as_strong_tier(): + agent = fake_agent() + router, callback = factory.build_query_routing(config(), agent, judge=FakeJudge()) + assert isinstance(router, RoutingLlm) + assert router.strong_model is agent.model + assert router.model == "strong" + assert [runtime.config.name for runtime in router.runtimes] == ["cheap"] + assert callable(callback) + + +def test_build_constructs_a_decision_client_from_ark_connection(monkeypatch): + captured = {} + + class Recorder: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(factory, "DecisionClient", Recorder) + factory.build_query_routing(config(), fake_agent()) + + assert captured == { + "endpoint": SYSTEMONE, + "api_key": "agent-key", + "model": "doubao-seed-2-1-lite-260915", + "timeout": 2.0, + } + + +def test_build_uses_explicit_strong_model(monkeypatch): + built = FakeStrongLlm(model="explicit") + monkeypatch.setattr(factory, "_build_model", lambda agent, name: built) + + router, _ = factory.build_query_routing( + config(strong_model="stronger"), fake_agent(), judge=FakeJudge() + ) + assert router.strong_model is built + + +def test_strong_model_defaults_to_agent_model(): + agent = fake_agent() + assert factory._strong_model(config(), agent) is agent.model + + +def test_strong_model_matching_primary_name_reuses_agent_model(): + agent = fake_agent(model_name="strong") + assert factory._strong_model(config(strong_model="strong"), agent) is agent.model + + +def test_strong_model_differs_from_primary(monkeypatch): + built = FakeStrongLlm(model="built") + calls = [] + + def build(agent, name): + calls.append((agent, name)) + return built + + monkeypatch.setattr(factory, "_build_model", build) + agent = fake_agent() + assert factory._strong_model(config(strong_model="stronger"), agent) is built + assert calls == [(agent, "stronger")] + + +def test_primary_model_name_accepts_a_list(): + assert factory._primary_model_name(fake_agent(model_name=["a", "b"])) == "a" + assert factory._primary_model_name(fake_agent(model_name=[])) is None + assert factory._primary_model_name(fake_agent(model_name="a")) == "a" + + +def test_build_model_prefixes_provider_for_plain_names(): + model = factory._build_model(fake_agent(), "cheap") + assert isinstance(model, RetryingLiteLlm) + assert model.model == "ark/cheap" + + +def test_build_model_keeps_qualified_names(): + model = factory._build_model(fake_agent(), "openai/cheap") + assert model.model == "openai/cheap" + + +def test_build_model_uses_ark_responses_without_cache(): + model = factory._build_model(fake_agent(enable_responses=True), "cheap") + assert isinstance(model, ArkLlm) + assert model.model == "ark/cheap" + assert model.enable_responses_cache is True + + +def resolve_for_agent(config, agent): + return resolve_decision_connection( + config, agent.model_api_base, agent.model_api_key + ) + + +def test_decision_connection_uses_explicit_url_and_key(): + agent = fake_agent(model_api_base="https://api.openai.com/v1", model_api_key="x") + decision = DecisionConfig( + **{ + "api_url": "https://ark.cn-beijing.volces.com/api/v1/systemone", + "api_key": "explicit", + } + ) + assert resolve_for_agent(decision, agent) == ( + "https://ark.cn-beijing.volces.com/api/v1/systemone", + "explicit", + ) + + +def test_decision_connection_refuses_cross_provider_endpoint(): + decision = DecisionConfig(**{"api_url": "https://example.com/systemone"}) + with pytest.raises(DecisionConfigError, match="not an Ark endpoint"): + resolve_for_agent(decision, fake_agent()) + + +def test_decision_connection_requires_ark_model_base_for_ark_url(): + decision = DecisionConfig(**{"api_url": SYSTEMONE}) + agent = fake_agent(model_api_base="https://api.openai.com/v1") + with pytest.raises(DecisionConfigError, match="fallback Ark API base"): + resolve_for_agent(decision, agent) + + +def test_decision_connection_requires_a_key_for_ark_url(): + decision = DecisionConfig(**{"api_url": SYSTEMONE}) + agent = fake_agent(model_api_key="") + with pytest.raises(DecisionConfigError, match="API key is not configured"): + resolve_for_agent(decision, agent) + + +def test_decision_connection_reuses_ark_key_for_ark_url(): + decision = DecisionConfig(**{"api_url": SYSTEMONE}) + assert resolve_for_agent(decision, fake_agent()) == ( + SYSTEMONE, + "agent-key", + ) + + +def test_decision_connection_rejects_non_ark_model_base_without_url(): + agent = fake_agent(model_api_base="https://api.openai.com/v1") + with pytest.raises(DecisionConfigError, match="cannot derive"): + resolve_for_agent(DecisionConfig(), agent) + + +def test_decision_connection_derives_systemone_from_ark_base(): + assert resolve_for_agent(DecisionConfig(), fake_agent()) == ( + SYSTEMONE, + "agent-key", + ) + + +def test_decision_connection_prefers_explicit_key_when_deriving(): + decision = DecisionConfig(**{"api_key": "explicit"}) + assert resolve_for_agent(decision, fake_agent()) == ( + SYSTEMONE, + "explicit", + ) + + +def test_decision_connection_requires_a_key_when_deriving(): + agent = fake_agent(model_api_key="") + with pytest.raises(DecisionConfigError, match="API key is not configured"): + resolve_for_agent(DecisionConfig(), agent) + + +def test_routing_rejects_underlying_model_fallbacks(): + model = ArkLlm(model="openai/strong", fallbacks=["openai/other"]) + with pytest.raises(QueryRoutingConfigError, match="entire turn"): + factory.build_query_routing( + config(), fake_agent(model=model), judge=FakeJudge() + ) + + +def test_independent_endpoint_never_inherits_credentials_or_provider_extras(): + from veadk.configs.routing_configs import SimpleModelConfig + + spec = SimpleModelConfig(name="cheap", api_base="https://other.example/v1") + with pytest.raises(QueryRoutingConfigError, match="independent api_base"): + factory._build_model(fake_agent(), spec) + spec = SimpleModelConfig( + name="cheap", + api_base="https://other.example/v1", + api_key="other-only", + api="chat_completions", + ) + model = factory._build_model( + fake_agent(model_extra_config={"extra_headers": {"provider-only": "yes"}}), spec + ) + assert model._additional_args["api_key"] == "other-only" + assert "provider-only" not in str(model._additional_args) + + +@pytest.mark.parametrize( + "extra", + [ + {"previous_response_id": "stale-id"}, + {"extra_body": {"model": "other-model"}}, + ], +) +def test_parent_transport_options_cannot_bypass_fixed_routing(extra): + model = ArkLlm(model="openai/strong", **extra) + with pytest.raises(QueryRoutingConfigError, match="override"): + factory.build_query_routing( + config(), fake_agent(model=model), judge=FakeJudge() + ) diff --git a/tests/routing/test_llm.py b/tests/routing/test_llm.py new file mode 100644 index 000000000..fb41a17c0 --- /dev/null +++ b/tests/routing/test_llm.py @@ -0,0 +1,169 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import AsyncGenerator + +import pytest +from google.adk.models.base_llm import BaseLlm +from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse +from google.genai import types +from pydantic import Field + +from veadk.configs.routing_configs import SimpleModelConfig +from veadk.routing.context import RouteSelection, set_route_selection +from veadk.routing.llm import RoutingLlm +from veadk.routing.registry import ModelRegistry, SimpleModelRuntime +from veadk.routing.store import RouteStore + + +class FakeLlm(BaseLlm): + model: str + fail_before: bool = False + fail_after: bool = False + error: Exception | None = Field(default=None, exclude=True) + seen_models: list[str] = Field(default_factory=list) + calls: int = 0 + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + self.calls += 1 + self.seen_models.append(llm_request.model) + if self.error is not None: + raise self.error + if self.fail_before: + raise RuntimeError("before") + yield LlmResponse( + content=types.Content(role="model", parts=[types.Part(text=self.model)]) + ) + if self.fail_after: + raise RuntimeError("after") + + +def build(fail_before=False, fail_after=False, error=None): + strong = FakeLlm(model="strong") + simple = FakeLlm( + model="simple", fail_before=fail_before, fail_after=fail_after, error=error + ) + runtime = SimpleModelRuntime(config=SimpleModelConfig(name="simple"), llm=simple) + registry = ModelRegistry(runtimes=[runtime]) + store = RouteStore() + router = RoutingLlm( + model="strong", + strong_model=strong, + runtimes=[runtime], + registry=registry, + store=store, + ) + return router, strong, simple, store + + +def request(): + return LlmRequest( + model="strong", + contents=[types.Content(role="user", parts=[types.Part(text="hello")])], + ) + + +async def collect(router, req): + return [item async for item in router.generate_content_async(req)] + + +@pytest.mark.asyncio +async def test_no_selection_uses_strong_model(): + router, strong, simple, _ = build() + set_route_selection(None) + responses = await collect(router, request()) + assert responses[0].content.parts[0].text == "strong" + assert strong.calls == 1 + assert simple.calls == 0 + + +@pytest.mark.asyncio +async def test_simple_selection_uses_configured_model(): + router, strong, simple, _ = build() + set_route_selection( + RouteSelection( + route="simple", + reason="ok", + invocation_id="inv", + model_name="simple", + ) + ) + responses = await collect(router, request()) + assert responses[0].content.parts[0].text == "simple" + assert simple.calls == 1 + assert strong.calls == 0 + assert simple.seen_models == ["simple"] + + +@pytest.mark.asyncio +async def test_unknown_target_is_rejected_without_switching(): + router, strong, _, _ = build() + set_route_selection( + RouteSelection( + route="simple", reason="ok", invocation_id="inv", model_name="missing" + ) + ) + with pytest.raises(ValueError, match="no longer configured"): + await collect(router, request()) + assert strong.calls == 0 + + +@pytest.mark.asyncio +async def test_pre_chunk_failure_propagates_without_switching(): + router, strong, simple, _ = build(fail_before=True) + set_route_selection( + RouteSelection( + route="simple", reason="ok", invocation_id="inv", model_name="simple" + ) + ) + with pytest.raises(RuntimeError, match="before"): + await collect(router, request()) + assert strong.calls == 0 + assert simple.calls == 1 + + +@pytest.mark.asyncio +async def test_post_chunk_failure_is_not_retried(): + router, strong, simple, _ = build(fail_after=True) + set_route_selection( + RouteSelection( + route="simple", + reason="ok", + invocation_id="inv", + model_name="simple", + ) + ) + with pytest.raises(RuntimeError, match="after"): + await collect(router, request()) + assert strong.calls == 0 + + +@pytest.mark.asyncio +async def test_shadow_route_uses_strong_even_with_model_name(): + router, strong, simple, _ = build() + set_route_selection( + RouteSelection( + route="strong", + reason="shadow", + invocation_id="inv", + model_name="simple", + shadow=True, + ) + ) + responses = await collect(router, request()) + assert responses[0].content.parts[0].text == "strong" + assert simple.calls == 0 diff --git a/tests/routing/test_policy.py b/tests/routing/test_policy.py new file mode 100644 index 000000000..ea33dc90c --- /dev/null +++ b/tests/routing/test_policy.py @@ -0,0 +1,146 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import pytest + +from veadk.configs.routing_configs import QueryRoutingConfig, RoutingQuestion +from veadk.routing.policy import ConservativePolicy + + +def answers(**values): + return {name: {"type": "noul", "noul": value} for name, value in values.items()} + + +def policy(): + return ConservativePolicy(QueryRoutingConfig(simple_models=["cheap"]).questions) + + +def all_low(**overrides): + values = { + "is_complex": 0.01, + "needs_tools": 0.01, + "needs_long_context": 0.01, + "is_high_stakes": 0.01, + } + values.update(overrides) + return answers(**values) + + +def test_questions_property_returns_the_gates(): + gates = policy().questions + assert [question.id for question in gates] == [ + "is_complex", + "needs_tools", + "needs_long_context", + "is_high_stakes", + ] + + +def test_all_questions_below_threshold_routes_simple(): + decision = policy().evaluate(all_low()) + assert decision.route == "simple" + assert decision.reason == "all_thresholds_passed" + assert decision.probabilities["is_complex"] == 0.01 + + +@pytest.mark.parametrize( + ("question", "value"), + [ + ("is_complex", 0.500001), + ("needs_tools", 0.500001), + ("needs_long_context", 0.500001), + ("is_high_stakes", 0.050001), + ], +) +def test_any_question_above_threshold_routes_strong(question, value): + decision = policy().evaluate(all_low(**{question: value})) + assert decision.route == "strong" + assert decision.reason == f"{question}_above_threshold" + + +@pytest.mark.parametrize( + "question", ["is_complex", "needs_tools", "needs_long_context"] +) +@pytest.mark.parametrize("value", [0.49, 0.5]) +def test_default_threshold_accepts_probability_up_to_half(question, value): + decision = policy().evaluate(all_low(**{question: value})) + assert decision.route == "simple" + + +def test_high_stakes_threshold_is_inclusive(): + decision = policy().evaluate(all_low(is_high_stakes=0.05)) + assert decision.route == "simple" + + +@pytest.mark.parametrize( + ("threshold", "value", "route"), + [ + (0.0, 0.0, "simple"), + (0.0, 0.000001, "strong"), + (1.0, 1.0, "simple"), + (0.5, 0.5, "simple"), + ], +) +def test_threshold_extremes(threshold, value, route): + gate = ConservativePolicy( + [RoutingQuestion(id="gate", instructions="是否复杂?", threshold=threshold)] + ) + decision = gate.evaluate({"gate": {"type": "noul", "noul": value}}) + assert decision.route == route + if route == "strong": + assert decision.reason == "gate_above_threshold" + + +def test_empty_answers_fail_closed(): + decision = policy().evaluate({}) + assert decision.route == "strong" + assert decision.reason == "invalid_answer:is_complex" + + +def test_missing_answer_fails_closed(): + values = all_low() + values.pop("needs_tools") + decision = policy().evaluate(values) + assert decision.route == "strong" + assert decision.reason == "invalid_answer:needs_tools" + + +def test_non_noul_type_fails_closed(): + values = all_low() + values["is_complex"] = {"type": "score", "score": 0.1} + decision = policy().evaluate(values) + assert decision.route == "strong" + assert decision.reason == "invalid_answer:is_complex" + + +@pytest.mark.parametrize("value", [-0.01, 1.01, math.nan, math.inf, True, "0.1"]) +def test_invalid_probability_fails_closed(value): + values = all_low() + values["is_complex"] = {"type": "noul", "noul": value} + decision = policy().evaluate(values) + assert decision.route == "strong" + assert decision.reason == "invalid_answer:is_complex" + + +def test_non_mapping_answers_fail_closed(): + decision = policy().evaluate([]) + assert decision.route == "strong" + assert decision.reason == "invalid_answers" + + +def test_empty_question_set_is_rejected(): + with pytest.raises(ValueError, match="at least one"): + ConservativePolicy([]) diff --git a/tests/routing/test_state.py b/tests/routing/test_state.py new file mode 100644 index 000000000..f5e3b8b4d --- /dev/null +++ b/tests/routing/test_state.py @@ -0,0 +1,128 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from google.adk.models.llm_request import LlmRequest +from google.genai import types + +from veadk.content_safety.state import DEFUSE_PREFIX, DEFUSE_SUFFIX +from veadk.routing.state import ( + QueryStateBuilder, + IncompleteRoutingState, + request_has_media, +) + +PNG_BYTES = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def make_request(*messages): + contents = [ + types.Content(role=role, parts=[types.Part(text=text)]) + for role, text in messages + ] + return LlmRequest(model="provider/model", contents=contents) + + +def text_of(state): + return "\n".join(part.get("text", "") for part in state) + + +def test_recent_turns_are_labelled_and_included(): + request = make_request(("user", "first"), ("model", "answer"), ("user", "second")) + state = QueryStateBuilder(defuse=False).build_for_routing(request) + assert text_of(state) == "[user]\nfirst\n[model]\nanswer\n[user]\nsecond" + + +def test_turn_window_reports_incomplete_history(): + request = make_request( + ("user", "one"), + ("model", "two"), + ("user", "three"), + ("model", "four"), + ) + with pytest.raises(IncompleteRoutingState, match="history"): + QueryStateBuilder(defuse=False, max_turns=2).build_for_routing(request) + + +def test_empty_request_returns_empty_state(): + request = LlmRequest(model="provider/model", contents=[]) + assert QueryStateBuilder().build_for_routing(request) == [] + + +def test_content_without_parts_is_skipped(): + request = LlmRequest( + model="provider/model", + contents=[ + types.Content(role="user", parts=[]), + types.Content(role="user", parts=[types.Part(text="kept")]), + ], + ) + state = QueryStateBuilder(defuse=False).build_for_routing(request) + assert text_of(state) == "[user]\nkept" + + +def test_thought_parts_are_skipped(): + request = LlmRequest( + model="provider/model", + contents=[ + types.Content( + role="model", + parts=[types.Part(text="hidden", thought=True), types.Part(text="ok")], + ) + ], + ) + state = QueryStateBuilder(defuse=False).build_for_routing(request) + assert "hidden" not in text_of(state) + assert "ok" in text_of(state) + + +def test_defuse_wrapper_is_added_by_default(): + state = QueryStateBuilder().build_for_routing(make_request(("user", "hello"))) + assert state[0]["text"] == DEFUSE_PREFIX + assert state[-1]["text"] == DEFUSE_SUFFIX + + +def test_text_budget_reports_incomplete_state(): + request = make_request(("user", "x" * 500)) + with pytest.raises(IncompleteRoutingState, match="character budget"): + QueryStateBuilder(max_chars=100, defuse=False).build_for_routing(request) + + +def test_media_part_is_converted_to_image_url(): + request = LlmRequest( + model="provider/model", + contents=[ + types.Content( + role="user", + parts=[ + types.Part( + inline_data=types.Blob(mime_type="image/png", data=PNG_BYTES) + ) + ], + ) + ], + ) + state = QueryStateBuilder(defuse=False).build_for_routing(request) + assert request_has_media(request) is True + media = [part for part in state if part.get("type") == "image"] + assert len(media) == 1 + assert media[0]["image_url"].startswith("data:image/png;base64,") + + +def test_text_only_request_has_no_media(): + assert request_has_media(make_request(("user", "hello"))) is False diff --git a/tests/routing/test_store.py b/tests/routing/test_store.py new file mode 100644 index 000000000..9e29236c1 --- /dev/null +++ b/tests/routing/test_store.py @@ -0,0 +1,129 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio + +import pytest + +from veadk.routing.context import RouteSelection +from veadk.routing.store import RouteStore + + +def selection(reason="ok"): + return RouteSelection(route="simple", reason=reason, invocation_id="inv") + + +class Clock: + def __init__(self): + self.now = 1_000.0 + + def __call__(self): + return self.now + + +@pytest.mark.asyncio +async def test_concurrent_callers_share_one_classification(): + store = RouteStore() + calls = 0 + started = asyncio.Event() + release = asyncio.Event() + + async def factory(): + nonlocal calls + calls += 1 + started.set() + await release.wait() + return selection() + + first = asyncio.create_task(store.get_or_create("k", factory)) + await started.wait() + second = asyncio.create_task(store.get_or_create("k", factory)) + await asyncio.sleep(0) + release.set() + + results = await asyncio.gather(first, second) + assert calls == 1 + assert results[0] == results[1] + + +@pytest.mark.asyncio +async def test_cached_decision_is_reused_within_ttl(): + clock = Clock() + store = RouteStore(decision_ttl=10, clock=clock) + calls = 0 + + async def factory(): + nonlocal calls + calls += 1 + return selection() + + assert await store.get_or_create("k", factory) == selection() + clock.now += 9 + assert await store.get_or_create("k", factory) == selection() + assert calls == 1 + + +@pytest.mark.asyncio +async def test_expired_decision_is_reclassified(): + clock = Clock() + store = RouteStore(decision_ttl=10, clock=clock) + calls = 0 + + async def factory(): + nonlocal calls + calls += 1 + return selection(str(calls)) + + await store.get_or_create("k", factory) + clock.now += 11 + result = await store.get_or_create("k", factory) + assert calls == 2 + assert result.reason == "2" + + +@pytest.mark.asyncio +async def test_factory_failure_propagates_and_clears_inflight(): + store = RouteStore() + + async def failing(): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + await store.get_or_create("k", failing) + + async def ok(): + return selection() + + assert await store.get_or_create("k", ok) == selection() + + +def test_clear_drops_all_state(): + store = RouteStore() + store.clear() + assert not store._decisions + + +@pytest.mark.asyncio +async def test_many_entries_are_pruned_when_expired(): + clock = Clock() + store = RouteStore(decision_ttl=10, clock=clock) + + async def factory(): + return selection() + + for index in range(4_100): + await store.get_or_create(f"k{index}", factory) + clock.now += 100 + await store.get_or_create("fresh", factory) + assert len(store._decisions) == 1 # noqa: SLF001 diff --git a/tests/routing/test_tool_arguments.py b/tests/routing/test_tool_arguments.py new file mode 100644 index 000000000..0ccea798c --- /dev/null +++ b/tests/routing/test_tool_arguments.py @@ -0,0 +1,387 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tool argument contracts through the real ADK execution boundary.""" + +import copy +from dataclasses import replace +from types import SimpleNamespace + +import pytest +from google.adk.models.llm_response import LlmResponse +from google.genai import types +from google.adk.tools import FunctionTool, AgentTool +from google.adk.agents import LlmAgent +from veadk import Agent + +from tests.routing.test_turn_contract import Judge, RecordingLlm, routed, setup, turn +from veadk.routing.context import RouteSelection, TurnContext, _TURN +from veadk.routing.tool_arguments import validate_tool_arguments + + +class ToolCaller(RecordingLlm): + arguments: list[dict] + + async def generate_content_async(self, llm_request, stream=False): + self.calls.append(copy.deepcopy(llm_request.contents)) + index = len(self.calls) - 1 + part = ( + types.Part( + function_call=types.FunctionCall( + name="multiply", id=f"call-{index}", args=self.arguments[index] + ) + ) + if index < len(self.arguments) + else types.Part(text="complete") + ) + yield LlmResponse(content=types.Content(role="model", parts=[part])) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "arguments", + [ + [{"quantity": "7", "price": "19"}], + [{"quantity": "7.5", "price": 19}, {"quantity": 7, "price": 19}], + [{"quantity": True, "price": 19}, {"quantity": 7, "price": 19}], + [{"quantity": 7}, {"quantity": 7, "price": 19}], + ], +) +async def test_invalid_arguments_corrected_without_switching_or_side_effect( + monkeypatch, arguments +): + effects = [] + + def multiply(quantity: int, price: int) -> dict: + effects.append((quantity, price)) + return {"total": quantity * price} + + cheap = ToolCaller(model="cheap", arguments=arguments) + strong, judge = RecordingLlm(model="strong"), Judge() + root = routed(monkeypatch, judge, strong, cheap, tools=[multiply]) + runner, session = await setup(root) + try: + events = await turn(runner, session, "Calculate the total") + assert effects == [(7, 19)] + assert all(type(value) is int for value in effects[0]) + assert len(judge.calls) == 1 and not strong.calls + results = [ + p.function_response.response + for e in events + if e.content + for p in e.content.parts or [] + if p.function_response + ] + assert results[-1] == {"total": 133} + if len(arguments) > 1: + assert results[0]["error"] == "invalid_tool_arguments" + assert results[0]["retryable"] is True + finally: + await runner.close() + + +def binding(invocation="root-turn"): + return TurnContext( + RouteSelection("simple", "test", invocation), None, None, "target" + ) + + +def declared(schema=None, *, genai=None): + declaration = types.FunctionDeclaration( + name="tool", parameters_json_schema=schema, parameters=genai + ) + return SimpleNamespace(_get_declaration=lambda: declaration) + + +@pytest.mark.asyncio +async def test_guard_opt_out_and_absent_declarations(): + args = {"count": "3"} + assert await validate_tool_arguments(object(), args, object()) is None + assert args == {"count": "3"} + token = _TURN.set(binding()) + try: + assert await validate_tool_arguments(declared(), args, object()) is None + assert ( + await validate_tool_arguments( + SimpleNamespace(_get_declaration=lambda: None), args, object() + ) + is None + ) + finally: + _TURN.reset(token) + + +@pytest.mark.asyncio +async def test_shadow_does_not_change_tool_arguments_or_execute_validation(): + context = binding() + context = replace(context, selection=replace(context.selection, shadow=True)) + token = _TURN.set(context) + args = {"count": "7"} + try: + assert await validate_tool_arguments(object(), args, object()) is None + assert args == {"count": "7"} + assert not context.tool_argument_errors + finally: + _TURN.reset(token) + + +@pytest.mark.asyncio +async def test_nested_schema_references_arrays_and_nullable(): + schema = { + "type": "object", + "properties": { + "items": {"type": "array", "items": {"$ref": "#/$defs/Item"}}, + "note": {"type": ["string", "null"]}, + }, + "$defs": { + "Item": {"type": "object", "properties": {"count": {"type": "integer"}}} + }, + } + args = {"items": [{"count": "17"}], "note": None} + token = _TURN.set(binding()) + try: + assert ( + await validate_tool_arguments( + declared(schema), args, SimpleNamespace(state={}) + ) + is None + ) + assert args == {"items": [{"count": 17}], "note": None} + args = {"note": None} + genai = types.Schema( + type="OBJECT", + properties={"note": types.Schema(type="STRING", nullable=True)}, + ) + assert ( + await validate_tool_arguments( + declared(genai=genai), args, SimpleNamespace(state={}) + ) + is None + ) + finally: + _TURN.reset(token) + + +@pytest.mark.asyncio +async def test_schema_constants_are_data_not_rewritten_schema(): + tool = declared( + { + "type": "object", + "properties": { + "value": { + "const": {"type": "INTEGER", "$ref": "https://example.test/data"} + } + }, + } + ) + args = {"value": {"type": "INTEGER", "$ref": "https://example.test/data"}} + token = _TURN.set(binding()) + try: + assert ( + await validate_tool_arguments(tool, args, SimpleNamespace(state={})) is None + ) + assert args["value"]["type"] == "INTEGER" + finally: + _TURN.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "invalid", ["7.0", "01", "+7", "1e3", True, 1.2, None, "9" * 130] +) +async def test_ambiguous_or_lossy_coercion_is_rejected_without_argument_leak(invalid): + tool = declared({"type": "object", "properties": {"count": {"type": "integer"}}}) + token = _TURN.set(binding()) + try: + args = {"count": invalid} + result = await validate_tool_arguments(tool, args, SimpleNamespace(state={})) + assert result["error"] == "invalid_tool_arguments" + assert result["issues"] == [ + {"schema_path": ["properties", "count", "type"], "rule": "type"} + ] + assert args == {"count": invalid} + finally: + _TURN.reset(token) + + +@pytest.mark.asyncio +async def test_budget_shared_by_nested_tasks_persisted_for_resume_and_reset_next_turn(): + tool = declared({"type": "object", "required": ["count"]}) + state = {} + token = _TURN.set(binding()) + try: + await validate_tool_arguments(tool, {}, SimpleNamespace(state=state)) + await validate_tool_arguments(tool, {}, SimpleNamespace(state={})) + with pytest.raises(ValueError, match="budget exhausted"): + await validate_tool_arguments(tool, {}, SimpleNamespace(state=state)) + finally: + _TURN.reset(token) + token = _TURN.set(binding()) + try: + with pytest.raises(ValueError, match="budget exhausted"): + await validate_tool_arguments(tool, {}, SimpleNamespace(state=state)) + finally: + _TURN.reset(token) + token = _TURN.set(binding("next-turn")) + try: + result = await validate_tool_arguments(tool, {}, SimpleNamespace(state=state)) + assert result["corrections_remaining"] == 1 + finally: + _TURN.reset(token) + + +@pytest.mark.asyncio +async def test_remote_schema_reference_does_not_trigger_network(): + token = _TURN.set(binding()) + try: + with pytest.raises(ValueError, match="local references"): + await validate_tool_arguments( + declared({"$ref": "https://example.test/schema"}), + {}, + SimpleNamespace(state={}), + ) + finally: + _TURN.reset(token) + + +@pytest.mark.asyncio +async def test_invalid_arguments_have_bounded_correction_budget(monkeypatch): + effects = [] + + def multiply(quantity: int, price: int) -> dict: + effects.append(True) + return {"total": quantity * price} + + cheap = ToolCaller( + model="cheap", arguments=[{"quantity": "invalid", "price": 19}] * 4 + ) + root = routed( + monkeypatch, Judge(), RecordingLlm(model="strong"), cheap, tools=[multiply] + ) + runner, session = await setup(root) + try: + with pytest.raises( + ValueError, match="tool argument correction budget exhausted" + ): + await turn(runner, session, "Calculate") + assert not effects + assert len(cheap.calls) == 3 + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_opt_out_retains_original_tool_behavior(): + def multiply(quantity: int, price: int) -> dict: + return {"total": quantity * price} + + agent = Agent( + name="plain", + model=ToolCaller(model="test", arguments=[{"quantity": "7", "price": "19"}]), + model_api_key="test-only", + query_routing=False, + tools=[multiply], + ) + runner, session = await setup(agent) + try: + with pytest.raises(TypeError): + await turn(runner, session, "Calculate") + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_user_callback_and_confirmation_are_preserved(monkeypatch): + effects = [] + callbacks = [] + + def multiply(quantity: int, price: int) -> dict: + effects.append(True) + return {"total": quantity * price} + + def before_tool(tool, args, tool_context): + callbacks.append(dict(args)) + args["quantity"] = "7" + + tool = FunctionTool(multiply, require_confirmation=True) + root = routed( + monkeypatch, + Judge(), + RecordingLlm(model="strong"), + ToolCaller(model="cheap", arguments=[{"quantity": "bad", "price": "19"}]), + tools=[tool], + before_tool_callback=before_tool, + ) + runner, session = await setup(root) + try: + events = await turn(runner, session, "Calculate") + assert callbacks == [{"quantity": "bad", "price": "19"}] + assert not effects + assert any(e.actions.requested_tool_confirmations for e in events) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_native_agent_tool_inherits_generic_argument_validation(monkeypatch): + from tests.routing.test_turn_contract import DelegatingLlm + + effects = [] + + def multiply(quantity: int, price: int) -> dict: + effects.append((quantity, price)) + return {"total": quantity * price} + + class NestedCaller(DelegatingLlm): + async def generate_content_async(self, llm_request, stream=False): + if "multiply" in llm_request.tools_dict: + self.calls.append(copy.deepcopy(llm_request.contents)) + has_result = any( + p.function_response + for c in llm_request.contents + for p in c.parts or [] + ) + part = ( + types.Part(text="done") + if has_result + else types.Part( + function_call=types.FunctionCall( + name="multiply", + id="nested", + args={"quantity": "7", "price": "19"}, + ) + ) + ) + yield LlmResponse(content=types.Content(role="model", parts=[part])) + else: + async for response in super().generate_content_async( + llm_request, stream + ): + yield response + + child = LlmAgent(name="child", model=RecordingLlm(model="unused"), tools=[multiply]) + cheap = NestedCaller(model="cheap") + root = routed( + monkeypatch, + Judge(), + RecordingLlm(model="strong"), + cheap, + tools=[AgentTool(child)], + ) + runner, session = await setup(root) + try: + await turn(runner, session, "Calculate") + assert effects == [(7, 19)] + assert not child.model.calls + finally: + await runner.close() diff --git a/tests/routing/test_turn_contract.py b/tests/routing/test_turn_contract.py new file mode 100644 index 000000000..ac2482690 --- /dev/null +++ b/tests/routing/test_turn_contract.py @@ -0,0 +1,299 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contracts across real ADK turns; only provider IO is substituted.""" + +import copy +import asyncio +from types import SimpleNamespace + +import pytest +from google.adk.models.base_llm import BaseLlm +from google.adk.agents import LlmAgent +from google.adk.models.llm_response import LlmResponse +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.adk.tools import AgentTool +from google.genai import types +from pydantic import Field + +from veadk import Agent +from veadk.models.ark_llm import ArkLlm, request_reorganization_by_ark +from veadk.routing import factory +from veadk.routing.context import get_turn_context + + +class Judge: + def __init__(self, values=(0.01,)): + self.values = iter(values) + self.calls = [] + + async def judge(self, *, state, questions): + self.calls.append(state) + value = next(self.values, 0.01) + return SimpleNamespace( + answers={key: {"type": "noul", "noul": value} for key in questions}, + raw={"usage": {"input_tokens": 10}}, + ) + + async def aclose(self): + pass + + +class RecordingLlm(BaseLlm): + calls: list = Field(default_factory=list) + failure: bool = False + + async def generate_content_async(self, llm_request, stream=False): + self.calls.append(copy.deepcopy(llm_request.contents)) + if self.failure: + raise RuntimeError("provider failed") + yield LlmResponse( + content=types.Content(role="model", parts=[types.Part(text=self.model)]) + ) + + +class RecordingArk(ArkLlm): + calls: list = Field(default_factory=list) + + async def generate_content_via_responses(self, responses_args, stream=False): + payload = request_reorganization_by_ark( + copy.deepcopy(responses_args), self.enable_responses_cache + ) + self.calls.append(payload) + yield LlmResponse( + content=types.Content(role="model", parts=[types.Part(text=self.model)]), + interaction_id=f"{self.model}-{len(self.calls)}", + ) + + +async def turn(runner, session, text): + return [ + event + async for event in runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=types.Content(role="user", parts=[types.Part(text=text)]), + ) + ] + + +async def setup(agent): + service = InMemorySessionService() + runner = Runner(app_name="routing_contract", agent=agent, session_service=service) + session = await service.create_session(app_name="routing_contract", user_id="user") + return runner, session + + +def routed(monkeypatch, judge, strong, cheap, **kwargs): + monkeypatch.setattr(factory, "DecisionClient", lambda **_: judge) + monkeypatch.setattr(factory, "_build_model", lambda *args, **kw: cheap) + return Agent( + name="root", + model=strong, + model_api_key="test-only", + model_api_base="https://ark.cn-beijing.volces.com/api/v3", + query_routing={ + "mode": "enforce", + "simple_models": [{"name": cheap.model, "supports_tools": True}], + }, + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_next_turn_redecides_and_retains_same_session_history(monkeypatch): + judge = Judge((0.9, 0.01)) + strong, cheap = RecordingLlm(model="strong"), RecordingLlm(model="cheap") + agent = routed(monkeypatch, judge, strong, cheap) + runner, session = await setup(agent) + try: + await turn(runner, session, "Remember project ORCHID") + await turn(runner, session, "Continue with the project") + assert len(judge.calls) == 2 + assert len(strong.calls) == len(cheap.calls) == 1 + assert "ORCHID" in str(cheap.calls[0]) + assert agent.canonical_model is strong + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_selected_model_failure_never_changes_model(monkeypatch): + strong, cheap = ( + RecordingLlm(model="strong"), + RecordingLlm(model="cheap", failure=True), + ) + agent = routed(monkeypatch, Judge(), strong, cheap) + runner, session = await setup(agent) + try: + with pytest.raises(RuntimeError, match="provider failed"): + await turn(runner, session, "hello") + assert not strong.calls + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_responses_reuse_and_model_switch_rebuild_history(monkeypatch): + judge = Judge((0.01, 0.01, 0.9, 0.01)) + strong, cheap = ( + RecordingArk(model="openai/strong"), + RecordingArk(model="openai/cheap"), + ) + agent = routed(monkeypatch, judge, strong, cheap, enable_responses=True) + runner, session = await setup(agent) + try: + for text in ("ORCHID", "continue", "new requirement MAPLE", "recall both"): + await turn(runner, session, text) + assert len(judge.calls) == 4 + assert cheap.calls[0].get("previous_response_id") is None + assert cheap.calls[1]["previous_response_id"] == "openai/cheap-1" + assert strong.calls[0].get("previous_response_id") is None + assert cheap.calls[2].get("previous_response_id") is None + assert "ORCHID" in str(cheap.calls[2]["input"]) + assert "MAPLE" in str(cheap.calls[2]["input"]) + finally: + await runner.close() + + +def test_independent_model_endpoint_and_protocol(monkeypatch): + monkeypatch.setattr(factory, "DecisionClient", lambda **_: Judge()) + agent = Agent( + name="root", + model_name="strong", + model_api_key="root-only", + model_api_base="https://ark.cn-beijing.volces.com/api/v3", + query_routing={ + "simple_models": [ + { + "name": "cheap", + "api_base": "https://other.example/v1", + "api_key": "other-only", + "provider": "openai", + "api": "chat_completions", + } + ], + }, + ) + runtime = agent.model.registry.runtimes[0] + assert runtime.llm._additional_args["api_base"] == "https://other.example/v1" + assert runtime.llm._additional_args["api_key"] == "other-only" + + +class DelegatingLlm(RecordingLlm): + scopes: list = Field(default_factory=list) + + async def generate_content_async(self, llm_request, stream=False): + self.calls.append(copy.deepcopy(llm_request.contents)) + self.scopes.append(get_turn_context().selection.invocation_id) + # Each nested Agent sees only its own task/tool history. + has_result = any( + p.function_response for c in llm_request.contents for p in c.parts or [] + ) + tools = llm_request.tools_dict + if tools and not has_result: + name = next(iter(tools)) + part = types.Part( + function_call=types.FunctionCall( + name=name, args={"request": "nested task"} + ) + ) + else: + part = types.Part(text="done") + yield LlmResponse(content=types.Content(role="model", parts=[part])) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", [False, True]) +async def test_two_nested_agent_tools_inherit_one_decision_and_keep_context_isolation( + monkeypatch, native +): + judge = Judge() + cheap = DelegatingLlm(model="cheap") + unused = RecordingLlm(model="unused") + cls = LlmAgent if native else Agent + kwargs = {} if native else {"model_api_key": "test-only", "query_routing": False} + leaf = cls(name="leaf", model=unused, **kwargs) + child = cls(name="child", model=unused, tools=[AgentTool(leaf)], **kwargs) + agent = routed(monkeypatch, judge, unused, cheap, tools=[AgentTool(child)]) + runner, session = await setup(agent) + try: + await turn(runner, session, "ROOT_PRIVATE_CONTEXT") + assert len(judge.calls) == 1 + assert len(cheap.calls) == 5 + assert len(set(cheap.scopes)) == 1 + assert not unused.calls + assert "ROOT_PRIVATE_CONTEXT" not in str(cheap.calls[1]) + assert "ROOT_PRIVATE_CONTEXT" not in str(cheap.calls[2]) + assert get_turn_context() is None + assert leaf.canonical_model is unused + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_concurrent_sessions_keep_distinct_fixed_models(monkeypatch): + judge = Judge((0.01, 0.9)) + strong, cheap = RecordingLlm(model="strong"), RecordingLlm(model="cheap") + agent = routed(monkeypatch, judge, strong, cheap) + runner, first = await setup(agent) + second = await runner.session_service.create_session( + app_name="routing_contract", user_id="user" + ) + try: + await asyncio.gather( + turn(runner, first, "FIRST"), turn(runner, second, "SECOND") + ) + assert len(judge.calls) == 2 + assert len(strong.calls) == len(cheap.calls) == 1 + assert "FIRST" in str(cheap.calls[0]) + assert "SECOND" in str(strong.calls[0]) + assert get_turn_context() is None + finally: + await runner.close() + + +class TransferringLlm(RecordingLlm): + async def generate_content_async(self, llm_request, stream=False): + self.calls.append(copy.deepcopy(llm_request.contents)) + if len(self.calls) == 1: + part = types.Part( + function_call=types.FunctionCall( + name="transfer_to_agent", args={"agent_name": "child"} + ) + ) + else: + part = types.Part(text=self.model) + yield LlmResponse(content=types.Content(role="model", parts=[part])) + + +@pytest.mark.asyncio +async def test_native_transfer_and_next_turn_share_root_routing(monkeypatch): + judge = Judge((0.01, 0.9)) + unused = RecordingLlm(model="unused") + cheap, strong = TransferringLlm(model="cheap"), RecordingLlm(model="strong") + child = LlmAgent(name="child", model=unused) + root = routed(monkeypatch, judge, strong, cheap, sub_agents=[child]) + runner, session = await setup(root) + try: + await turn(runner, session, "Transfer task ORCHID") + await turn(runner, session, "Continue ORCHID") + assert len(judge.calls) == 2 + assert len(cheap.calls) == 2 + assert len(strong.calls) == 1 + assert not unused.calls + assert child.canonical_model is unused + assert "ORCHID" in str(strong.calls) + finally: + await runner.close() diff --git a/tests/test_cloud.py b/tests/test_cloud.py index b1698efd3..bb45f3ea0 100644 --- a/tests/test_cloud.py +++ b/tests/test_cloud.py @@ -14,6 +14,7 @@ import os import tempfile +import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from threading import Thread @@ -565,6 +566,7 @@ def send_archive(**kwargs): raise requests.ConnectionError() return Mock(status_code=200) + sleep = Mock() with ( patch( "veadk.integrations.ve_faas.ve_faas.zip_and_encode_folder", @@ -575,8 +577,9 @@ def send_archive(**kwargs): side_effect=send_archive, ) as upload, patch("veadk.integrations.ve_faas.ve_faas.signed_request") as callback, - patch("veadk.integrations.ve_faas.ve_faas.time.sleep") as sleep, + patch("veadk.integrations.ve_faas.ve_faas.time", SimpleNamespace(sleep=sleep)), ): + assert time.sleep is not sleep service._upload_and_mount_code("function-id", ".") assert upload.call_count == 2 diff --git a/veadk/agent.py b/veadk/agent.py index 6a0595e74..e26f26be2 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -39,6 +39,7 @@ from pydantic import AliasChoices, BaseModel, ConfigDict, Field from typing_extensions import Any +from veadk.decision import DecisionConfig, load_decision_config from veadk.config import settings from veadk.consts import DEFAULT_AGENT_NAME, DEFAULT_MODEL_EXTRA_CONFIG from veadk.knowledgebase import KnowledgeBase @@ -364,6 +365,9 @@ class Agent(LlmAgent): as tools for this agent. The cloud app must also mount the tunnel routes via `veadk.tunnel.mount_tunnel`/`mount_tunnel_if_enabled`.""" + decision: DecisionConfig | None = Field(default=None, exclude=True) + """Shared SystemOne connection; does not enable routing or content safety.""" + content_safety: Any = Field(default=None, exclude=True) """Opt-in content-safety configuration. @@ -377,6 +381,20 @@ class Agent(LlmAgent): ``plugins=agent.content_safety_plugins``. """ + query_routing: Any = Field(default=None, exclude=True) + """Opt-in query routing between a strong model and manually listed cheap models. + + ``None`` (default) uses a ``query_routing`` block from ``config.yaml`` when + one exists, so a project without that block keeps its existing behavior. + ``False`` disables routing explicitly. ``True`` requires the YAML block. + A mapping or :class:`veadk.configs.routing_configs.QueryRoutingConfig` + supplies the configuration directly. + + Routing is only supported by the ADK runtime. It preserves the Agent's + configured model as the strong tier and delegates cheap-tier calls to + prebuilt per-tier model instances. + """ + def model_post_init(self, __context: Any) -> None: # Snapshot before anything below assigns fields on ``self``: pydantic # adds every assignment to ``model_fields_set``, so this is the only @@ -644,6 +662,8 @@ def model_post_init(self, __context: Any) -> None: else: self.after_agent_callback = dataset_auto_gen_callback + self._setup_query_routing() + logger.info(f"VeADK version: {VERSION}") logger.info(f"{self.__class__.__name__} `{self.name}` init done.") @@ -661,6 +681,71 @@ def model_post_init(self, __context: Any) -> None: check_agent_runtime_support(self, self.runtime) + def _setup_query_routing(self) -> None: + """Wrap the built model when query routing is configured. + + The configuration is resolved after the strong model exists, so a + project without a ``query_routing`` block is untouched and keeps the + exact model instance it built above. + """ + + from veadk.routing.factory import ( + build_query_routing, + resolve_query_routing_config, + ) + + config = resolve_query_routing_config(self.query_routing) + if config is None: + return + if self.runtime != "adk": + raise ValueError( + "query_routing is only supported by the ADK runtime; " + f"got runtime={self.runtime}" + ) + + if self.decision is None: + self.decision = load_decision_config() + router, _ = build_query_routing(config, self, decision=self.decision) + self.query_routing = config + self.model = router + from veadk.routing.agent_tool import bind_sub_agents + + bind_sub_agents(self) + + @property + def canonical_model(self): + from veadk.routing.context import get_turn_context + from veadk.routing.llm import RoutingLlm + + turn = get_turn_context() + if turn is not None: + return turn.model + if isinstance(self.model, RoutingLlm): + return self.model.strong_model + return super().canonical_model + + @property + def canonical_before_model_callbacks(self): + from veadk.routing.context import get_turn_context + from veadk.routing.execution import bind_request + + callbacks = super().canonical_before_model_callbacks + return ( + [*callbacks, bind_request] if get_turn_context() is not None else callbacks + ) + + @property + def canonical_before_tool_callbacks(self): + from veadk.routing.context import get_turn_context + from veadk.routing.tool_arguments import validate_tool_arguments + + callbacks = super().canonical_before_tool_callbacks + return ( + [*callbacks, validate_tool_arguments] + if get_turn_context() is not None + else callbacks + ) + def update_model(self, model_name: str): """Point the agent at a different model. @@ -674,6 +759,13 @@ def update_model(self, model_name: str): Args: model_name (str): The new model name, without a provider prefix. """ + from veadk.routing.llm import RoutingLlm + + if isinstance(self.model, RoutingLlm): + raise ValueError( + "update_model is not supported while query_routing is enabled; " + "configure strong_model in query_routing instead" + ) logger.info(f"Updating model to {model_name}") self.model_name = model_name self.model = self.model.model_copy( @@ -878,6 +970,7 @@ def content_safety_plugins(self) -> list[Any]: @property def _llm_flow(self) -> BaseLlmFlow: + from veadk.routing.execution import prepare_flow from google.adk.flows.llm_flows.auto_flow import AutoFlow from google.adk.flows.llm_flows.single_flow import SingleFlow @@ -890,16 +983,16 @@ def _llm_flow(self) -> BaseLlmFlow: if self.enable_supervisor: logger.debug(f"Enable supervisor flow for agent: {self.name}") - return SupervisorSingleFlow(supervised_agent=self) + return prepare_flow(SupervisorSingleFlow(supervised_agent=self)) else: - return SingleFlow() + return prepare_flow(SingleFlow()) else: from veadk.flows.supervise_auto_flow import SupervisorAutoFlow if self.enable_supervisor: logger.debug(f"Enable supervisor flow for agent: {self.name}") - return SupervisorAutoFlow(supervised_agent=self) - return AutoFlow() + return prepare_flow(SupervisorAutoFlow(supervised_agent=self)) + return prepare_flow(AutoFlow()) async def _run_async_impl( self, ctx: "InvocationContext" @@ -917,6 +1010,16 @@ async def _run_async_impl( return if self.runtime == "adk": + from veadk.routing.context import get_turn_context + from veadk.routing.execution import run_routed, routing_owner + + if routing_owner(self) is not None or get_turn_context() is not None: + async with aclosing( + run_routed(self, ctx, super()._run_async_impl) + ) as events: + async for event in events: + yield event + return # A transfer can close this wrapper before the LLM stream ends. # Close it in the same task so tracing contexts do not leak into # async-generator finalization or the receiving sub-agent. diff --git a/veadk/configs/routing_configs.py b/veadk/configs/routing_configs.py new file mode 100644 index 000000000..b6ea4ffe5 --- /dev/null +++ b/veadk/configs/routing_configs.py @@ -0,0 +1,415 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Opt-in query routing configuration. + +Query routing is disabled unless a ``query_routing`` block exists in the +project ``config.yaml`` (or the caller passes an explicit configuration to +:class:`veadk.Agent`). The block is intentionally small: the only required +field is ``simple_models``. + +Canonical shape:: + + query_routing: + simple_models: + - doubao-seed-2-1-lite-260915 + questions: + - question: 回答是否需要多步推理或长链分析? + threshold: 0.15 + +User questions are appended to the four built-in conservative questions. A +question only needs its text and threshold; the Decision API type is always +``noul`` and the answer id is generated as ``custom_1``, ``custom_2``, ... +""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Literal + +import yaml +from pydantic import ( + AliasChoices, + BaseModel, + ConfigDict, + Field, + SecretStr, + field_validator, + model_validator, +) + + +from veadk.decision import DecisionConfig + +DEFAULT_QUERY_ROUTING_THRESHOLD = 0.15 +DEFAULT_QUERY_ROUTING_TIMEOUT_SECONDS = 2.0 +MAX_QUERY_ROUTING_CONTEXT_CHARS = 32_000 +_REQUEST_OVERRIDE_FIELDS = frozenset( + { + "previous_response_id", + "previous_interaction_id", + "input", + "messages", + "instructions", + "tools", + "functions", + } +) +_CONNECTION_FIELDS = frozenset({"model", "api_key", "api_base", "fallbacks"}) + +_BUILTIN_QUESTIONS: tuple[tuple[str, str, float], ...] = ( + ( + "is_complex", + "完成当前用户任务是否需要超出候选模型能力的复杂推理、规划、代码或数学推导?" + "简单事实问答和单步改写判为否。", + 0.5, + ), + ( + "needs_tools", + "完成当前任务是否需要超出候选模型已声明能力的工具编排或错误恢复?" + "候选模型支持的简单工具查询判为否。", + 0.5, + ), + ( + "needs_long_context", + "当前用户任务是否必须综合较长历史或多份材料进行全局分析?" + "仅从少量最近消息提取已明确给出的代号、姓名、数值等事实,判为否;" + "不能因为任务需要回看历史就判为是。只判断当前任务实际需要的分析。", + 0.5, + ), + ( + "is_high_stakes", + "当前用户任务是否要求医疗、法律、金融、安全等高风险结论,或不可撤销的操作?" + "一般咨询、创作、草稿和单纯提取已给出的非敏感事实判为否。" + "历史中曾讨论高风险主题不代表当前任务也是高风险;" + "但当前任务仍在追问相关建议或操作时判为是。", + 0.05, + ), +) + + +class QueryRoutingConfigError(ValueError): + """Raised when query-routing configuration cannot be trusted.""" + + +class RoutingQuestion(BaseModel): + """One ``noul`` question sent to the Decision API.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + strict=True, + populate_by_name=True, + ) + + id: str = Field(min_length=1, max_length=128) + instructions: str = Field( + min_length=1, + validation_alias=AliasChoices("instructions", "question"), + ) + threshold: float = Field( + default=DEFAULT_QUERY_ROUTING_THRESHOLD, + ge=0.0, + le=1.0, + ) + + @property + def type(self) -> Literal["noul"]: + """The Decision API question type is fixed to ``noul``.""" + + return "noul" + + +class SimpleModelConfig(BaseModel): + """A manually listed cheap model and its declared capabilities. + + Unknown capabilities are denied by default: a bare model name is treated as + text-only, without tools, structured output, or media support. Callers can + opt into more capability per model when the model genuinely supports it. + """ + + model_config = ConfigDict( + extra="forbid", + frozen=True, + strict=True, + populate_by_name=True, + ) + + name: str = Field( + min_length=1, + validation_alias=AliasChoices("name", "model", "model_name"), + ) + supports_tools: bool = False + capability_description: str | None = Field( + default=None, min_length=1, max_length=2000 + ) + provider: str | None = None + api_base: str | None = None + api_key: SecretStr | None = Field(default=None, repr=False) + api_key_env: str | None = None + api: Literal["chat_completions", "responses"] | None = None + extra_config: dict[str, Any] = Field(default_factory=dict, repr=False) + supports_media: bool = False + supports_structured_output: bool = False + max_context_chars: int = Field( + default=MAX_QUERY_ROUTING_CONTEXT_CHARS, + ge=1, + ) + + @property + def api_key_value(self) -> str | None: + if self.api_key_env: + value = os.getenv(self.api_key_env) + if not value: + raise QueryRoutingConfigError( + f"missing credential variable: {self.api_key_env}" + ) + return value + return self.api_key.get_secret_value() if self.api_key else None + + @model_validator(mode="after") + def _connection_contract(self) -> "SimpleModelConfig": + if self.api_base: + DecisionConfig(api_url=self.api_base) + if self.api_key and self.api_key_env: + raise ValueError("use only one of api_key and api_key_env") + forbidden = _CONNECTION_FIELDS | _REQUEST_OVERRIDE_FIELDS + body = self.extra_config.get("extra_body") or {} + if forbidden.intersection(self.extra_config) or forbidden.intersection(body): + raise ValueError( + "connection, request and fallback fields are not allowed in extra_config" + ) + return self + + +def default_questions() -> list[RoutingQuestion]: + """Return the four built-in conservative questions, in fixed order.""" + + return [ + RoutingQuestion(id=name, instructions=instructions, threshold=threshold) + for name, instructions, threshold in _BUILTIN_QUESTIONS + ] + + +class QueryRoutingConfig(BaseModel): + """Validated, immutable query-routing configuration. + + ``questions`` always contains the built-in questions followed by the + caller's questions. Built-in ids are reserved, so a custom question can + never replace a built-in question. Thresholds have explicit overrides. + """ + + model_config = ConfigDict( + hide_input_in_errors=True, + extra="forbid", + frozen=True, + populate_by_name=True, + ) + + simple_models: list[SimpleModelConfig] = Field(min_length=1) + strong_model: str | None = None + mode: Literal["enforce", "shadow"] = "shadow" + builtin_thresholds: dict[str, float] = Field(default_factory=dict) + questions: list[RoutingQuestion] = Field(default_factory=list) + decision_timeout: float = Field( + default=DEFAULT_QUERY_ROUTING_TIMEOUT_SECONDS, + gt=0.0, + le=30.0, + ) + + @field_validator("simple_models", mode="before") + @classmethod + def _normalize_simple_models(cls, value: Any) -> Any: + if not isinstance(value, list): + return value + normalized: list[Any] = [] + for item in value: + normalized.append({"name": item} if isinstance(item, str) else item) + return normalized + + @model_validator(mode="before") + @classmethod + def _merge_questions(cls, data: Any) -> Any: + if not isinstance(data, dict): + return data + provided = data.get("questions") or [] + if not isinstance(provided, list): + return data + + merged: list[Any] = [ + { + "id": question.id, + "instructions": question.instructions, + "threshold": question.threshold, + } + for question in default_questions() + ] + thresholds = data.get("builtin_thresholds", {}) + if set(thresholds) - {q["id"] for q in merged}: + raise ValueError("unknown built-in threshold id") + for entry in merged: + if entry["id"] in thresholds: + entry["threshold"] = thresholds[entry["id"]] + seen = {question["id"] for question in merged} + for index, item in enumerate(provided, start=1): + if isinstance(item, str): + entry: dict[str, Any] = {"instructions": item} + elif isinstance(item, RoutingQuestion): + entry = { + "id": item.id, + "instructions": item.instructions, + "threshold": item.threshold, + } + elif isinstance(item, dict): + entry = dict(item) + else: + raise ValueError( + "questions items must be strings, mappings, or RoutingQuestion" + ) + + entry.setdefault("id", f"custom_{index}") + question_id = entry["id"] + if question_id in seen: + raise ValueError( + f"question id `{question_id}` is reserved or duplicated" + ) + seen.add(question_id) + merged.append(entry) + + return {**data, "questions": merged} + + @field_validator("strong_model") + @classmethod + def _validate_strong_model(cls, value: str | None) -> str | None: + if value is not None and not value.strip(): + raise ValueError("strong_model must be non-empty when set") + return value + + @model_validator(mode="after") + def _validate_unique_models(self) -> "QueryRoutingConfig": + names = [model.name for model in self.simple_models] + if len(names) != len(set(names)): + raise ValueError("simple_models must not contain duplicate names") + return self + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "QueryRoutingConfig": + """Build from explicitly supplied code values; environment is ignored.""" + + return cls.model_validate(data) + + def questions_hash(self) -> str: + """Stable hash of the question set, for diagnostics and cache keys.""" + + payload = [ + { + "id": question.id, + "type": question.type, + "instructions": question.instructions, + "threshold": question.threshold, + } + for question in self.questions + ] + encoded = json.dumps( + payload, ensure_ascii=False, sort_keys=True, separators=(",", ":") + ) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def _resolve_config_path(path: str | Path | None, *, optional: bool) -> Path | None: + from dotenv import find_dotenv + + if path is not None: + return Path(path) + found = find_dotenv(filename="config.yaml", usecwd=True) + if not found: + if optional: + return None + raise QueryRoutingConfigError("no config.yaml found; pass an explicit path") + return Path(found) + + +def _read_config_document(config_path: Path) -> dict[str, Any]: + try: + text = config_path.read_text(encoding="utf-8") + except OSError as exc: + raise QueryRoutingConfigError( + f"cannot read query-routing config: {config_path}" + ) from exc + + try: + document = yaml.safe_load(text) + except yaml.YAMLError as exc: + raise QueryRoutingConfigError(f"invalid YAML in {config_path}: {exc}") from exc + + if document is None: + return {} + if not isinstance(document, dict): + raise QueryRoutingConfigError("config root must be a mapping") + return document + + +def _build_config(block: dict[str, Any], config_path: Path) -> QueryRoutingConfig: + try: + return QueryRoutingConfig.model_validate(block) + except Exception as exc: + raise QueryRoutingConfigError( + f"invalid `query_routing` block in {config_path}: {exc}" + ) from exc + + +def load_query_routing_config( + path: str | Path | None = None, +) -> QueryRoutingConfig: + """Load ``query_routing`` from ``config.yaml``. + + When ``path`` is omitted, the same ``config.yaml`` discovery used by VeADK + (``find_dotenv`` from the current directory) applies. A missing block is an + error; use :func:`load_query_routing_config_if_present` for optional + discovery. + """ + + config_path = _resolve_config_path(path, optional=False) + assert config_path is not None + document = _read_config_document(config_path) + block = document.get("query_routing") + if not isinstance(block, dict): + raise QueryRoutingConfigError("config must contain a `query_routing` mapping") + return _build_config(block, config_path) + + +def load_query_routing_config_if_present( + path: str | Path | None = None, +) -> QueryRoutingConfig | None: + """Load ``query_routing`` when configured, otherwise return ``None``. + + Missing files and a missing ``query_routing`` block mean "not configured" + and leave routing disabled. A block that exists but is invalid raises, so + misconfiguration fails at startup instead of silently disabling routing. + """ + + config_path = _resolve_config_path(path, optional=True) + if config_path is None: + return None + document = _read_config_document(config_path) + block = document.get("query_routing") + if block is None: + return None + if not isinstance(block, dict): + raise QueryRoutingConfigError( + f"`query_routing` in {config_path} must be a mapping" + ) + return _build_config(block, config_path) diff --git a/veadk/content_safety/__init__.py b/veadk/content_safety/__init__.py index 8bc89ac88..cba93d138 100644 --- a/veadk/content_safety/__init__.py +++ b/veadk/content_safety/__init__.py @@ -35,13 +35,11 @@ from .config import ( DEFAULT_BLOCK_MESSAGE, - DEFAULT_DECISION_MODEL, DEFAULT_QUESTION_INSTRUCTIONS, DEFAULT_QUESTION_THRESHOLD, DEFAULT_QUESTION_THRESHOLDS, ContentSafetyConfig, ContentSafetyConfigError, - DecisionConfig, QuestionConfig, default_questions, load_content_safety_config, @@ -72,7 +70,6 @@ "ContentSafetyConfig", "ContentSafetyConfigError", "DecisionPlugin", - "DecisionConfig", "Provider", "QuestionConfig", "QuestionResult", @@ -80,7 +77,6 @@ "StatePayload", "STATE_POLICY_PREFIX", "DEFAULT_BLOCK_MESSAGE", - "DEFAULT_DECISION_MODEL", "DEFAULT_QUESTION_INSTRUCTIONS", "DEFAULT_QUESTION_THRESHOLD", "DEFAULT_QUESTION_THRESHOLDS", diff --git a/veadk/content_safety/config.py b/veadk/content_safety/config.py index 4c42bd882..4866fd37c 100644 --- a/veadk/content_safety/config.py +++ b/veadk/content_safety/config.py @@ -23,10 +23,6 @@ tool: content_safety: provider: decision - decision: - api_url: https://ark.example/api/v1/systemone # optional - api_key: ... # optional - model: doubao-seed-2-1-lite-260915 # optional questions: # optional; overrides/extends the built-ins prompt_injection: instructions: ... @@ -35,21 +31,17 @@ from __future__ import annotations -import copy import hashlib import json import math -import os from pathlib import Path from typing import Any, Literal -from urllib.parse import urlsplit import yaml from pydantic import ( BaseModel, ConfigDict, Field, - SecretStr, field_validator, model_validator, ) @@ -60,17 +52,10 @@ logger = get_logger(__name__) -DEFAULT_DECISION_MODEL = "doubao-seed-2-1-lite-260915" DEFAULT_BLOCK_MESSAGE = ( "Your request was blocked by the content safety policy. " "Please modify your input and try again." ) -_ENV_PREFIX = "TOOL_CONTENT_SAFETY_DECISION_" -_ENV_FIELDS = { - "api_url": "API_URL", - "api_key": "API_KEY", - "model": "MODEL", -} DEFAULT_QUESTION_INSTRUCTIONS = { # Resource and input manipulation "compute_abuse": ( @@ -188,57 +173,6 @@ class ContentSafetyConfigError(ValueError): """Raised when a content-safety configuration file cannot be trusted.""" -class DecisionConfig(BaseModel): - """Connection details for the System One decision API. - - ``api_url`` is a complete System One URL and is used verbatim. When it is - absent, the factory derives ``/api/v1/systemone`` from the effective Ark - API base. ``api_key`` defaults to the existing Ark credential resolution and - is stored as a ``SecretStr`` so repr and serialization do not expose it. - """ - - model_config = ConfigDict(extra="forbid", frozen=True) - - api_url: str | None = None - api_key: SecretStr | None = None - model: str = Field(default=DEFAULT_DECISION_MODEL, min_length=1) - - @field_validator("api_url") - @classmethod - def _validate_api_url(cls, value: str | None) -> str | None: - if value is None: - return None - if not value.startswith(("http://", "https://")): - raise ValueError("decision.api_url must start with http:// or https://") - try: - parsed = urlsplit(value) - except ValueError as exc: - raise ValueError("decision.api_url must be a valid URL") from exc - if not parsed.hostname: - raise ValueError("decision.api_url must include a host") - if parsed.scheme == "http" and parsed.hostname.lower() not in { - "localhost", - "127.0.0.1", - "::1", - }: - raise ValueError( - "decision.api_url must use https unless it targets loopback" - ) - return value - - @field_validator("api_key") - @classmethod - def _validate_api_key(cls, value: SecretStr | None) -> SecretStr | None: - if value is not None and not value.get_secret_value(): - raise ValueError("decision.api_key must be non-empty when set") - return value - - @property - def api_key_value(self) -> str | None: - """Return the plaintext key for the HTTP layer, never for display.""" - return self.api_key.get_secret_value() if self.api_key is not None else None - - class QuestionConfig(BaseModel): """One question sent to System One. @@ -294,12 +228,11 @@ class ContentSafetyConfig(BaseModel): behind an adapter. """ - model_config = ConfigDict(extra="forbid", frozen=True) + model_config = ConfigDict(hide_input_in_errors=True, extra="forbid", frozen=True) provider: Provider = Provider.NONE on_error: Literal["allow", "block"] = "allow" block_message: str = Field(default=DEFAULT_BLOCK_MESSAGE, min_length=1) - decision: DecisionConfig = Field(default_factory=DecisionConfig) use_default_questions: bool = True """Merge the built-in questions under ``questions``. @@ -448,7 +381,6 @@ def _read_config_document(config_path: Path) -> dict[str, Any]: def _build_config(block: dict[str, Any], config_path: Path) -> ContentSafetyConfig: - block = _apply_decision_env(block) try: return ContentSafetyConfig.model_validate(block) except Exception as exc: @@ -478,23 +410,6 @@ def _resolve_config_path(path: str | Path | None, *, optional: bool) -> Path | N return Path(found) -def _apply_decision_env(block: dict[str, Any]) -> dict[str, Any]: - """Apply process/.env overrides to the decision connection fields.""" - block = copy.deepcopy(block) - decision = block.get("decision") - if decision is None: - decision = {} - if not isinstance(decision, dict): - return block - for field, suffix in _ENV_FIELDS.items(): - value = os.getenv(f"{_ENV_PREFIX}{suffix}") - if value: - decision[field] = value - if decision: - block["decision"] = decision - return block - - def load_content_safety_config( path: str | Path | None = None, ) -> ContentSafetyConfig: @@ -506,9 +421,8 @@ def load_content_safety_config( This is deliberately opt-in: importing :mod:`veadk.content_safety` never looks for this file, and the rest of VeADK keeps its existing behavior. - Precedence for the three decision connection fields is process environment - (which includes a loaded ``.env``) over YAML; explicit code values passed - through :meth:`ContentSafetyConfig.from_dict` are never overridden. + Shared connection fields are loaded separately by + :func:`veadk.decision.load_decision_config` from root ``decision``. """ _load_dotenv() diff --git a/veadk/content_safety/factory.py b/veadk/content_safety/factory.py index 5a73f8766..76d5533c2 100644 --- a/veadk/content_safety/factory.py +++ b/veadk/content_safety/factory.py @@ -16,106 +16,28 @@ from __future__ import annotations from typing import Any -from urllib.parse import urlsplit from google.adk.plugins import BasePlugin -from veadk.decision import DecisionClient, DecisionJudge +from veadk.decision import ( + DecisionClient, + DecisionJudge, + DecisionConfig, + load_decision_config, + resolve_decision_connection, +) -from .config import ContentSafetyConfig, ContentSafetyConfigError +from .config import ContentSafetyConfig from .decider import ThresholdDecider from .plugin import DecisionPlugin from .state import StateBuilder from .types import Provider -def _systemone_url(api_base: str) -> str: - """Derive the Beta System One URL from an Ark API base. - - The design inherits the region/domain/gateway prefix and appends - ``/api/v1/systemone``; it explicitly must not produce ``/api/v3/systemone``. - """ - base = api_base.rstrip("/") - for suffix in ("/api/v3", "/api/v1"): - if base.endswith(suffix): - base = base[: -len(suffix)] - break - return f"{base}/api/v1/systemone" - - -def _ark_defaults(api_base: str | None, api_key: str | None) -> tuple[str, str]: - if api_base and api_key: - return api_base, api_key - from veadk.config import settings - - return api_base or settings.model.api_base, api_key or settings.model.api_key - - -_ARK_HOST_SUFFIXES = (".volces.com", ".bytepluses.com") - - -def _is_ark_api_base(api_base: str) -> bool: - """Return whether ``api_base`` is a known Ark HTTPS endpoint. - - The content-safety fallback must never derive a System One URL from an - arbitrary model endpoint: doing so could send audited content to a - non-Ark provider. Custom gateways must be configured explicitly through - ``decision.api_url``. - """ - try: - parsed = urlsplit(api_base) - except ValueError: - return False - host = (parsed.hostname or "").lower() - return ( - parsed.scheme == "https" - and host.startswith("ark.") - and host.endswith(_ARK_HOST_SUFFIXES) - ) - - -def _require_ark_api_base(api_base: str) -> None: - if not _is_ark_api_base(api_base): - raise ContentSafetyConfigError( - "Cannot derive the System One endpoint from a non-Ark model API base. " - "Set decision.api_url explicitly to use a custom endpoint." - ) - - -def _resolve_decision_connection( - config: ContentSafetyConfig, - ark_api_base: str | None, - ark_api_key: str | None, -) -> tuple[str, str]: - """Resolve the System One URL and key without an unsafe cross-provider mix.""" - api_url = config.decision.api_url - api_key = config.decision.api_key_value - - if api_url: - if api_key: - return api_url, api_key - if not _is_ark_api_base(api_url): - raise ContentSafetyConfigError( - "decision.api_key is required when decision.api_url is not an Ark " - "endpoint; refusing to reuse the Ark credential for a custom endpoint." - ) - base, fallback_key = _ark_defaults(ark_api_base, ark_api_key) - if not _is_ark_api_base(base): - raise ContentSafetyConfigError( - "decision.api_key is required when decision.api_url is an Ark " - "endpoint but the fallback Ark API base is not an Ark endpoint; " - "refusing to reuse a non-Ark credential for Ark." - ) - return api_url, fallback_key - - base, fallback_key = _ark_defaults(ark_api_base, ark_api_key) - _require_ark_api_base(base) - return _systemone_url(base), api_key or fallback_key - - def build_content_safety( config: ContentSafetyConfig, *, + decision: DecisionConfig | None = None, client: DecisionJudge | None = None, decider: ThresholdDecider | None = None, state_builder: StateBuilder | None = None, @@ -141,13 +63,14 @@ def build_content_safety( if config.provider is Provider.DECISION: judge = client if judge is None: - api_url, api_key = _resolve_decision_connection( - config, ark_api_base, ark_api_key + connection = decision if decision is not None else load_decision_config() + api_url, api_key = resolve_decision_connection( + connection, ark_api_base, ark_api_key ) judge = DecisionClient( endpoint=api_url, api_key=api_key, - model=config.decision.model, + model=connection.model, ) return DecisionPlugin( config=config, diff --git a/veadk/content_safety/integration.py b/veadk/content_safety/integration.py index 6c27e8398..14bf192d9 100644 --- a/veadk/content_safety/integration.py +++ b/veadk/content_safety/integration.py @@ -97,6 +97,7 @@ def resolve_content_safety_plugin( return build_content_safety( config, + decision=getattr(agent, "decision", None), ark_api_base=getattr(agent, "model_api_base", None), ark_api_key=getattr(agent, "model_api_key", None), ) diff --git a/veadk/decision/__init__.py b/veadk/decision/__init__.py index b4e71cf17..442ab37c8 100644 --- a/veadk/decision/__init__.py +++ b/veadk/decision/__init__.py @@ -19,6 +19,13 @@ DecisionClientError, DecisionJudge, ) +from .config import ( + DEFAULT_DECISION_MODEL, + DecisionConfig, + DecisionConfigError, + load_decision_config, +) +from .connection import resolve_decision_connection from .types import ( DecisionQuestion, DecisionRequest, @@ -27,6 +34,11 @@ ) __all__ = [ + "DEFAULT_DECISION_MODEL", + "DecisionConfig", + "DecisionConfigError", + "load_decision_config", + "resolve_decision_connection", "DEFAULT_TIMEOUT_SECONDS", "DecisionClient", "DecisionClientError", diff --git a/veadk/decision/config.py b/veadk/decision/config.py new file mode 100644 index 000000000..3ee020bd1 --- /dev/null +++ b/veadk/decision/config.py @@ -0,0 +1,157 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared, immutable SystemOne connection configuration; enables no consumer.""" + +from pathlib import Path +import os +from urllib.parse import urlsplit + +import yaml +from pydantic import ( + BaseModel, + ConfigDict, + Field, + SecretStr, + ValidationError, + field_validator, +) + +DEFAULT_DECISION_MODEL = "doubao-seed-2-1-lite-260915" + + +class DecisionConfigError(ValueError): + """Invalid common decision connection configuration.""" + + +class DecisionConfig(BaseModel): + """Connection details for the System One decision API. + + ``api_url`` is a complete System One URL and is used verbatim. When it is + absent, the factory derives ``/api/v1/systemone`` from the effective Ark + API base. ``api_key`` defaults to the existing Ark credential resolution and + is stored as a ``SecretStr`` so repr and serialization do not expose it. + """ + + model_config = ConfigDict(hide_input_in_errors=True, extra="forbid", frozen=True) + + api_url: str | None = None + api_key: SecretStr | None = None + model: str = Field(default=DEFAULT_DECISION_MODEL, min_length=1) + + @field_validator("api_url") + @classmethod + def _validate_api_url(cls, value: str | None) -> str | None: + if value is None: + return None + if not value.startswith(("http://", "https://")): + raise ValueError("decision.api_url must start with http:// or https://") + try: + parsed = urlsplit(value) + except ValueError as exc: + raise ValueError("decision.api_url must be a valid URL") from exc + if not parsed.hostname: + raise ValueError("decision.api_url must include a host") + if parsed.scheme == "http" and parsed.hostname.lower() not in { + "localhost", + "127.0.0.1", + "::1", + }: + raise ValueError( + "decision.api_url must use https unless it targets loopback" + ) + return value + + @field_validator("api_key") + @classmethod + def _validate_api_key(cls, value: SecretStr | None) -> SecretStr | None: + if value is not None and not value.get_secret_value(): + raise ValueError("decision.api_key must be non-empty when set") + return value + + @property + def api_key_value(self) -> str | None: + """Return the plaintext key for the HTTP layer, never for display.""" + return self.api_key.get_secret_value() if self.api_key is not None else None + + +class _UniqueKeyLoader(yaml.SafeLoader): + pass + + +def _unique_mapping(loader, node, deep=False): + result = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + if not isinstance(key, str) or key in result: + raise DecisionConfigError("decision config requires unique string keys") + result[key] = loader.construct_object(value_node, deep=deep) + return result + + +_UniqueKeyLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, _unique_mapping +) + + +def load_decision_config(path: str | Path | None = None) -> DecisionConfig: + """Load root ``decision``; DECISION_* environment values override YAML. + + No file/block means defaults. An explicit unreadable file, invalid block, + or ambiguous destination/key override fails at startup. Code construction + bypasses this loader, so explicit values are never overwritten. + """ + from dotenv import find_dotenv, load_dotenv + + dotenv_path = find_dotenv(filename=".env", usecwd=True) + if dotenv_path: + load_dotenv(dotenv_path, override=False) + location = ( + str(path) + if path is not None + else find_dotenv(filename="config.yaml", usecwd=True) + ) + block = {} + if location: + try: + document = yaml.load( + Path(location).read_text(encoding="utf-8"), Loader=_UniqueKeyLoader + ) + except (OSError, yaml.YAMLError): + raise DecisionConfigError("cannot read decision configuration") from None + if not isinstance(document, dict): + raise DecisionConfigError("decision config root must be a mapping") + block = document.get("decision", {}) + if not isinstance(block, dict): + raise DecisionConfigError("decision must be a mapping") + block = dict(block) + endpoint = os.getenv("DECISION_API_URL") + key = os.getenv("DECISION_API_KEY") + if ( + endpoint + and endpoint != block.get("api_url") + and block.get("api_key") + and not key + ): + raise DecisionConfigError( + "DECISION_API_KEY is required when overriding a configured decision destination" + ) + for field in ("api_url", "api_key", "model"): + value = os.getenv(f"DECISION_{field.upper()}") + if value is not None: + block[field] = value + try: + return DecisionConfig.model_validate(block) + except ValidationError: + # ValidationError includes raw input; never put credentials into logs. + raise DecisionConfigError("invalid decision configuration") from None diff --git a/veadk/decision/connection.py b/veadk/decision/connection.py new file mode 100644 index 000000000..a1d0c0967 --- /dev/null +++ b/veadk/decision/connection.py @@ -0,0 +1,87 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Common SystemOne endpoint and credential resolution.""" + +from urllib.parse import urlsplit + +from .config import DecisionConfig, DecisionConfigError + + +def _systemone_url(api_base: str) -> str: + """Derive the Beta System One URL from an Ark API base. + + The design inherits the region/domain/gateway prefix and appends + ``/api/v1/systemone``; it explicitly must not produce ``/api/v3/systemone``. + """ + base = api_base.rstrip("/") + for suffix in ("/api/v3", "/api/v1"): + if base.endswith(suffix): + base = base[: -len(suffix)] + break + return f"{base}/api/v1/systemone" + + +_ARK_HOST_SUFFIXES = (".volces.com", ".bytepluses.com") + + +def _is_ark_api_base(api_base: str) -> bool: + """Return whether ``api_base`` is a known Ark HTTPS endpoint. + + The content-safety fallback must never derive a System One URL from an + arbitrary model endpoint: doing so could send audited content to a + non-Ark provider. Custom gateways must be configured explicitly through + ``decision.api_url``. + """ + try: + parsed = urlsplit(api_base) + except ValueError: + return False + host = (parsed.hostname or "").lower() + return ( + parsed.scheme == "https" + and host.startswith("ark.") + and host.endswith(_ARK_HOST_SUFFIXES) + ) + + +def resolve_decision_connection( + config: DecisionConfig, + ark_api_base: str | None = None, + ark_api_key: str | None = None, +) -> tuple[str, str]: + """Resolve credentials without crossing an implicit provider boundary.""" + endpoint, explicit_key = config.api_url, config.api_key_value + if endpoint and explicit_key: + return endpoint, explicit_key + if endpoint and not _is_ark_api_base(endpoint): + raise DecisionConfigError( + "decision.api_key is required when decision.api_url is not an Ark endpoint" + ) + if ark_api_base is None: + from veadk.config import settings + + ark_api_base = settings.model.api_base + base = ark_api_base + if not _is_ark_api_base(base): + raise DecisionConfigError( + "cannot derive System One connection: fallback Ark API base is not an Ark endpoint; set decision.api_url and decision.api_key explicitly" + ) + key = explicit_key if explicit_key is not None else ark_api_key + if key is None: + from veadk.config import settings + + key = settings.model.api_key + if not key: + raise DecisionConfigError("decision API key is not configured") + return endpoint or _systemone_url(base), key diff --git a/veadk/models/ark_llm.py b/veadk/models/ark_llm.py index 130a96a3b..f5b2f3567 100644 --- a/veadk/models/ark_llm.py +++ b/veadk/models/ark_llm.py @@ -262,6 +262,8 @@ def _get_content( def _content_to_input_item( content: types.Content, + *, + text_before_tool_calls: bool = False, ) -> Union[ResponseInputItemParam, List[ResponseInputItemParam]]: role = _to_ark_role(content.role) @@ -286,6 +288,11 @@ def _content_to_input_item( if input_content: return input_content else: # model + # Portable replay keeps a tool call adjacent to its result. Otherwise + # explanatory text can prematurely close the tool-call message group. + if input_content and text_before_tool_calls: + input_list.append(input_content) + input_content = None # 3. Processing model messages for part in content.parts: if part.function_call: @@ -375,6 +382,8 @@ def _responses_schema_to_text( def _get_responses_inputs( llm_request: LlmRequest, + *, + text_before_tool_calls: bool = False, ) -> Tuple[ Optional[str], Optional[List[ResponseInputItemParam]], @@ -392,7 +401,9 @@ def _get_responses_inputs( # Each content represents `one conversation`. # This `one conversation` may contain `multiple pieces of content`, # but it cannot contain `multiple conversations`. - input_item_or_list = _content_to_input_item(content) + input_item_or_list = _content_to_input_item( + content, text_before_tool_calls=text_before_tool_calls + ) if isinstance(input_item_or_list, list): input_params.extend(input_item_or_list) elif input_item_or_list: @@ -744,8 +755,15 @@ async def generate_content_async( self._maybe_append_user_content(llm_request) # logger.debug(_build_request_log(llm_request)) + from veadk.routing.context import get_turn_context + + turn = get_turn_context() + instructions, input_param, tools, text_format, generation_params = ( - _get_responses_inputs(llm_request) + _get_responses_inputs( + llm_request, + text_before_tool_calls=turn is not None and not turn.selection.shadow, + ) ) if "functions" in self._additional_args: diff --git a/veadk/routing/__init__.py b/veadk/routing/__init__.py new file mode 100644 index 000000000..0ff5983db --- /dev/null +++ b/veadk/routing/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Opt-in per-invocation query routing for VeADK agents.""" + +from veadk.routing.context import ( + RouteSelection, + get_route_selection, + set_route_selection, +) +from veadk.routing.llm import RoutingLlm +from veadk.routing.store import RouteStore, get_route_store + +__all__ = [ + "RouteSelection", + "RouteStore", + "RoutingLlm", + "get_route_selection", + "get_route_store", + "set_route_selection", +] diff --git a/veadk/routing/agent_tool.py b/veadk/routing/agent_tool.py new file mode 100644 index 000000000..18cb907f9 --- /dev/null +++ b/veadk/routing/agent_tool.py @@ -0,0 +1,107 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Execution-local views for native ADK AgentTool children. + +The tool, original agents and their sessions remain independently owned. +Only the view used for this call inherits the root's model binding. +""" + +from __future__ import annotations + +import copy +from contextlib import aclosing +from typing import Any + +from google.adk.agents import LlmAgent +from google.adk.tools import AgentTool + +from veadk.routing.context import get_turn_context +from veadk.routing.execution import ( + bind_request, + prepare_flow, + run_routed, + routing_owner, +) + + +class RoutedLlmAgent(LlmAgent): + @property + def canonical_model(self): + turn = get_turn_context() + return turn.model if turn else super().canonical_model + + @property + def canonical_before_model_callbacks(self): + callbacks = super().canonical_before_model_callbacks + return [*callbacks, bind_request] if get_turn_context() else callbacks + + @property + def canonical_before_tool_callbacks(self): + from veadk.routing.tool_arguments import validate_tool_arguments + + callbacks = super().canonical_before_tool_callbacks + return ( + [*callbacks, validate_tool_arguments] if get_turn_context() else callbacks + ) + + @property + def _llm_flow(self): + return prepare_flow(super()._llm_flow) + + async def _run_async_impl(self, ctx): + execute = super()._run_async_impl + stream = ( + run_routed(self, ctx, execute) + if get_turn_context() or routing_owner(self) + else execute(ctx) + ) + async with aclosing(stream) as events: + async for event in events: + yield event + + +def _view(agent: Any) -> Any: + from veadk.agent import Agent + + if isinstance(agent, Agent): + result = agent.clone(update={"sub_agents": []}) + elif isinstance(agent, LlmAgent): + if type(agent) is not LlmAgent: + raise ValueError( + "custom ADK agents must implement turn-scoped model binding" + ) + values = {name: getattr(agent, name) for name in type(agent).model_fields} + result = RoutedLlmAgent.model_construct(**values) + result.__pydantic_private__ = copy.copy(agent.__pydantic_private__) + else: + result = agent.clone(update={"sub_agents": []}) + result.sub_agents = [_view(child) for child in agent.sub_agents] + for child in result.sub_agents: + child.parent_agent = result + return result + + +def inherit_tool(tool: Any) -> Any: + if not isinstance(tool, AgentTool) or get_turn_context() is None: + return tool + result = copy.copy(tool) + result.agent = _view(tool.agent) + return result + + +def bind_sub_agents(agent: Any) -> None: + """Own adapted copies so direct transfers also retain the root router.""" + agent.sub_agents = [_view(child) for child in agent.sub_agents] + for child in agent.sub_agents: + child.parent_agent = agent diff --git a/veadk/routing/callback.py b/veadk/routing/callback.py new file mode 100644 index 000000000..be1c88470 --- /dev/null +++ b/veadk/routing/callback.py @@ -0,0 +1,22 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Final request binding; semantic classification happens at root-turn entry.""" + +from typing import Any + + +def build_before_model_callback(classifier: Any, store: Any, config: Any) -> Any: + from veadk.routing.execution import bind_request + + return bind_request diff --git a/veadk/routing/capability.py b/veadk/routing/capability.py new file mode 100644 index 000000000..e67da72b8 --- /dev/null +++ b/veadk/routing/capability.py @@ -0,0 +1,75 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Local, deterministic capability checks before downgrading a request.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from veadk.configs.routing_configs import SimpleModelConfig +from veadk.routing.state import request_has_media, request_text_chars + + +@dataclass(frozen=True) +class CapabilityCheck: + """Whether a cheap model can serve one request.""" + + allowed: bool + reason: str = "" + + +class CapabilityGuard: + """Reject a cheap model when the request needs more than it declares. + + Unknown capability is denied by default. This is a hard, local gate: it + never calls the Decision API and never depends on model-name heuristics. + """ + + def check( + self, + llm_request: Any, + model: SimpleModelConfig, + ) -> CapabilityCheck: + if request_has_media(llm_request) and not model.supports_media: + return CapabilityCheck(False, "media_not_supported") + + if self._has_tools(llm_request) and not model.supports_tools: + return CapabilityCheck(False, "tools_not_supported") + + if ( + self._has_structured_output(llm_request) + and not model.supports_structured_output + ): + return CapabilityCheck(False, "structured_output_not_supported") + + if request_text_chars(llm_request) > model.max_context_chars: + return CapabilityCheck(False, "context_too_large") + + return CapabilityCheck(True) + + @staticmethod + def _has_tools(llm_request: Any) -> bool: + if getattr(llm_request, "tools_dict", None): + return True + config = getattr(llm_request, "config", None) + return bool(getattr(config, "tools", None)) + + @staticmethod + def _has_structured_output(llm_request: Any) -> bool: + config = getattr(llm_request, "config", None) + for field in ("response_schema", "response_json_schema", "output_schema"): + if getattr(config, field, None) is not None: + return True + return False diff --git a/veadk/routing/classifier.py b/veadk/routing/classifier.py new file mode 100644 index 000000000..3378d0547 --- /dev/null +++ b/veadk/routing/classifier.py @@ -0,0 +1,215 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Classify one root user turn and apply the conservative routing gates.""" + +from __future__ import annotations + +import asyncio +import time +import json +from typing import Any + +from veadk.configs.routing_configs import QueryRoutingConfig +from veadk.decision import DecisionClientError, DecisionQuestion +from veadk.routing.context import RouteSelection +from veadk.routing.policy import ConservativePolicy +from veadk.routing.registry import ModelRegistry +from veadk.routing.state import QueryStateBuilder, IncompleteRoutingState +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + + +class QueryClassifier: + """Run local gates, call the Decision API, and return one route label. + + Any timeout, transport failure, malformed answer, or missing capability + fails closed to the strong model. Classification has a bounded timeout; cancellation propagates. + """ + + def __init__( + self, + config: QueryRoutingConfig, + judge: Any, + registry: ModelRegistry, + *, + state_builder: QueryStateBuilder | None = None, + policy: ConservativePolicy | None = None, + ): + self._config = config + self._judge = judge + self._registry = registry + self._state_builder = state_builder or QueryStateBuilder() + self._policy = policy or ConservativePolicy(config.questions) + self._questions = { + question.id: DecisionQuestion( + type=question.type, + instructions=question.instructions, + ) + for question in config.questions + } + + @property + def registry(self) -> ModelRegistry: + """The model pool this classifier selects from.""" + + return self._registry + + async def aclose(self) -> None: + close = getattr(self._judge, "aclose", None) + if close is not None: + await close() + + async def classify( + self, + llm_request: Any, + *, + invocation_id: str, + session_id: str = "", + ) -> RouteSelection: + started_at = time.monotonic() + + runtime = self._registry.first_available(llm_request) + if runtime is None: + return self._strong("no_capable_simple_model", invocation_id, session_id) + + try: + state = self._state_builder.build_for_routing(llm_request) + except IncompleteRoutingState: + return self._strong("incomplete_state", invocation_id, session_id) + if not state: + return self._strong("empty_state", invocation_id, session_id) + + state.insert( + 0, + { + "type": "text", + "text": "Judge the CURRENT user turn using the history and task rules. " + "The LAST [user] message defines the current task. Earlier user " + "questions and assistant replies are context, not active requests. " + "Use earlier messages only to resolve what the current task asks for. " + "Candidate capabilities: " + + json.dumps( + runtime.config.model_dump( + include={ + "name", + "capability_description", + "supports_tools", + "supports_media", + "supports_structured_output", + "max_context_chars", + } + ) + ), + }, + ) + + try: + response = await asyncio.wait_for( + self._judge.judge(state=state, questions=self._questions), + timeout=self._config.decision_timeout, + ) + except asyncio.TimeoutError: + return self._strong("decision_timeout", invocation_id, session_id) + except DecisionClientError as exc: + logger.info( + "query routing decision failed: invocation_id=%s error_type=%s", + invocation_id, + exc.error_type, + ) + return self._strong("decision_error", invocation_id, session_id) + except Exception as exc: # noqa: BLE001 - fail closed, never fail the turn + logger.warning( + "query routing decision raised: invocation_id=%s error=%s", + invocation_id, + type(exc).__name__, + ) + return self._strong("decision_error", invocation_id, session_id) + + decision = self._policy.evaluate(getattr(response, "answers", None)) + elapsed_ms = round((time.monotonic() - started_at) * 1000, 2) + + raw = getattr(response, "raw", None) or {} + usage = raw.get("usage") or {} + input_tokens = usage.get("input_tokens") + if ( + not isinstance(input_tokens, int) + or isinstance(input_tokens, bool) + or input_tokens < 0 + ): + input_tokens = None + + if decision.route == "strong": + logger.info( + "query routing decision: invocation_id=%s route=strong reason=%s " + "latency_ms=%s probabilities=%s", + invocation_id, + decision.reason, + elapsed_ms, + decision.probabilities, + ) + return RouteSelection( + route="strong", + reason=decision.reason, + invocation_id=invocation_id, + session_id=session_id, + latency_ms=elapsed_ms, + input_tokens=input_tokens, + ) + + if self._config.mode == "shadow": + logger.info( + "query routing decision: invocation_id=%s route=strong " + "shadow_route=simple reason=shadow latency_ms=%s probabilities=%s", + invocation_id, + elapsed_ms, + decision.probabilities, + ) + return RouteSelection( + route="strong", + reason="shadow", + invocation_id=invocation_id, + session_id=session_id, + model_name=runtime.config.name, + shadow=True, + latency_ms=elapsed_ms, + input_tokens=input_tokens, + ) + + logger.info( + "query routing decision: invocation_id=%s route=simple model=%s " + "latency_ms=%s probabilities=%s", + invocation_id, + runtime.config.name, + elapsed_ms, + decision.probabilities, + ) + return RouteSelection( + route="simple", + reason=decision.reason, + invocation_id=invocation_id, + session_id=session_id, + model_name=runtime.config.name, + latency_ms=elapsed_ms, + input_tokens=input_tokens, + ) + + @staticmethod + def _strong(reason: str, invocation_id: str, session_id: str) -> RouteSelection: + return RouteSelection( + route="strong", + reason=reason, + invocation_id=invocation_id, + session_id=session_id, + ) diff --git a/veadk/routing/context.py b/veadk/routing/context.py new file mode 100644 index 000000000..6c7580691 --- /dev/null +++ b/veadk/routing/context.py @@ -0,0 +1,76 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Root-turn model selection propagated through nested execution. + +ADK's ``BaseLlm.generate_content_async`` only receives an ``LlmRequest``, so +the invocation-scoped decision is carried in a context variable. Each ADK +invocation runs in its own asyncio task, which keeps parallel sub-agents +isolated without mutating the shared Agent. +""" + +from __future__ import annotations + +from contextvars import ContextVar +from dataclasses import dataclass, field +from typing import Any, Literal + + +@dataclass(frozen=True) +class RouteSelection: + """The route chosen for one root user turn.""" + + route: Literal["simple", "strong"] + reason: str + invocation_id: str + session_id: str = "" + model_name: str | None = None + shadow: bool = False + latency_ms: float = 0.0 + input_tokens: int | None = None + + +@dataclass(frozen=True) +class TurnContext: + """One immutable model binding inherited by nested execution tasks.""" + + selection: RouteSelection + model: Any + router: Any + target: str + # Shared by nested tasks, owned by this root turn; the model binding stays fixed. + tool_argument_errors: dict[str, int] = field(default_factory=dict, compare=False) + + +_TURN: ContextVar[TurnContext | None] = ContextVar("veadk_query_turn", default=None) + + +def get_turn_context() -> TurnContext | None: + return _TURN.get() + + +_ROUTE_SELECTION: ContextVar[RouteSelection | None] = ContextVar( + "veadk_route_selection", default=None +) + + +def set_route_selection(selection: RouteSelection | None) -> None: + """Set the route selection for the current task.""" + + _ROUTE_SELECTION.set(selection) + + +def get_route_selection() -> RouteSelection | None: + """Return the route selection for the current task, if any.""" + + return _ROUTE_SELECTION.get() diff --git a/veadk/routing/execution.py b/veadk/routing/execution.py new file mode 100644 index 000000000..e9b88bdc0 --- /dev/null +++ b/veadk/routing/execution.py @@ -0,0 +1,351 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Turn-scoped model binding and portable Responses continuation. + +Only opted-in execution enters this module. The concrete model is visible to +ADK before its processors run; the full request survives until the transport +chooses a validated continuation or a new chain. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +from contextlib import aclosing, asynccontextmanager +from contextvars import ContextVar +from dataclasses import asdict, dataclass +from typing import Any + +from google.adk.events import Event, EventActions +from google.adk.flows.llm_flows import ( + contents, + interactions_processor, + instructions, + basic, +) +from google.adk.models.llm_request import LlmRequest +from google.genai import types +from google.adk.tools import BaseTool, FunctionTool +from google.adk.tools.base_toolset import BaseToolset + +from veadk.routing.context import RouteSelection, TurnContext, _TURN, get_turn_context + +_METADATA = "veadk_query_routing" +_RECORD = "veadk:query_routing:turn" + + +def digest(value: Any) -> str: + encoded = json.dumps(value, sort_keys=True, default=str, ensure_ascii=False) + return hashlib.sha256(encoded.encode()).hexdigest() + + +def target_id(model: Any) -> str: + # Only the digest leaves this function; credentials never enter metadata. + return digest( + ( + type(model).__module__, + type(model).__qualname__, + model.model, + getattr(model, "_additional_args", {}), + getattr(model, "fallbacks", None), + getattr(model, "enable_responses_cache", None), + ) + ) + + +def _content_data(items: list[Any]) -> list[Any]: + return [item.model_dump(exclude_none=True) for item in items] + + +@dataclass +class RequestState: + contents: list[Any] + config_hash: str = "" + + +_REQUEST: ContextVar[RequestState | None] = ContextVar( + "veadk_routing_request", default=None +) + + +def prepare_flow(flow: Any) -> Any: + if get_turn_context() is not None: + # Keep contents in full, including function IDs. Continuation is chosen + # after all processors and user callbacks, never before history replay. + flow.request_processors = [ + p + for p in flow.request_processors + if p is not interactions_processor.request_processor + ] + return flow + + +def _previous_id( + callback_context: Any, request: LlmRequest, turn: TurnContext, config_hash: str +) -> str | None: + ctx = callback_context._invocation_context + for event in reversed(ctx.session.events): + if event.branch != ctx.branch or event.isolation_scope != ctx.isolation_scope: + continue + if not event.content or event.content.role != "model" or event.partial: + continue + # An intervening model/agent response invalidates the old chain. + metadata = (event.custom_metadata or {}).get(_METADATA, {}) + size = metadata.get("prefix_size", 0) + if ( + event.author != ctx.agent.name + or metadata.get("target") != turn.target + or metadata.get("config_hash") != config_hash + or not size + ): + return None + if digest(_content_data(request.contents[:size])) != metadata.get( + "prefix_hash" + ): + return None + return event.interaction_id + return None + + +def _portable_contents(ctx: Any, items: list[Any], target: str) -> list[Any]: + """Replay business content without another target's opaque reasoning state.""" + owned_parts = { + digest(part.model_dump(exclude_none=True)) + for event in ctx.session.events + if (event.custom_metadata or {}).get(_METADATA, {}).get("target") == target + and event.content + for part in event.content.parts or [] + } + result = [] + for content in copy.deepcopy(items): + portable = [] + for part in content.parts or []: + if digest(part.model_dump(exclude_none=True)) not in owned_parts: + part.thought_signature = None + if part.thought: + if not (part.function_call or part.function_response): + continue + part.thought = None + if part.model_dump(exclude_none=True): + portable.append(part) + if portable: + content.parts = portable + result.append(content) + return result + + +async def bind_request(callback_context: Any, llm_request: LlmRequest) -> None: + turn = get_turn_context() + state = _REQUEST.get() + if turn is None or state is None: + return + from veadk.routing.agent_tool import inherit_tool + + if target_id(turn.model) != turn.target: + raise ValueError("routing model configuration changed during the turn") + replacements = { + id(tool): inherit_tool(tool) for tool in llm_request.tools_dict.values() + } + llm_request.tools_dict = { + name: replacements[id(tool)] for name, tool in llm_request.tools_dict.items() + } + ctx = callback_context._invocation_context + if ctx.canonical_tools_cache is not None: + ctx.canonical_tools_cache = [ + replacements.get(id(tool), tool) for tool in ctx.canonical_tools_cache + ] + if turn.selection.route == "simple": + runtime = turn.router.registry.get(turn.selection.model_name) + check = turn.router.registry.check(llm_request, runtime) + if not check.allowed: + raise ValueError( + f"fixed routing model cannot serve this request: {check.reason}" + ) + llm_request.contents = _portable_contents(ctx, llm_request.contents, turn.target) + llm_request.model = turn.model.model + llm_request.previous_interaction_id = None + state.contents = copy.deepcopy(llm_request.contents) + state.config_hash = digest(llm_request.config.model_dump(exclude_none=True)) + if getattr(turn.model, "enable_responses_cache", False): + llm_request.previous_interaction_id = _previous_id( + callback_context, llm_request, turn, state.config_hash + ) + + +def _tag_event(event: Event, ctx: Any, turn: TurnContext) -> None: + state = _REQUEST.get() + if ( + state is None + or event.author != ctx.agent.name + or event.partial + or not event.content + or event.content.role != "model" + ): + return + prefix = [*state.contents, event.content] + metadata = { + "root_turn_id": turn.selection.invocation_id, + "route": turn.selection.route, + "model": turn.model.model, + "target": turn.target, + "prefix_size": len(prefix), + "prefix_hash": digest(_content_data(prefix)), + "config_hash": state.config_hash, + } + event.custom_metadata = {**(event.custom_metadata or {}), _METADATA: metadata} + + +def _agents(root: Any) -> list[Any]: + pending, found, seen = [root], [], set() + while pending: + agent = pending.pop() + if id(agent) in seen: + continue + seen.add(id(agent)) + found.append(agent) + pending.extend(getattr(agent, "sub_agents", ())) + pending.extend( + t.agent for t in getattr(agent, "tools", ()) if hasattr(t, "agent") + ) + return found + + +def routing_owner(agent: Any) -> Any: + """Find the outermost configured router, including resumed transfer targets.""" + from veadk.routing.llm import RoutingLlm + + owner = None + while agent is not None: + if isinstance(getattr(agent, "model", None), RoutingLlm): + owner = agent + agent = agent.parent_agent + return owner + + +async def routing_request(agent: Any, ctx: Any) -> LlmRequest: + request = LlmRequest() + async for _ in contents.request_processor.run_async(ctx, request): + pass + async for _ in basic.request_processor.run_async(ctx, request): + pass + async for _ in instructions.request_processor.run_async(ctx, request): + pass + for member in _agents(routing_owner(agent) or agent): + for tool in getattr(member, "tools", ()): + name = getattr(tool, "name", None) or getattr( + tool, "__name__", type(tool).__name__ + ) + request.tools_dict[f"{member.name}.{name}"] = tool + if isinstance(tool, BaseToolset): + # Tools discovered during execution cannot be classified from a + # guessed schema; keep the strong model for this turn. + request.config.labels = {"routing_incomplete": "dynamic_toolset"} + else: + concrete = tool if isinstance(tool, BaseTool) else FunctionTool(tool) + declaration = concrete._get_declaration() + if declaration: + request.config.tools = [ + *(request.config.tools or []), + types.Tool(function_declarations=[declaration]), + ] + if getattr(member, "output_schema", None): + request.config.response_json_schema = ( + member.output_schema.model_json_schema() + ) + return request + + +async def _select(agent: Any, ctx: Any, router: Any) -> TurnContext: + config_hash = configuration_id(router) + record = ctx.session.state.get(_RECORD, {}) + if record.get("invocation_id") == ctx.invocation_id: + if record.get("config_hash") != config_hash: + raise ValueError("routing configuration changed during a resumed turn") + selection = RouteSelection(**record["selection"]) + else: + request = await routing_request(agent, ctx) + key = f"{ctx.app_name}:{ctx.user_id}:{ctx.session.id}:{ctx.invocation_id}:{config_hash}" + selection = await router.store.get_or_create( + key, + lambda: router.classifier.classify( + request, + invocation_id=ctx.invocation_id, + session_id=ctx.session.id, + ), + ) + model = router.selected_model(selection) + return TurnContext(selection, model, router, target_id(model)) + + +def configuration_id(router: Any) -> str: + return digest( + ( + router.config.model_dump(mode="json"), + target_id(router.strong_model), + [target_id(r.llm) for r in router.runtimes], + ) + ) + + +@asynccontextmanager +async def _scope(agent: Any, ctx: Any): + inherited = get_turn_context() + router = inherited.router if inherited else routing_owner(agent).model + if inherited is None: + router.active_runs += 1 + try: + turn = inherited or await _select(agent, ctx, router) + token = _TURN.set(turn) + request_token = _REQUEST.set(RequestState([])) + try: + yield turn, inherited is None + finally: + _REQUEST.reset(request_token) + _TURN.reset(token) + finally: + if inherited is None: + router.active_runs -= 1 + if router.active_runs == 0: + await router.classifier.aclose() + + +async def run_routed(agent: Any, ctx: Any, execute: Any): + async with _scope(agent, ctx) as (turn, is_root): + if is_root: + record = { + "invocation_id": ctx.invocation_id, + "config_hash": configuration_id(turn.router), + "selection": asdict(turn.selection), + } + yield Event( + author=agent.name, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + actions=EventActions(state_delta={_RECORD: record}), + custom_metadata={ + _METADATA: { + "root_turn_id": ctx.invocation_id, + "route": turn.selection.route, + "model": turn.model.model, + "reason": turn.selection.reason, + "latency_ms": turn.selection.latency_ms, + "input_tokens": turn.selection.input_tokens, + } + }, + ) + async with aclosing(execute(ctx)) as events: + async for event in events: + _tag_event(event, ctx, turn) + yield event diff --git a/veadk/routing/factory.py b/veadk/routing/factory.py new file mode 100644 index 000000000..f5f2c4a01 --- /dev/null +++ b/veadk/routing/factory.py @@ -0,0 +1,197 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Build routing components from configuration and Agent model settings.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from google.adk.models.base_llm import BaseLlm + +from veadk.configs.routing_configs import ( + QueryRoutingConfig, + QueryRoutingConfigError, + load_query_routing_config, + load_query_routing_config_if_present, +) +from veadk.decision import ( + DecisionClient, + DecisionConfig, + load_decision_config, + resolve_decision_connection, +) +from veadk.routing.callback import build_before_model_callback +from veadk.routing.classifier import QueryClassifier +from veadk.routing.llm import RoutingLlm +from veadk.routing.registry import ModelRegistry, SimpleModelRuntime +from veadk.routing.store import RouteStore, get_route_store + + +def resolve_query_routing_config( + value: Any, +) -> QueryRoutingConfig | None: + """Resolve the Agent's ``query_routing`` value. + + ``None`` auto-detects a ``query_routing`` block in ``config.yaml``; + ``False`` disables routing explicitly; ``True`` requires the YAML block. + A mapping or a :class:`QueryRoutingConfig` supplies configuration directly. + """ + + if value is False: + return None + if value is None: + return load_query_routing_config_if_present() + if value is True: + return load_query_routing_config() + if isinstance(value, QueryRoutingConfig): + return value + if isinstance(value, Mapping): + return QueryRoutingConfig.model_validate(dict(value)) + raise QueryRoutingConfigError( + "query_routing must be None, True, False, a mapping, or QueryRoutingConfig" + ) + + +def build_query_routing( + config: QueryRoutingConfig, + agent: Any, + *, + decision: DecisionConfig | None = None, + store: RouteStore | None = None, + judge: Any | None = None, +) -> tuple[RoutingLlm, Any]: + """Build the model pool and final request binder. + + ``judge`` is injectable for tests; production uses a ``DecisionClient`` + whose Ark endpoint is derived from the Agent's effective model connection. + """ + + strong_model = _strong_model(config, agent) + _validate_model_contract(strong_model) + runtimes = [ + SimpleModelRuntime(config=model, llm=_build_model(agent, model)) + for model in config.simple_models + ] + registry = ModelRegistry(runtimes=runtimes) + route_store = store or get_route_store() + + if judge is None: + connection = decision if decision is not None else load_decision_config() + endpoint, api_key = resolve_decision_connection( + connection, + getattr(agent, "model_api_base", None), + getattr(agent, "model_api_key", None), + ) + judge = DecisionClient( + endpoint=endpoint, + api_key=api_key, + model=connection.model, + timeout=config.decision_timeout, + ) + + classifier = QueryClassifier(config, judge, registry) + router = RoutingLlm( + model=getattr(strong_model, "model", config.strong_model or ""), + strong_model=strong_model, + runtimes=runtimes, + registry=registry, + store=route_store, + classifier=classifier, + config=config, + ) + callback = build_before_model_callback(classifier, route_store, config) + return router, callback + + +def _strong_model(config: QueryRoutingConfig, agent: Any) -> BaseLlm: + strong_name = config.strong_model + primary = _primary_model_name(agent) + if strong_name is None or strong_name == primary: + return agent.model + return _build_model(agent, strong_name) + + +def _primary_model_name(agent: Any) -> str | None: + model_name = getattr(agent, "model_name", None) + if isinstance(model_name, list): + return model_name[0] if model_name else None + return model_name + + +def _validate_model_contract(model: BaseLlm) -> None: + from veadk.configs.routing_configs import _REQUEST_OVERRIDE_FIELDS + + extra = getattr(model, "_additional_args", {}) or {} + body = extra.get("extra_body") or {} + if _REQUEST_OVERRIDE_FIELDS.intersection(extra) or ( + _REQUEST_OVERRIDE_FIELDS | {"model"} + ).intersection(body): + raise QueryRoutingConfigError( + "model transport options must not override routed requests" + ) + if getattr(model, "fallbacks", None) or extra.get("fallbacks"): + raise QueryRoutingConfigError( + "query routing fixes the model for the entire turn; remove model_fallbacks" + ) + + +def _build_model(agent: Any, specification: Any) -> BaseLlm: + from veadk.configs.routing_configs import SimpleModelConfig + from veadk.models.ark_llm import ArkLlm + from veadk.models.retrying_lite_llm import RetryingLiteLlm + + spec = ( + SimpleModelConfig(name=specification) + if isinstance(specification, str) + else specification + ) + provider = spec.provider or getattr(agent, "model_provider", "openai") + full_name = spec.name if "/" in spec.name else f"{provider}/{spec.name}" + api_base = spec.api_base or getattr(agent, "model_api_base", "") + explicit_key = spec.api_key_value + if ( + spec.api_base + and spec.api_base.rstrip("/") + != getattr(agent, "model_api_base", "").rstrip("/") + and not explicit_key + ): + raise QueryRoutingConfigError( + "an independent api_base requires api_key or api_key_env" + ) + api_key = explicit_key or getattr(agent, "model_api_key", "") + # Endpoint-specific defaults must not cross provider boundaries. + extra = ( + dict(getattr(agent, "model_extra_config", {}) or {}) + if not spec.api_base + else {} + ) + extra.update(spec.extra_config) + responses = ( + spec.api == "responses" + if spec.api + else getattr(agent, "enable_responses", False) + ) + kwargs = dict( + model=full_name, api_key=api_key, api_base=api_base, fallbacks=None, **extra + ) + if responses: + model = ArkLlm( + enable_responses_cache=getattr(agent, "enable_responses_cache", True), + **kwargs, + ) + else: + model = RetryingLiteLlm(**kwargs) + _validate_model_contract(model) + return model diff --git a/veadk/routing/llm.py b/veadk/routing/llm.py new file mode 100644 index 000000000..bfd4f93b9 --- /dev/null +++ b/veadk/routing/llm.py @@ -0,0 +1,53 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Configured model pool; execution resolves a concrete adapter before preprocessing.""" + +from __future__ import annotations + +from typing import Any + +from google.adk.models.base_llm import BaseLlm +from pydantic import Field + +from veadk.routing.context import get_route_selection + + +class RoutingLlm(BaseLlm): + """Own the configured adapters without storing a mutable current model.""" + + strong_model: BaseLlm + runtimes: list[Any] = Field(default_factory=list) + registry: Any = None + store: Any = None + classifier: Any = None + config: Any = None + active_runs: int = 0 + + def selected_model(self, selection: Any) -> BaseLlm: + if selection.route == "strong": + return self.strong_model + runtime = self.registry.get(selection.model_name) + if runtime is None: + raise ValueError("the fixed routing target is no longer configured") + return runtime.llm + + async def generate_content_async(self, llm_request, stream=False): + # Direct use remains possible; Agent resolves the concrete model earlier. + selection = get_route_selection() + model = self.selected_model(selection) if selection else self.strong_model + request = llm_request.model_copy(deep=True) + request.tools_dict = dict(llm_request.tools_dict) + request.model = model.model + async for response in model.generate_content_async(request, stream=stream): + yield response diff --git a/veadk/routing/policy.py b/veadk/routing/policy.py new file mode 100644 index 000000000..ee96636e7 --- /dev/null +++ b/veadk/routing/policy.py @@ -0,0 +1,99 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conservative gate applied to Decision API answers.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Literal + +from veadk.configs.routing_configs import RoutingQuestion + + +@dataclass(frozen=True) +class RoutingDecision: + """The outcome of the conservative gate.""" + + route: Literal["simple", "strong"] + reason: str + probabilities: dict[str, float] = field(default_factory=dict) + + +class ConservativePolicy: + """Require every question to stay at or below its threshold. + + A missing, malformed, or non-``noul`` answer fails closed: the request goes + to the strong model instead of guessing. + """ + + def __init__(self, questions: list[RoutingQuestion]): + if not questions: + raise ValueError("at least one routing question is required") + self._questions = tuple(questions) + + @property + def questions(self) -> tuple[RoutingQuestion, ...]: + return self._questions + + def evaluate(self, answers: Any) -> RoutingDecision: + probabilities: dict[str, float] = {} + + if not isinstance(answers, Mapping): + return RoutingDecision( + route="strong", + reason="invalid_answers", + probabilities={}, + ) + + for question in self._questions: + answer = answers.get(question.id) + value = self._probability(answer) + if value is None: + return RoutingDecision( + route="strong", + reason=f"invalid_answer:{question.id}", + probabilities=probabilities, + ) + probabilities[question.id] = value + if value > question.threshold: + return RoutingDecision( + route="strong", + reason=f"{question.id}_above_threshold", + probabilities=probabilities, + ) + + return RoutingDecision( + route="simple", + reason="all_thresholds_passed", + probabilities=probabilities, + ) + + @staticmethod + def _probability(answer: Any) -> float | None: + if not isinstance(answer, Mapping): + return None + if answer.get("type") != "noul": + return None + value = answer.get("noul") + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + try: + probability = float(value) + except OverflowError: + return None + if not math.isfinite(probability) or not 0.0 <= probability <= 1.0: + return None + return probability diff --git a/veadk/routing/registry.py b/veadk/routing/registry.py new file mode 100644 index 000000000..a5673054d --- /dev/null +++ b/veadk/routing/registry.py @@ -0,0 +1,71 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Per-tier model instances and their local health state.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Any + +from google.adk.models.base_llm import BaseLlm + +from veadk.configs.routing_configs import SimpleModelConfig +from veadk.routing.capability import CapabilityGuard, CapabilityCheck + + +@dataclass +class SimpleModelRuntime: + """A configured cheap model and the client instance bound to it.""" + + config: SimpleModelConfig + llm: BaseLlm + unhealthy_until: float = 0.0 + + +@dataclass +class ModelRegistry: + """Ordered cheap-model pool plus per-model capability/health checks. + + Selection is deterministic: the first model in configuration order that + passes ``CapabilityGuard`` and is not circuit-broken wins. There is no + random or round-robin dispatch. + """ + + runtimes: list[SimpleModelRuntime] + guard: CapabilityGuard = field(default_factory=CapabilityGuard) + _clock: Any = field(default=time.monotonic, repr=False) + + def first_available(self, llm_request: Any) -> SimpleModelRuntime | None: + now = self._clock() + for runtime in self.runtimes: + if runtime.unhealthy_until > now: + continue + if self.guard.check(llm_request, runtime.config).allowed: + return runtime + return None + + def check(self, llm_request: Any, runtime: SimpleModelRuntime) -> CapabilityCheck: + return self.guard.check(llm_request, runtime.config) + + def mark_failed( + self, runtime: SimpleModelRuntime, *, cooldown: float = 60.0 + ) -> None: + runtime.unhealthy_until = self._clock() + max(0.0, cooldown) + + def get(self, name: str) -> SimpleModelRuntime | None: + for runtime in self.runtimes: + if runtime.config.name == name: + return runtime + return None diff --git a/veadk/routing/state.py b/veadk/routing/state.py new file mode 100644 index 000000000..b3f6530d6 --- /dev/null +++ b/veadk/routing/state.py @@ -0,0 +1,129 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Build the bounded state sent to the Decision API for query routing. + +The builder reuses the content-safety media conversion (image/audio/video, +``file_id://`` inlining, 20 MiB and 32K-character limits) so routing and +auditing speak the same System One state format. +""" + +from __future__ import annotations + +from typing import Any +import json + +from veadk.content_safety.state import ( + MAX_STATE_CHARS, + StateBuilder, + _media_state_part, + _parts_segments, +) + +DEFAULT_MAX_TURNS = 10_000 +_MEDIA_CHAR_COST = 4_000 + + +class IncompleteRoutingState(ValueError): + """The decision input cannot faithfully represent the effective context.""" + + +class QueryStateBuilder(StateBuilder): + """Build routing state from the recent conversation and media parts. + + Unlike the content-safety builder, which inspects only the latest user + message, routing needs enough history to judge whether the request depends + on a long context. The window is bounded by turn count and characters. + """ + + def __init__( + self, + *, + max_chars: int = MAX_STATE_CHARS, + max_turns: int = DEFAULT_MAX_TURNS, + defuse: bool = True, + ): + super().__init__(max_chars=max_chars, defuse=defuse) + self._max_turns = max(1, max_turns) + + def build_for_routing(self, llm_request: Any) -> list[dict[str, Any]]: + contents = list(getattr(llm_request, "contents", None) or []) + segments: list[str | dict[str, Any]] = [] + if len(contents) > self._max_turns: + raise IncompleteRoutingState("history exceeds the routing window") + config = getattr(llm_request, "config", None) + if (getattr(config, "labels", None) or {}).get("routing_incomplete"): + raise IncompleteRoutingState("dynamic tool declarations are not available") + instruction = getattr(config, "system_instruction", None) + if instruction: + segments.append(f"[agent_instructions]\n{instruction}") + names = list(getattr(llm_request, "tools_dict", {}) or {}) + if names: + segments.append("[available_tools]\n" + json.dumps(names)) + for field in ("tools", "response_schema", "response_json_schema"): + value = getattr(config, field, None) + if value: + segments.append(f"[{field}]\n{value}") + for content in contents[-self._max_turns :]: + for part in getattr(content, "parts", None) or []: + if ( + getattr(part, "file_data", None) + or getattr(part, "inline_data", None) + ) and _media_state_part(part) is None: + raise IncompleteRoutingState("media cannot be represented") + parts = _parts_segments(getattr(content, "parts", None) or []) + if not parts: + continue + role = getattr(content, "role", None) or "user" + segments.append(f"[{role}]") + segments.extend(parts) + if sum(len(s) for s in segments if isinstance(s, str)) > self._max_chars: + raise IncompleteRoutingState("routing text exceeds the character budget") + return self._finalize_segments(segments) + + +def request_has_media(llm_request: Any) -> bool: + """Return whether the request carries image, audio, video, or file input.""" + + for content in getattr(llm_request, "contents", None) or []: + for part in getattr(content, "parts", None) or []: + if ( + getattr(part, "inline_data", None) + or getattr(part, "file_data", None) + or _media_state_part(part) is not None + ): + return True + return False + + +def request_text_chars(llm_request: Any) -> int: + """Approximate the request size in characters for the context gate.""" + + total = 0 + config = getattr(llm_request, "config", None) + if config is not None: + total += len(str(getattr(config, "system_instruction", "") or "")) + total += len(str(getattr(config, "tools", "") or "")) + total += len(str(getattr(config, "response_json_schema", "") or "")) + for content in getattr(llm_request, "contents", None) or []: + for part in getattr(content, "parts", None) or []: + text = getattr(part, "text", None) + if isinstance(text, str): + total += len(text) + elif getattr(part, "function_call", None) or getattr( + part, "function_response", None + ): + total += len(str(part)) + elif _media_state_part(part) is not None: + total += _MEDIA_CHAR_COST + return total diff --git a/veadk/routing/store.py b/veadk/routing/store.py new file mode 100644 index 000000000..58d52615c --- /dev/null +++ b/veadk/routing/store.py @@ -0,0 +1,106 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Root-invocation decision sharing.""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Awaitable, Callable + +from veadk.routing.context import RouteSelection + +DEFAULT_DECISION_TTL_SECONDS = 300.0 +_MAX_ENTRIES = 4_096 + + +class RouteStore: + """Coalesce concurrent decisions; active turns keep their own immutable binding.""" + + def __init__( + self, + *, + decision_ttl: float = DEFAULT_DECISION_TTL_SECONDS, + clock: Callable[[], float] = time.monotonic, + ): + self._decision_ttl = decision_ttl + self._clock = clock + self._decisions: dict[str, tuple[float, RouteSelection]] = {} + self._inflight: dict[str, asyncio.Future[RouteSelection]] = {} + + async def get_or_create( + self, + key: str, + factory: Callable[[], Awaitable[RouteSelection]], + ) -> RouteSelection: + """Return a cached decision, join an in-flight one, or classify once.""" + + self._purge() + cached = self._decisions.get(key) + if cached is not None: + created_at, selection = cached + if self._clock() - created_at <= self._decision_ttl: + return selection + self._decisions.pop(key, None) + + inflight = self._inflight.get(key) + if inflight is not None: + return await asyncio.shield(inflight) + + loop = asyncio.get_running_loop() + future: asyncio.Future[RouteSelection] = loop.create_future() + self._inflight[key] = future + try: + selection = await factory() + except BaseException as exc: + if not future.done(): + future.set_exception(exc) + # Mark the exception as retrieved so an unawaited future does + # not log "exception was never retrieved" at GC time. + future.exception() + self._inflight.pop(key, None) + raise + + self._decisions[key] = (self._clock(), selection) + self._inflight.pop(key, None) + if not future.done(): + future.set_result(selection) + return selection + + def clear(self) -> None: + """Drop all cached state. Intended for tests and process shutdown.""" + + self._decisions.clear() + self._inflight.clear() + + def _purge(self) -> None: + if len(self._decisions) <= _MAX_ENTRIES: + return + now = self._clock() + self._decisions = { + key: value + for key, value in self._decisions.items() + if now - value[0] <= self._decision_ttl + } + while len(self._decisions) > _MAX_ENTRIES: + self._decisions.pop(next(iter(self._decisions))) + + +_DEFAULT_STORE = RouteStore() + + +def get_route_store() -> RouteStore: + """Return the process-wide route store used by Agents by default.""" + + return _DEFAULT_STORE diff --git a/veadk/routing/tool_arguments.py b/veadk/routing/tool_arguments.py new file mode 100644 index 000000000..98edb61a9 --- /dev/null +++ b/veadk/routing/tool_arguments.py @@ -0,0 +1,143 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Validate model-produced tool arguments against the advertised contract. + +The guard runs after user callbacks and before execution, only in routed turns. +It never retries a tool that has already executed or changes the selected model. +""" + +from __future__ import annotations + +import copy +import re +from typing import Any + +from jsonschema import Draft202012Validator + +from veadk.routing.context import get_turn_context + +_BUDGET_KEY = "veadk:query_routing:tool_argument_errors" +_MAX_CORRECTIONS = 2 + + +def _schema(value: Any) -> Any: + """Translate GenAI schema types and nullable into JSON Schema.""" + if isinstance(value, list): + return [_schema(item) for item in value] + if not isinstance(value, dict): + return value + result = copy.deepcopy(value) + # Only schema-bearing keywords are recursive. Defaults, const and enum + # may contain arbitrary objects whose own 'type' keys are user data. + for key in ( + "properties", + "patternProperties", + "$defs", + "definitions", + "dependentSchemas", + ): + if isinstance(result.get(key), dict): + result[key] = {name: _schema(item) for name, item in result[key].items()} + for key in ( + "items", + "prefixItems", + "additionalProperties", + "unevaluatedProperties", + "unevaluatedItems", + "contains", + "propertyNames", + "not", + "if", + "then", + "else", + "anyOf", + "allOf", + "oneOf", + ): + if key in result: + result[key] = _schema(result[key]) + for key in ("$ref", "$dynamicRef"): + if key in result and not result[key].startswith("#"): + raise ValueError("routing tool schemas must use local references") + if isinstance(result.get("type"), str): + result["type"] = result["type"].lower() + if result.pop("nullable", False): + return {"anyOf": [result, {"type": "null"}]} + return result + + +async def validate_tool_arguments(tool, args, tool_context): + turn = get_turn_context() + if turn is None or turn.selection.shadow: + return None + declaration = tool._get_declaration() + if declaration is None: + return None + schema = declaration.parameters_json_schema + if schema is None and declaration.parameters is not None: + schema = _schema(declaration.parameters.model_dump(exclude_none=True)) + if schema is None: + return None + # Check references in JSON schemas as well; validation must not perform IO. + schema = _schema(schema) + Draft202012Validator.check_schema(schema) + validator = Draft202012Validator(schema) + normalized = copy.deepcopy(args) + for error in validator.iter_errors(normalized): + # Only an unambiguous integer declaration permits normalization. In + # particular, never truncate fractions or convert booleans to integers. + if ( + error.validator == "type" + and error.validator_value == "integer" + and isinstance(error.instance, str) + and re.fullmatch(r"-?(?:0|[1-9][0-9]{0,127})", error.instance) + and error.absolute_path + ): + parent = normalized + path = list(error.absolute_path) + for key in path[:-1]: + parent = parent[key] + parent[path[-1]] = int(error.instance) + errors = list(validator.iter_errors(normalized)) + if not errors: + args.clear() + args.update(normalized) + return None + + record = tool_context.state.get(_BUDGET_KEY, {}) + previous = ( + record.get("count", 0) + if record.get("invocation_id") == turn.selection.invocation_id + else 0 + ) + count = max(previous, turn.tool_argument_errors.get("count", 0)) + 1 + turn.tool_argument_errors["count"] = count + tool_context.state[_BUDGET_KEY] = { + "invocation_id": turn.selection.invocation_id, + "count": count, + } + if count > _MAX_CORRECTIONS: + raise ValueError("routing tool argument correction budget exhausted") + return { + "error": "invalid_tool_arguments", + "retryable": True, + "corrections_remaining": _MAX_CORRECTIONS - count, + "instruction": "Correct the arguments using the tool schema before retrying.", + # Error messages/instances can contain secrets. Only schema locations + # and rule names are returned; argument values never enter this result. + "issues": [ + {"schema_path": list(error.absolute_schema_path), "rule": error.validator} + for error in errors[:8] + ], + } diff --git a/veadk/utils/misc.py b/veadk/utils/misc.py index 4a0d1b1e1..548aed60e 100644 --- a/veadk/utils/misc.py +++ b/veadk/utils/misc.py @@ -137,16 +137,18 @@ def set_envs(config_yaml_path: str, env_from_dotenv: dict = None) -> tuple[dict, with open(config_yaml_path, "r", encoding="utf-8") as yaml_file: config_dict = safe_load(yaml_file) - # The content-safety subtree is structured runtime configuration, not a - # source of environment variables. Keep its questions/decision values out - # of the flattened env so nested data is not lossy or surprising. + # Structured feature and shared connection blocks stay local to each + # consumer; exporting them would leak configuration into later agents. env_config_dict = config_dict if isinstance(config_dict, dict): tool_config = config_dict.get("tool") has_tool_block = ( isinstance(tool_config, dict) and "content_safety" in tool_config ) - has_root_block = "content_safety" in config_dict + has_root_block = any( + key in config_dict + for key in ("content_safety", "query_routing", "decision") + ) if has_tool_block or has_root_block: env_config_dict = dict(config_dict) if has_tool_block: @@ -156,7 +158,8 @@ def set_envs(config_yaml_path: str, env_from_dotenv: dict = None) -> tuple[dict, env_config_dict["tool"] = remaining_tool else: env_config_dict.pop("tool", None) - env_config_dict.pop("content_safety", None) + for key in ("content_safety", "query_routing", "decision"): + env_config_dict.pop(key, None) flatten_config_dict = flatten_dict(env_config_dict) config_upper_map = {k.upper(): v for k, v in flatten_config_dict.items()}