diff --git a/.env.example b/.env.example index 4cff4556a..67368a805 100644 --- a/.env.example +++ b/.env.example @@ -70,6 +70,21 @@ POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_NORMALIZATION=unit POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BATCH_SIZE=10 POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_TIMEOUT_SECONDS=30 +# Optional workload-specific endpoint overrides. +# POWERCONTEXT_SERVER_INFERENCE_GENERATION_BASE_URL=http://127.0.0.1:8080/v1 +# POWERCONTEXT_SERVER_INFERENCE_GENERATION_HEADERS={"Authorization":"Bearer replace-me"} +# POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS={"max_tokens":4096} +# POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BASE_URL=http://127.0.0.1:8081/v1 +# POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_HEADERS={"Authorization":"Bearer replace-me"} +# POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL_SETTINGS={"dimensions":2560} +# A dedicated reranker LLM is optional; without it, reranking reuses the generation model. +# POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL=openai-chat:local-reranker +# POWERCONTEXT_SERVER_INFERENCE_RERANK_BASE_URL=http://127.0.0.1:8082/v1 +# POWERCONTEXT_SERVER_INFERENCE_RERANK_HEADERS={"Authorization":"Bearer replace-me"} +# POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL_SETTINGS={"max_tokens":256} +# POWERCONTEXT_SERVER_INFERENCE_RERANK_TIMEOUT_SECONDS=30 +# POWERCONTEXT_SERVER_INFERENCE_RERANK_MAX_REQUESTS=2 + # Provider A: OpenAI (enabled). Set OPENAI_API_KEY in the Server shell. POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL=openai:gpt-4.1-mini POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL=openai:text-embedding-3-small diff --git a/docs/en/development/pydantic-ai-inference.md b/docs/en/development/pydantic-ai-inference.md index 1d8a9ec1f..35cd51a4f 100644 --- a/docs/en/development/pydantic-ai-inference.md +++ b/docs/en/development/pydantic-ai-inference.md @@ -46,8 +46,49 @@ export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_PROFILE_ID="project-embedding-v1" export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_DIMENSION="1536" ``` -Provider credentials remain in the environment variables understood by the selected Pydantic AI provider. They are -not fields on PowerContext models. +Each workload can target a different model service. Custom base URLs use the provider interface named by the model +identifier; use `openai-chat:` for an OpenAI-compatible Chat Completions service, `openai:` for an +OpenAI-compatible Responses or embeddings service, and `anthropic:` for an Anthropic-compatible generation +service. The built-in reranker is an LLM listwise reranker, so its independent endpoint is also a Pydantic AI generation +endpoint rather than a cross-encoder `/rerank` API: + +```bash +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL="openai-chat:generator" +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_BASE_URL="http://127.0.0.1:8080/v1" +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_HEADERS='{"Authorization":"Bearer generation-secret"}' +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS='{"max_tokens":4096}' + +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL="openai:embedding" +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BASE_URL="http://127.0.0.1:8081/v1" +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_HEADERS='{"Authorization":"Bearer embedding-secret"}' +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL_SETTINGS='{"dimensions":1536}' + +export POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL="openai-chat:reranker" +export POWERCONTEXT_SERVER_INFERENCE_RERANK_BASE_URL="http://127.0.0.1:8082/v1" +export POWERCONTEXT_SERVER_INFERENCE_RERANK_HEADERS='{"Authorization":"Bearer rerank-secret"}' +export POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL_SETTINGS='{"max_tokens":256}' +``` + +The header and model-settings values are JSON objects. Header values are treated as secrets by settings models and are +installed as static headers on the workload's provider client. They are not included in Pydantic AI request settings. +Do not put `extra_headers` inside a model-settings object; use the dedicated headers variable so configuration and log +redaction remain effective. Pydantic AI passes the remaining model settings through to the selected provider. The +reranker always fixes `temperature` to zero. A custom embedding base URL currently requires the OpenAI-compatible +embeddings interface. + +A base URL may contain a gateway path prefix. The selected Pydantic AI provider still owns the operation suffix, such +as `/chat/completions`, `/responses`, or `/embeddings`; arbitrary operation-path rewriting is not supported. +Custom base URLs and static headers require an explicit OpenAI- or Anthropic-compatible model identifier so +PowerContext can construct the corresponding provider client. + +When `RERANK_MODEL` is unset, LLM reranking reuses the generation model and base URL. Reranker headers and model +settings can still extend or override the generation configuration. A header override creates a separate provider +client for the rerank workload while retaining the generation model identifier and base URL. A separate reranker base +URL requires an explicit reranker model. The reranker timeout and request limit inherit their generation counterparts +unless they are set explicitly. + +When no custom base URL or headers are needed, provider credentials remain in the environment variables understood by +the selected Pydantic AI provider. The Server rejects a partial embedding profile. `embedding_model`, `embedding_profile_id`, and `embedding_dimension` must be configured together. SQLite vector search uses that embedding configuration because the index dimension and diff --git a/docs/en/docs/reference/configuration.md b/docs/en/docs/reference/configuration.md index c0ee374e0..3205f4f7b 100644 --- a/docs/en/docs/reference/configuration.md +++ b/docs/en/docs/reference/configuration.md @@ -74,15 +74,26 @@ Server settings use the `POWERCONTEXT_SERVER_` prefix. | `POWERCONTEXT_SERVER_RUNTIME_MEMORY_RERANK_CANDIDATE_LIMIT` | `30` | Coarse candidate pool supplied to the reranker | | `POWERCONTEXT_SERVER_RUNTIME_SCHEDULE_SECONDS` | unset | Scheduler interval; unset disables scheduling | | `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL` | unset | Pydantic AI model used by configured extraction, generation, Handoff, and reranking operations | -| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS` | `{}` | JSON object of Pydantic AI model settings shared by generation and reranking | +| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_BASE_URL` | provider default | Custom generation provider base URL | +| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_HEADERS` | `{}` | JSON object of static generation client headers; values are secrets | +| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS` | `{}` | JSON object of Pydantic AI generation model settings | | `POWERCONTEXT_SERVER_INFERENCE_GENERATION_TIMEOUT_SECONDS` | `30` | Timeout in seconds for one structured generation operation | | `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MAX_REQUESTS` | `2` | Maximum provider requests for one structured generation operation, including retries | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL` | unset | Pydantic AI embedding model; requires profile ID and dimension | +| `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BASE_URL` | provider default | Custom OpenAI-compatible embeddings base URL | +| `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_HEADERS` | `{}` | JSON object of static embedding client headers; values are secrets | +| `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL_SETTINGS` | `{}` | JSON object of Pydantic AI embedding model settings | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_PROFILE_ID` | unset | Stable identity for the model, dimension, and normalization used by the vector index | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_DIMENSION` | unset | Positive output dimension requested from and validated against the embedding model | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_NORMALIZATION` | `unit` | Vector normalization: `unit` or `none` | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_TIMEOUT_SECONDS` | `30` | Timeout in seconds for one embedding request | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BATCH_SIZE` | `10` | Maximum texts sent in one embedding request | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL` | generation model | Optional dedicated Pydantic AI model for LLM reranking | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_BASE_URL` | inherited/provider default | Custom LLM reranker provider base URL | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_HEADERS` | `{}` | JSON object of static LLM reranker client headers; values are secrets | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL_SETTINGS` | `{}` | JSON object of Pydantic AI reranker model settings | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_TIMEOUT_SECONDS` | generation timeout | LLM reranker timeout | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_MAX_REQUESTS` | generation request limit | Maximum model requests in one rerank operation | | `POWERCONTEXT_SERVER_RUNTIME_EXPERIENCE_SCHEDULE_SECONDS` | unset | Experience incubation interval; unset disables that job | | `POWERCONTEXT_SERVER_EXTERNAL_SKILLS` | unset | JSON object containing the host identity and explicit Agent Skill targets | @@ -128,19 +139,6 @@ command-line arguments, documentation, or Memory. Replace `provider:model-name` Pydantic AI. Scheduled extraction requires both a generation model and `POWERCONTEXT_SERVER_RUNTIME_SCHEDULE_SECONDS`. An explicit Memory write does not require either. -For provider-specific request parameters, set model settings as one JSON object. For example, an OpenAI-compatible -endpoint that supports Qwen's thinking switch can receive `chat_template_kwargs.enable_thinking=false` through the -portable Pydantic AI `extra_body` setting: - -```bash -export POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS='{"extra_body":{"chat_template_kwargs":{"enable_thinking":false}}}' -``` - -The Server applies these settings to extraction, Experience and Skill generation, Handoff generation, optional LLM -reranking, and the generation readiness probe. The readiness probe always overrides `max_tokens` to `1`, and reranking -always overrides `temperature` to `0`. Only settings supported by the selected Pydantic AI model and provider are -meaningful. Keep credentials and static headers in the selected provider's configuration rather than this JSON object. - The default `coding` extraction profile keeps cross-task work context such as preferences, decisions, constraints, expensive facts, and unfinished progress. Select `conversation` when the product must preserve independently answerable personal facts, relationships, events, exact dates, lists, and historical states from dialogue evidence: @@ -166,6 +164,10 @@ change stored Memory or indexes. Provider and structured-output failures remain reranking when search must remain independent of model availability. See [RFC 0080](/en/rfcs/0080_memory_search_reranking/) for the algorithm, concurrency, and API boundaries. +The built-in reranker is an LLM listwise reranker, not a dedicated cross-encoder protocol. By default it reuses the +generation model and its provider settings. Set `POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL` to give that LLM operation +an independent model, base URL, headers, settings, timeout, and request limit. + The same configured generation model gates explicit Experience generation, managed Skill generation and evolution, and external Skill import or fork. Without it, these operations return a capability error before persisting a Candidate. Candidate Review, exact reads, and external Skill scan/list/resolve continue to work. diff --git a/docs/zh/development/pydantic-ai-inference.md b/docs/zh/development/pydantic-ai-inference.md index 18097f660..62491060f 100644 --- a/docs/zh/development/pydantic-ai-inference.md +++ b/docs/zh/development/pydantic-ai-inference.md @@ -45,7 +45,46 @@ export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_PROFILE_ID="project-embedding-v1" export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_DIMENSION="1536" ``` -provider credential 仍使用所选 Pydantic AI provider 支持的环境变量，不属于 PowerContext model 字段。 +每类 workload 可以连接不同的模型服务。自定义 base URL 使用 model identifier 指定的 provider 接口： +OpenAI-compatible Chat Completions 服务使用 `openai-chat:`，OpenAI-compatible Responses 或 embedding +服务使用 `openai:`，Anthropic-compatible generation 服务使用 `anthropic:`。内置 reranker 是 LLM +listwise reranker，因此它的独立 endpoint 也是 Pydantic AI generation endpoint，而不是 cross-encoder `/rerank` +API： + +```bash +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL="openai-chat:generator" +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_BASE_URL="http://127.0.0.1:8080/v1" +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_HEADERS='{"Authorization":"Bearer generation-secret"}' +export POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS='{"max_tokens":4096}' + +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL="openai:embedding" +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BASE_URL="http://127.0.0.1:8081/v1" +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_HEADERS='{"Authorization":"Bearer embedding-secret"}' +export POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL_SETTINGS='{"dimensions":1536}' + +export POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL="openai-chat:reranker" +export POWERCONTEXT_SERVER_INFERENCE_RERANK_BASE_URL="http://127.0.0.1:8082/v1" +export POWERCONTEXT_SERVER_INFERENCE_RERANK_HEADERS='{"Authorization":"Bearer rerank-secret"}' +export POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL_SETTINGS='{"max_tokens":256}' +``` + +header 和 model settings 都使用 JSON object。settings model 会将 header value 作为 secret 处理，并将它们作为 +workload provider client 的静态 header；这些值不会进入 Pydantic AI request settings。不要在 model settings 中配置 +`extra_headers`；使用独立的 headers 变量才能保留配置与日志脱敏语义。其余 model settings 由 Pydantic AI 传递给 +所选 provider。reranker 始终将 `temperature` 固定为零。自定义 embedding base URL 目前要求服务实现 +OpenAI-compatible embeddings 接口。 + +base URL 可以包含 gateway path prefix，但具体 operation suffix 仍由所选 Pydantic AI provider 决定，例如 +`/chat/completions`、`/responses` 或 `/embeddings`；不支持任意改写 operation path。 +自定义 base URL 或静态 header 时必须使用显式的 OpenAI- 或 Anthropic-compatible model identifier，PowerContext +才能创建对应的 provider client。 + +未设置 `RERANK_MODEL` 时，LLM rerank 复用 generation model 和 base URL；仍可通过 reranker headers 和 model +settings 扩展或覆盖 generation 配置。覆盖 header 时会为 rerank workload 创建独立 provider client，但保留 +generation model identifier 和 base URL。独立的 reranker base URL 必须同时配置显式 reranker model。reranker +timeout 和 request limit 未显式设置时继承 generation 的对应配置。 + +不需要自定义 base URL 或 header 时，provider credential 仍使用所选 Pydantic AI provider 支持的环境变量。 Server 会拒绝不完整的 embedding profile。`embedding_model`、`embedding_profile_id` 和 `embedding_dimension` 必须一起配置。SQLite vector search 使用这组配置，因为 index dimension 必须与持久化向量一致。 diff --git a/docs/zh/docs/reference/configuration.md b/docs/zh/docs/reference/configuration.md index e7c6d3c4c..4c42b6bf7 100644 --- a/docs/zh/docs/reference/configuration.md +++ b/docs/zh/docs/reference/configuration.md @@ -71,15 +71,26 @@ Server 配置使用 `POWERCONTEXT_SERVER_` 前缀。 | `POWERCONTEXT_SERVER_RUNTIME_MEMORY_RERANK_CANDIDATE_LIMIT` | `30` | 交给 reranker 的粗排候选池大小 | | `POWERCONTEXT_SERVER_RUNTIME_SCHEDULE_SECONDS` | 未设置 | Scheduler 间隔；未设置即不启用 | | `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL` | 未设置 | 配置的 extraction、generation、Handoff 和 rerank 操作共用的 Pydantic AI 模型 | -| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS` | `{}` | generation 与 rerank 共用的 Pydantic AI model settings JSON object | +| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_BASE_URL` | provider 默认值 | 自定义 generation provider base URL | +| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_HEADERS` | `{}` | generation client 静态 header JSON object；value 按 secret 处理 | +| `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS` | `{}` | Pydantic AI generation model settings JSON object | | `POWERCONTEXT_SERVER_INFERENCE_GENERATION_TIMEOUT_SECONDS` | `30` | 单次结构化 generation 操作的超时秒数 | | `POWERCONTEXT_SERVER_INFERENCE_GENERATION_MAX_REQUESTS` | `2` | 单次结构化 generation 操作最多发起的 provider 请求数，包含重试 | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL` | 未设置 | Pydantic AI embedding model；必须同时设置 profile ID 和 dimension | +| `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BASE_URL` | provider 默认值 | 自定义 OpenAI-compatible embeddings base URL | +| `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_HEADERS` | `{}` | embedding client 静态 header JSON object；value 按 secret 处理 | +| `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL_SETTINGS` | `{}` | Pydantic AI embedding model settings JSON object | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_PROFILE_ID` | 未设置 | vector index 使用的模型、dimension 和 normalization 的稳定标识 | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_DIMENSION` | 未设置 | 向 embedding model 请求并校验的正整数输出维度 | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_NORMALIZATION` | `unit` | vector normalization：`unit` 或 `none` | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_TIMEOUT_SECONDS` | `30` | 单次 embedding 请求的超时秒数 | | `POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BATCH_SIZE` | `10` | 单次 embedding 请求最多发送的文本数量 | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL` | generation model | LLM rerank 可选的独立 Pydantic AI model | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_BASE_URL` | 继承值或 provider 默认值 | 自定义 LLM reranker provider base URL | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_HEADERS` | `{}` | LLM reranker client 静态 header JSON object；value 按 secret 处理 | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL_SETTINGS` | `{}` | Pydantic AI reranker model settings JSON object | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_TIMEOUT_SECONDS` | generation 超时 | LLM reranker 超时 | +| `POWERCONTEXT_SERVER_INFERENCE_RERANK_MAX_REQUESTS` | generation request limit | 单次 rerank operation 的最大 model request 数量 | | `POWERCONTEXT_SERVER_RUNTIME_EXPERIENCE_SCHEDULE_SECONDS` | 未设置 | Experience 孵化间隔；未设置即不启用该 job | | `POWERCONTEXT_SERVER_EXTERNAL_SKILLS` | 未设置 | 包含 host identity 和显式 Agent Skill targets 的 JSON object | @@ -120,19 +131,6 @@ powercontext server run Memory。请把 `provider:model-name` 替换为 Pydantic AI 支持的模型标识。定时提取需要同时配置 generation model 和 `POWERCONTEXT_SERVER_RUNTIME_SCHEDULE_SECONDS`；显式 Memory 写入不需要这两项配置。 -provider-specific request parameter 使用一个 JSON object 配置。例如，OpenAI-compatible endpoint 支持 Qwen 的 -thinking switch 时，可以通过 Pydantic AI 的通用 `extra_body` setting 发送 -`chat_template_kwargs.enable_thinking=false`： - -```bash -export POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS='{"extra_body":{"chat_template_kwargs":{"enable_thinking":false}}}' -``` - -Server 会将这些 settings 用于 extraction、Experience 与 Skill generation、Handoff generation、可选的 LLM -rerank 以及 generation readiness probe。readiness probe 始终将 `max_tokens` 覆盖为 `1`，rerank 始终将 -`temperature` 覆盖为 `0`。只有所选 Pydantic AI model 与 provider 支持的 setting 才有意义。credential 和 -static header 应保留在所选 provider 的配置中，不要放入这个 JSON object。 - 默认的 `coding` 抽取 profile 保留跨任务工作上下文，例如偏好、决策、约束、昂贵事实和未完成进度。当产品 需要从对话证据中保留可独立回答的人物事实、关系、事件、精确日期、列表和历史状态时，可选择 `conversation`： @@ -157,6 +155,10 @@ search request 最终 `limit` 的结果。它不会修改已存储 Memory 或索 显式返回；如果搜索必须独立于模型可用性，请关闭 rerank。算法、并发与 API 边界见 [RFC 0080](/zh/rfcs/0080_memory_search_reranking/)。 +内置 reranker 是 LLM listwise reranker，不是独立的 cross-encoder protocol。默认复用 generation model 及其 provider +settings。设置 `POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL` 后，该 LLM operation 可以使用独立的 model、base URL、 +headers、settings、timeout 和 request limit。 + 同一个 generation model 也控制显式 Experience generation、managed Skill generation/evolution，以及 external Skill import/fork。未配置模型时，这些 operation 会在持久化 Candidate 前返回 capability error； Candidate Review、exact read 和 external Skill scan/list/resolve 仍可使用。 diff --git a/src/powercontext/builtin/runtime/composition.py b/src/powercontext/builtin/runtime/composition.py index 3867a28e9..88271c23c 100644 --- a/src/powercontext/builtin/runtime/composition.py +++ b/src/powercontext/builtin/runtime/composition.py @@ -16,12 +16,13 @@ from __future__ import annotations -from collections.abc import AsyncIterator +import os +from collections.abc import AsyncIterator, Callable, Mapping from contextlib import AsyncExitStack, asynccontextmanager from pathlib import Path -from typing import TYPE_CHECKING, TypeVar, cast +from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast -from pydantic import JsonValue +from pydantic import AnyHttpUrl, JsonValue, SecretStr from typing_extensions import override from powercontext.builtin.artifacts.experience import ExperienceCandidatePipeline, ExperienceGenerator @@ -76,7 +77,9 @@ from powercontext.sources import Source if TYPE_CHECKING: + from pydantic_ai.models import Model from pydantic_ai.models.instrumented import InstrumentationSettings + from pydantic_ai.providers import Provider ValueT = TypeVar("ValueT") @@ -87,8 +90,11 @@ class BuiltinConfigurationError(RuntimeError): def __init__(self, issue: str) -> None: messages = { "external-skill-host": "external Skill roots require a host identity", + "inference-endpoint-provider": ( + "custom inference endpoints require an OpenAI- or Anthropic-compatible model identifier" + ), "inference-profile": "validated inference profile is incomplete", - "memory-reranker": "Memory reranking requires a configured generation model or injected reranker", + "memory-reranker": "Memory reranking requires a configured generation or rerank model, or injected reranker", "scheduled-experience-pipeline": "scheduled Experience incubation requires a candidate pipeline", "scheduled-pipeline": "scheduled Source processing requires a candidate pipeline", "database": "unsupported built-in database", @@ -182,6 +188,7 @@ async def open_builtin_runtime( generated_handoff, generated_reranker, generation_readiness, + rerank_readiness, ) = ( await _generation_pipelines(config.inference, config.runtime, resources, instrumentation) if ( @@ -192,7 +199,7 @@ async def open_builtin_runtime( or handoff_pipeline is None or (config.runtime.memory_rerank_enabled and memory_reranker is None) ) - else (None, None, None, None, None, None, None) + else (None, None, None, None, None, None, None, None) ) configured_pipeline = generated_memory if candidate_pipeline is None else candidate_pipeline configured_incubation = generated_incubation if experience_pipeline is None else experience_pipeline @@ -245,16 +252,17 @@ async def open_builtin_runtime( blocking=True, ), } - if generation_readiness is not None: - readiness_probes["inference.generation"] = ReadinessProbeDefinition( - probe=generation_readiness, - blocking=False, - ) - if readiness_embedding is not None: - readiness_probes["inference.embedding"] = ReadinessProbeDefinition( - probe=_embedding_readiness_probe(readiness_embedding), - blocking=False, - ) + inference_readiness = ( + ("inference.generation", generation_readiness), + ("inference.rerank", rerank_readiness), + ( + "inference.embedding", + None if readiness_embedding is None else _embedding_readiness_probe(readiness_embedding), + ), + ) + for name, readiness_probe in inference_readiness: + if readiness_probe is not None: + readiness_probes[name] = ReadinessProbeDefinition(probe=readiness_probe, blocking=False) runtime = await resources.enter_async_context( BuiltinRuntime( provider=contexts, @@ -400,12 +408,11 @@ async def _generation_pipelines( HandoffGenerationPipeline | None, MemoryReranker | None, ReadinessProbe | None, + ReadinessProbe | None, ]: - if settings.generation_model is None: - return None, None, None, None, None, None, None + if settings.generation_model is None and (not runtime.memory_rerank_enabled or settings.rerank_model is None): + return None, None, None, None, None, None, None, None - from pydantic_ai.models import infer_model - from pydantic_ai.models.instrumented import InstrumentedModel from pydantic_ai.settings import ModelSettings, merge_model_settings from powercontext.builtin.artifacts.experience import ( @@ -445,95 +452,273 @@ async def _generation_pipelines( probe_pydantic_ai_model, ) - provider_model = await resources.enter_async_context(infer_model(settings.generation_model)) - model = provider_model if instrumentation is None else InstrumentedModel(provider_model, instrumentation) - model_settings = cast(ModelSettings, dict(settings.generation_model_settings)) or None - - async def probe_generation() -> None: - # Readiness probing runs outside any operation span; keep it out of traces. - await probe_pydantic_ai_model( - provider_model, - timeout_seconds=READINESS_PROBE_TIMEOUT_SECONDS, - model_settings=model_settings, + generated_memory: CandidatePipeline | None = None + generated_incubation: ExperienceCandidatePipeline | None = None + generated_experience: ExperienceGenerator | None = None + generated_skill: SkillGenerator | None = None + generated_handoff: HandoffGenerationPipeline | None = None + generated_reranker: MemoryReranker | None = None + generation_readiness: ReadinessProbe | None = None + rerank_readiness: ReadinessProbe | None = None + + generation_provider_model: Model | None = None + generation_model: Model | None = None + generation_request_settings: ModelSettings | None = None + if settings.generation_model is not None: + generation_provider_model, generation_model = await _open_pydantic_ai_model( + settings.generation_model, + base_url=settings.generation_base_url, + headers=settings.generation_headers, + resources=resources, + instrumentation=instrumentation, ) - - limits = InferenceLimits( - timeout_seconds=settings.generation_timeout_seconds, - max_requests=settings.generation_max_requests, - ) - memory_generator = PydanticAIStructuredGenerator( - model=model, - instructions=memory_extraction_instructions(runtime.memory_extraction_profile), - input_type=MemoryExtractionInput, - output_type=MemoryExtractionOutput, - limits=limits, - model_settings=model_settings, - name="memory_extraction", - ) - experience_generator = PydanticAIStructuredGenerator( - model=model, - instructions=EXPERIENCE_INCUBATION_INSTRUCTIONS, - input_type=ExperienceIncubationInput, - output_type=ExperienceIncubationOutput, - limits=limits, - model_settings=model_settings, - name="experience_incubation", - ) - explicit_experience_generator = PydanticAIStructuredGenerator( - model=model, - instructions=EXPERIENCE_GENERATION_INSTRUCTIONS, - input_type=ArtifactGenerationInput, - output_type=ExperienceGenerationOutput, - limits=limits, - model_settings=model_settings, - name="experience_generation", - ) - skill_generator = PydanticAIStructuredGenerator( - model=model, - instructions=SKILL_GENERATION_INSTRUCTIONS, - input_type=ArtifactGenerationInput, - output_type=SkillGenerationOutput, - limits=limits, - model_settings=model_settings, - name="skill_generation", - ) - handoff_generator = PydanticAIStructuredGenerator( - model=model, - instructions=HANDOFF_GENERATION_INSTRUCTIONS, - input_type=HandoffGenerationInput, - output_type=HandoffGenerationOutput, - limits=limits, - model_settings=model_settings, - name="handoff_generation", - ) - rerank_generator = ( - PydanticAIStructuredGenerator( - model=model, - instructions=MEMORY_RERANK_INSTRUCTIONS, - input_type=MemoryRerankInput, - output_type=MemoryRerankOutput, - limits=limits, - model_settings=merge_model_settings(model_settings, ModelSettings(temperature=0.0)), - name="memory_rerank", + generation_request_settings = cast(ModelSettings, dict(settings.generation_model_settings)) + generation_limits = InferenceLimits( + timeout_seconds=settings.generation_timeout_seconds, + max_requests=settings.generation_max_requests, ) - if runtime.memory_rerank_enabled - else None - ) - return ( - LLMMemoryCandidatePipeline( + memory_generator = PydanticAIStructuredGenerator( + model=generation_model, + instructions=memory_extraction_instructions(runtime.memory_extraction_profile), + input_type=MemoryExtractionInput, + output_type=MemoryExtractionOutput, + limits=generation_limits, + model_settings=generation_request_settings, + name="memory_extraction", + ) + experience_generator = PydanticAIStructuredGenerator( + model=generation_model, + instructions=EXPERIENCE_INCUBATION_INSTRUCTIONS, + input_type=ExperienceIncubationInput, + output_type=ExperienceIncubationOutput, + limits=generation_limits, + model_settings=generation_request_settings, + name="experience_incubation", + ) + explicit_experience_generator = PydanticAIStructuredGenerator( + model=generation_model, + instructions=EXPERIENCE_GENERATION_INSTRUCTIONS, + input_type=ArtifactGenerationInput, + output_type=ExperienceGenerationOutput, + limits=generation_limits, + model_settings=generation_request_settings, + name="experience_generation", + ) + skill_generator = PydanticAIStructuredGenerator( + model=generation_model, + instructions=SKILL_GENERATION_INSTRUCTIONS, + input_type=ArtifactGenerationInput, + output_type=SkillGenerationOutput, + limits=generation_limits, + model_settings=generation_request_settings, + name="skill_generation", + ) + handoff_generator = PydanticAIStructuredGenerator( + model=generation_model, + instructions=HANDOFF_GENERATION_INSTRUCTIONS, + input_type=HandoffGenerationInput, + output_type=HandoffGenerationOutput, + limits=generation_limits, + model_settings=generation_request_settings, + name="handoff_generation", + ) + generated_memory = LLMMemoryCandidatePipeline( UsageReportingStructuredGenerator(memory_generator), evidence_projector=_ContentEvidenceProjector(), - ), - LLMExperienceCandidatePipeline(UsageReportingStructuredGenerator(experience_generator)), - LLMExperienceGenerator(UsageReportingStructuredGenerator(explicit_experience_generator)), - LLMSkillGenerator(UsageReportingStructuredGenerator(skill_generator)), - LLMHandoffGenerationPipeline( + ) + generated_incubation = LLMExperienceCandidatePipeline(UsageReportingStructuredGenerator(experience_generator)) + generated_experience = LLMExperienceGenerator(UsageReportingStructuredGenerator(explicit_experience_generator)) + generated_skill = LLMSkillGenerator(UsageReportingStructuredGenerator(skill_generator)) + generated_handoff = LLMHandoffGenerationPipeline( UsageReportingStructuredGenerator(handoff_generator), evidence_projector=_ContentHandoffEvidenceProjector(), - ), - (None if rerank_generator is None else LLMMemoryReranker(UsageReportingStructuredGenerator(rerank_generator))), - CachedReadinessProbe(dependency_readiness_probe(probe_generation)), + ) + + async def probe_generation() -> None: + # Readiness probing runs outside any operation span; keep it out of traces. + await probe_pydantic_ai_model( + generation_provider_model, + timeout_seconds=READINESS_PROBE_TIMEOUT_SECONDS, + model_settings=generation_request_settings, + ) + + generation_readiness = CachedReadinessProbe(dependency_readiness_probe(probe_generation)) + + if runtime.memory_rerank_enabled: + rerank_provider_model = generation_provider_model + rerank_model = generation_model + inherits_generation = settings.rerank_model is None + rerank_headers = ( + _merge_headers(settings.generation_headers, settings.rerank_headers) + if inherits_generation + else settings.rerank_headers + ) + separate_rerank_model = settings.rerank_model is not None or bool(settings.rerank_headers) + if separate_rerank_model: + rerank_model_name = settings.rerank_model or settings.generation_model + if rerank_model_name is None: + raise BuiltinConfigurationError("memory-reranker") + rerank_provider_model, rerank_model = await _open_pydantic_ai_model( + rerank_model_name, + base_url=settings.rerank_base_url + if settings.rerank_model is not None + else settings.generation_base_url, + headers=rerank_headers, + resources=resources, + instrumentation=instrumentation, + ) + if rerank_provider_model is not None and rerank_model is not None: + rerank_values = ( + settings.generation_model_settings | settings.rerank_model_settings + if inherits_generation + else settings.rerank_model_settings + ) + rerank_request_settings = cast(ModelSettings, dict(rerank_values)) + rerank_request_settings = merge_model_settings( + rerank_request_settings, + ModelSettings(temperature=0.0), + ) + rerank_generator = PydanticAIStructuredGenerator( + model=rerank_model, + instructions=MEMORY_RERANK_INSTRUCTIONS, + input_type=MemoryRerankInput, + output_type=MemoryRerankOutput, + limits=InferenceLimits( + timeout_seconds=settings.rerank_timeout_seconds or settings.generation_timeout_seconds, + max_requests=settings.rerank_max_requests or settings.generation_max_requests, + ), + model_settings=rerank_request_settings, + name="memory_rerank", + ) + generated_reranker = LLMMemoryReranker(UsageReportingStructuredGenerator(rerank_generator)) + + if separate_rerank_model or settings.rerank_model_settings: + + async def probe_rerank() -> None: + await probe_pydantic_ai_model( + rerank_provider_model, + timeout_seconds=READINESS_PROBE_TIMEOUT_SECONDS, + model_settings=rerank_request_settings, + ) + + rerank_readiness = CachedReadinessProbe(dependency_readiness_probe(probe_rerank)) + + return ( + generated_memory, + generated_incubation, + generated_experience, + generated_skill, + generated_handoff, + generated_reranker, + generation_readiness, + rerank_readiness, + ) + + +async def _open_pydantic_ai_model( + model_name: str, + *, + base_url: AnyHttpUrl | None, + headers: Mapping[str, SecretStr], + resources: AsyncExitStack, + instrumentation: InstrumentationSettings | None, +) -> tuple[Model, Model]: + from pydantic_ai.models import infer_model + from pydantic_ai.models.instrumented import InstrumentedModel + + if (base_url is not None or headers) and ":" not in model_name: + raise BuiltinConfigurationError("inference-endpoint-provider") + inferred_model = ( + infer_model(model_name) + if base_url is None and not headers + else infer_model( + model_name, + provider_factory=_provider_factory( + base_url, + headers, + workload="generation", + resources=resources, + ), + ) ) + provider_model = await resources.enter_async_context(inferred_model) + model = provider_model if instrumentation is None else InstrumentedModel(provider_model, instrumentation) + return provider_model, model + + +def _provider_factory( + base_url: AnyHttpUrl | None, + headers: Mapping[str, SecretStr], + *, + workload: Literal["generation", "embedding"], + resources: AsyncExitStack, +) -> Callable[[str], Provider[Any]]: + from pydantic_ai.providers import infer_provider + + if base_url is None and not headers: + return infer_provider + + from pydantic_ai.providers.openai import OpenAIProvider + + def create_provider(provider_name: str) -> Provider[Any]: + if provider_name in {"openai", "openai-chat", "openai-responses"}: + if headers: + from openai import AsyncOpenAI + + client = AsyncOpenAI( + base_url=None if base_url is None else str(base_url), + api_key=os.getenv("OPENAI_API_KEY") or "api-key-not-set", + default_headers=_resolve_headers(headers), + ) + resources.push_async_callback(client.close) + return OpenAIProvider(openai_client=client) + return OpenAIProvider(base_url=str(base_url)) + if provider_name == "anthropic" and workload == "generation": + from anthropic import AsyncAnthropic + from pydantic_ai.providers.anthropic import AnthropicProvider + + if headers: + default_headers = _resolve_headers(headers) + api_key = _pop_header(default_headers, "x-api-key") + client = AsyncAnthropic( + base_url=None if base_url is None else str(base_url), + api_key=api_key or os.getenv("ANTHROPIC_API_KEY") or "api-key-not-set", + default_headers=default_headers, + ) + resources.push_async_callback(client.close) + return AnthropicProvider(anthropic_client=client) + return AnthropicProvider( + base_url=str(base_url), + api_key=os.getenv("ANTHROPIC_API_KEY") or "api-key-not-set", + ) + raise BuiltinConfigurationError("inference-endpoint-provider") + + return create_provider + + +def _resolve_headers(headers: Mapping[str, SecretStr]) -> dict[str, str]: + return {name: value.get_secret_value() for name, value in headers.items()} + + +def _pop_header(headers: dict[str, str], name: str) -> str | None: + expected = name.casefold() + for existing_name in headers: + if existing_name.casefold() == expected: + return headers.pop(existing_name) + return None + + +def _merge_headers(*values: Mapping[str, SecretStr]) -> dict[str, SecretStr]: + merged: dict[str, SecretStr] = {} + names: dict[str, str] = {} + for headers in values: + for name, value in headers.items(): + normalized_name = name.casefold() + if previous_name := names.get(normalized_name): + del merged[previous_name] + merged[name] = value + names[normalized_name] = name + return merged async def _embedding_models( @@ -545,16 +730,21 @@ async def _embedding_models( return None, None from pydantic_ai import Embedder - from pydantic_ai.embeddings import infer_embedding_model - from pydantic_ai.providers import Provider, infer_provider + from pydantic_ai.embeddings import EmbeddingSettings, infer_embedding_model from powercontext.builtin.artifacts.memory import EmbeddingProfile from powercontext.builtin.inference.pydantic_ai import InferenceLimits, PydanticAIEmbeddingModel - providers: list[Provider[object]] = [] + providers: list[Provider[Any]] = [] + create_provider = _provider_factory( + settings.embedding_base_url, + settings.embedding_headers, + workload="embedding", + resources=resources, + ) - def provider_factory(provider_name: str) -> Provider[object]: - provider = infer_provider(provider_name) + def provider_factory(provider_name: str) -> Provider[Any]: + provider = create_provider(provider_name) providers.append(provider) return provider @@ -569,14 +759,14 @@ def provider_factory(provider_name: str) -> Provider[object]: normalization=settings.embedding_normalization, ) limits = InferenceLimits(timeout_seconds=settings.embedding_timeout_seconds) + embedding_settings = cast( + EmbeddingSettings, + settings.embedding_model_settings | {"dimensions": profile.dimension}, + ) def adapter(instrument: InstrumentationSettings | bool | None) -> EmbeddingModel: return PydanticAIEmbeddingModel( - embedder=Embedder( - model, - settings={"dimensions": profile.dimension}, - instrument=instrument, - ), + embedder=Embedder(model, settings=embedding_settings, instrument=instrument), batch_size=settings.embedding_batch_size, profile=profile, limits=limits, diff --git a/src/powercontext/builtin/runtime/config.py b/src/powercontext/builtin/runtime/config.py index c41989927..d734c65fa 100644 --- a/src/powercontext/builtin/runtime/config.py +++ b/src/powercontext/builtin/runtime/config.py @@ -16,10 +16,11 @@ from __future__ import annotations +import re from collections.abc import Mapping from typing import Any, Literal, Self -from pydantic import BaseModel, Field, JsonValue, field_validator, model_validator +from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, JsonValue, SecretStr, field_validator, model_validator from powercontext.builtin.artifacts.memory.prompts import MemoryExtractionProfile from powercontext.builtin.artifacts.skill import AgentSkillTarget, CodexSkillRoot @@ -28,6 +29,8 @@ from powercontext.builtin.persistence.sqlite import SQLiteConfig from powercontext.builtin.runtime._scope_cache import DEFAULT_SCOPE_CACHE_SIZE +_HTTP_FIELD_NAME_PATTERN = re.compile(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+") + class RuntimeConfig(BaseModel): """Built-in runtime policy and scheduler configuration.""" @@ -48,20 +51,33 @@ class HandoffReportConfig(BaseModel): class InferenceConfig(BaseModel): - """Optional generation and embedding configuration.""" + """Optional generation, embedding, and LLM reranking configuration.""" + + model_config = ConfigDict(hide_input_in_errors=True) generation_model: str | None = None + generation_base_url: AnyHttpUrl | None = None + generation_headers: dict[str, SecretStr] = Field(default_factory=dict, repr=False) generation_model_settings: dict[str, JsonValue] = Field(default_factory=dict) generation_timeout_seconds: float = Field(default=30.0, gt=0) generation_max_requests: int = Field(default=2, ge=1) embedding_model: str | None = None + embedding_base_url: AnyHttpUrl | None = None + embedding_headers: dict[str, SecretStr] = Field(default_factory=dict, repr=False) + embedding_model_settings: dict[str, JsonValue] = Field(default_factory=dict) embedding_profile_id: str | None = None embedding_dimension: int | None = Field(default=None, ge=1) embedding_normalization: Literal["none", "unit"] = "unit" embedding_timeout_seconds: float = Field(default=30.0, gt=0) embedding_batch_size: int = Field(default=10, ge=1) - - @field_validator("generation_model", "embedding_model", "embedding_profile_id") + rerank_model: str | None = None + rerank_base_url: AnyHttpUrl | None = None + rerank_headers: dict[str, SecretStr] = Field(default_factory=dict, repr=False) + rerank_model_settings: dict[str, JsonValue] = Field(default_factory=dict) + rerank_timeout_seconds: float | None = Field(default=None, gt=0) + rerank_max_requests: int | None = Field(default=None, ge=1) + + @field_validator("generation_model", "embedding_model", "embedding_profile_id", "rerank_model") @classmethod def validate_optional_identifier(cls, value: str | None) -> str | None: if value is None: @@ -81,18 +97,29 @@ def validate_normalization(cls, value: object) -> object: raise ValueError("embedding normalization must be 'none' or 'unit'") # noqa: TRY003 return normalized - @field_validator("generation_model_settings") + @field_validator("generation_headers", "embedding_headers", "rerank_headers") @classmethod - def reserve_generation_headers(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: - if "extra_headers" in value: - raise ValueError("configure credentials and static headers through the selected inference provider") # noqa: TRY003 + def validate_headers(cls, value: dict[str, SecretStr]) -> dict[str, SecretStr]: + normalized_names: set[str] = set() + for name, secret in value.items(): + normalized_name = name.casefold() + if _HTTP_FIELD_NAME_PATTERN.fullmatch(name) is None: + raise ValueError("inference header names must be non-empty HTTP field names") # noqa: TRY003 + if normalized_name in normalized_names: + raise ValueError("inference header names must be unique ignoring case") # noqa: TRY003 + if not secret.get_secret_value(): + raise ValueError("inference header values must not be empty") # noqa: TRY003 + normalized_names.add(normalized_name) return value - @model_validator(mode="after") - def validate_generation_model_settings(self) -> Self: - if self.generation_model_settings and self.generation_model is None: - raise ValueError("generation_model_settings requires generation_model") # noqa: TRY003 - return self + @field_validator("generation_model_settings", "embedding_model_settings", "rerank_model_settings") + @classmethod + def reserve_headers_field(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: + if "extra_headers" in value: + raise ValueError( # noqa: TRY003 + "configure credentials and static headers through the dedicated headers field" + ) + return value @model_validator(mode="after") def validate_embedding_profile(self) -> Self: @@ -103,6 +130,26 @@ def validate_embedding_profile(self) -> Self: ) return self + @model_validator(mode="after") + def validate_workload_overrides(self) -> Self: + if self.generation_model is None and self.generation_model_settings: + raise ValueError("generation_model_settings requires generation_model") # noqa: TRY003 + if self.generation_model is None and (self.generation_base_url is not None or self.generation_headers): + raise ValueError("generation overrides require generation_model") # noqa: TRY003 + if self.embedding_model is None and ( + self.embedding_base_url is not None or self.embedding_headers or self.embedding_model_settings + ): + raise ValueError("embedding overrides require a complete embedding profile") # noqa: TRY003 + if self.rerank_base_url is not None and self.rerank_model is None: + raise ValueError("rerank_base_url requires rerank_model") # noqa: TRY003 + if ( + self.rerank_model is None + and self.generation_model is None + and (self.rerank_headers or self.rerank_model_settings) + ): + raise ValueError("rerank overrides require rerank_model or generation_model") # noqa: TRY003 + return self + class ExternalSkillsConfig(BaseModel): """Explicit host-local targets used by Agent-native Skill providers.""" diff --git a/tests/builtin/runtime/test_composition_embedding.py b/tests/builtin/runtime/test_composition_embedding.py index f4ce4f916..b463495ec 100644 --- a/tests/builtin/runtime/test_composition_embedding.py +++ b/tests/builtin/runtime/test_composition_embedding.py @@ -25,7 +25,7 @@ class _SpyEmbedder(Embedder): - settings_seen: ClassVar[list[dict[str, int] | None]] = [] + settings_seen: ClassVar[list[dict[str, object] | None]] = [] def __init__(self, model, *, settings=None, defer_model_check=True, instrument=None): super().__init__(model, settings=settings, defer_model_check=defer_model_check, instrument=instrument) @@ -40,6 +40,7 @@ def test_embedding_models_send_the_configured_dimension_to_the_provider(monkeypa async def scenario() -> None: config = InferenceConfig( embedding_model="openai:text-embedding-3-small", + embedding_model_settings={"dimensions": 512, "extra_body": {"route": "embedding"}}, embedding_profile_id="bailian-1536-v1", embedding_dimension=1536, ) @@ -47,7 +48,8 @@ async def scenario() -> None: operational, readiness = await _embedding_models(config, resources, None) assert operational is not None assert readiness is not None - assert _SpyEmbedder.settings_seen == [{"dimensions": 1536}, {"dimensions": 1536}] + expected_settings = {"dimensions": 1536, "extra_body": {"route": "embedding"}} + assert _SpyEmbedder.settings_seen == [expected_settings, expected_settings] asyncio.run(scenario()) diff --git a/tests/test_inference_endpoints.py b/tests/test_inference_endpoints.py new file mode 100644 index 000000000..32f95de99 --- /dev/null +++ b/tests/test_inference_endpoints.py @@ -0,0 +1,311 @@ +# Copyright (c) 2026 OceanBase. +# +# 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 __future__ import annotations + +import asyncio +import json +import logging +import threading +from collections.abc import Iterator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +import pytest +from pydantic import AnyHttpUrl, SecretStr, ValidationError + +from powercontext.builtin.artifacts.memory import MemoryEntryInput +from powercontext.builtin.persistence.sqlite import SQLiteConfig +from powercontext.builtin.runtime import ( + BuiltinConfig, + InferenceConfig, + RememberMemoryRequest, + RuntimeConfig, + SearchMemoryRequest, + open_builtin_runtime, +) +from powercontext.server.settings import ServerSettings + + +class _RecordingModelServer(ThreadingHTTPServer): + requests: list[dict[str, Any]] + + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), _ModelHandler) + self.requests = [] + + +class _ModelHandler(BaseHTTPRequestHandler): + server: _RecordingModelServer + + def do_POST(self) -> None: + content_length = int(self.headers.get("Content-Length", "0")) + body = json.loads(self.rfile.read(content_length)) + path = self.path.partition("?")[0] + self.server.requests.append({ + "path": path, + "headers": {name.lower(): value for name, value in self.headers.items()}, + "body": body, + }) + if path == "/v1/embeddings": + response = { + "object": "list", + "model": body["model"], + "data": [ + {"object": "embedding", "index": index, "embedding": [1.0, 0.0, 0.0]} + for index, _value in enumerate(body["input"]) + ], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + elif path == "/v1/chat/completions": + response = { + "id": "chatcmpl-readiness", + "object": "chat.completion", + "created": 0, + "model": body["model"], + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": '{"selected_ranks":[1]}'}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + elif path == "/v1/messages": + response = { + "id": "msg-readiness", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "ok"}], + "model": body["model"], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + else: + self.send_error(404) + return + encoded = json.dumps(response).encode("utf-8") + 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, format: str, *args: Any) -> None: # noqa: A002 + return + + +@contextmanager +def _model_server() -> Iterator[tuple[_RecordingModelServer, str]]: + server = _RecordingModelServer() + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + address = server.server_address + host, port = str(address[0]), int(address[1]) + try: + yield server, f"http://{host}:{port}/v1" + finally: + server.shutdown() + server.server_close() + thread.join() + + +def test_inference_workload_settings_load_from_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL", "openai-chat:generator") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_GENERATION_BASE_URL", "http://generation.test/v1") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_GENERATION_HEADERS", '{"X-Workload":"generation-secret"}') + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_GENERATION_MODEL_SETTINGS", '{"max_tokens":256}') + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL", "openai:embedding") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_BASE_URL", "http://embedding.test/v1") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_HEADERS", '{"X-Workload":"embedding-secret"}') + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_MODEL_SETTINGS", '{"dimensions":3}') + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_PROFILE_ID", "embedding-v1") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_EMBEDDING_DIMENSION", "3") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL", "openai-chat:reranker") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_RERANK_BASE_URL", "http://rerank.test/v1") + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_RERANK_HEADERS", '{"X-Workload":"rerank-secret"}') + monkeypatch.setenv("POWERCONTEXT_SERVER_INFERENCE_RERANK_MODEL_SETTINGS", '{"top_p":0.25}') + + inference = ServerSettings().inference + + assert str(inference.generation_base_url) == "http://generation.test/v1" + assert inference.generation_headers["X-Workload"].get_secret_value() == "generation-secret" + assert inference.generation_model_settings == {"max_tokens": 256} + assert str(inference.embedding_base_url) == "http://embedding.test/v1" + assert inference.embedding_headers["X-Workload"].get_secret_value() == "embedding-secret" + assert inference.embedding_model_settings == {"dimensions": 3} + assert inference.rerank_model == "openai-chat:reranker" + assert str(inference.rerank_base_url) == "http://rerank.test/v1" + assert inference.rerank_headers["X-Workload"].get_secret_value() == "rerank-secret" + assert inference.rerank_model_settings == {"top_p": 0.25} + assert "generation-secret" not in repr(inference) + assert "embedding-secret" not in repr(inference) + assert "rerank-secret" not in repr(inference) + + +@pytest.mark.parametrize( + "values", + [ + {"generation_headers": {"X-Workload": "secret"}}, + {"embedding_model_settings": {"dimensions": 3}}, + {"rerank_base_url": "http://rerank.test/v1"}, + ], +) +def test_inference_settings_reject_orphaned_workload_overrides(values: dict[str, object]) -> None: + with pytest.raises(ValidationError): + InferenceConfig.model_validate(values) + + +def test_inference_settings_keep_headers_out_of_model_settings() -> None: + with pytest.raises(ValidationError, match="dedicated headers field"): + InferenceConfig( + generation_model="openai-chat:generator", + generation_model_settings={"extra_headers": {"X-Workload": "secret"}}, + ) + + +def test_inference_settings_hide_header_values_in_validation_errors() -> None: + with pytest.raises(ValidationError) as captured: + InferenceConfig( + generation_model="openai-chat:generator", + generation_headers={"Bad:Name": SecretStr("validation-secret")}, + ) + + assert "validation-secret" not in str(captured.value) + + +@pytest.mark.parametrize("header_name", ["Bad Header", "Bad\tHeader", "X-Ünicode"]) +def test_inference_settings_reject_invalid_http_field_names(header_name: str) -> None: + with pytest.raises(ValidationError, match="HTTP field names"): + InferenceConfig( + generation_model="openai-chat:generator", + generation_headers={header_name: SecretStr("secret")}, + ) + + +def test_anthropic_custom_endpoint_uses_standard_environment_api_key( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-standard-key") + + async def scenario() -> None: + with _model_server() as (server, base_url): + config = BuiltinConfig( + database=SQLiteConfig(url=f"sqlite+aiosqlite:///{tmp_path / 'runtime.db'}"), + inference=InferenceConfig( + generation_model="anthropic:tiny-generator", + generation_base_url=AnyHttpUrl(base_url.removesuffix("/v1")), + ), + ) + + async with open_builtin_runtime(config) as runtime: + readiness = await runtime.readiness() + + assert readiness.checks["inference.generation"] == "ready" + assert server.requests + assert server.requests[0]["path"] == "/v1/messages" + assert server.requests[0]["headers"]["x-api-key"] == "anthropic-standard-key" + + asyncio.run(scenario()) + + +def test_generation_embedding_and_llm_rerank_models_receive_their_own_settings( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + caplog.set_level(logging.DEBUG, logger="openai._base_client") + + async def scenario() -> None: + with ( + _model_server() as (generation_server, generation_url), + _model_server() as (embedding_server, embedding_url), + _model_server() as (rerank_server, rerank_url), + ): + config = BuiltinConfig( + database=SQLiteConfig(url=f"sqlite+aiosqlite:///{tmp_path / 'runtime.db'}"), + runtime=RuntimeConfig(memory_rerank_enabled=True), + inference=InferenceConfig( + generation_model="openai-chat:tiny-generator", + generation_base_url=AnyHttpUrl(generation_url), + generation_headers={"X-Workload": SecretStr("generation-secret")}, + generation_model_settings={"top_p": 0.5, "extra_body": {"route": "generation"}}, + embedding_model="openai:tiny-embedding", + embedding_base_url=AnyHttpUrl(embedding_url), + embedding_headers={"X-Workload": SecretStr("embedding-secret")}, + embedding_model_settings={"dimensions": 3, "extra_body": {"route": "embedding"}}, + embedding_profile_id="tiny-embedding-v1", + embedding_dimension=3, + rerank_model="openai-chat:tiny-reranker", + rerank_base_url=AnyHttpUrl(rerank_url), + rerank_headers={"X-Workload": SecretStr("rerank-secret")}, + rerank_model_settings={"top_p": 0.25, "extra_body": {"route": "rerank"}}, + ), + ) + + async with open_builtin_runtime(config) as runtime: + readiness = await runtime.readiness() + memory = runtime.memory.for_scope("custom-inference") + await memory.remember( + RememberMemoryRequest( + entries=( + MemoryEntryInput(kind="fact", text="Deployment uses the blue environment."), + MemoryEntryInput(kind="fact", text="Deployment rollback uses the green environment."), + ) + ) + ) + search = await memory.search(SearchMemoryRequest(query="deployment environment", mode="fts", limit=1)) + + assert readiness.status.value == "ready" + assert readiness.checks["inference.generation"] == "ready" + assert readiness.checks["inference.embedding"] == "ready" + assert readiness.checks["inference.rerank"] == "ready" + assert search.rerank is not None + assert search.rerank.selected_ranks == (1,) + + assert generation_server.requests + for generation_request in generation_server.requests: + assert generation_request["path"] == "/v1/chat/completions" + assert generation_request["headers"]["x-workload"] == "generation-secret" + assert generation_request["body"]["model"] == "tiny-generator" + assert generation_request["body"]["top_p"] == 0.5 + assert generation_request["body"]["route"] == "generation" + + assert embedding_server.requests + for embedding_request in embedding_server.requests: + assert embedding_request["path"] == "/v1/embeddings" + assert embedding_request["headers"]["x-workload"] == "embedding-secret" + assert embedding_request["body"]["model"] == "tiny-embedding" + assert embedding_request["body"]["dimensions"] == 3 + assert embedding_request["body"]["route"] == "embedding" + + assert rerank_server.requests + for rerank_request in rerank_server.requests: + assert rerank_request["path"] == "/v1/chat/completions" + assert rerank_request["headers"]["x-workload"] == "rerank-secret" + assert rerank_request["body"]["model"] == "tiny-reranker" + assert rerank_request["body"]["top_p"] == 0.25 + assert rerank_request["body"]["temperature"] == 0.0 + assert rerank_request["body"]["route"] == "rerank" + + log_output = "\n".join(record.getMessage() for record in caplog.records) + assert "generation-secret" not in log_output + assert "embedding-secret" not in log_output + assert "rerank-secret" not in log_output + + asyncio.run(scenario())