Skip to content

Normalize completion responses at the provider boundary - #2254

Draft
gold-silver-copper wants to merge 6 commits into
mainfrom
refactor/normalize-completion-responses
Draft

Normalize completion responses at the provider boundary#2254
gold-silver-copper wants to merge 6 commits into
mainfrom
refactor/normalize-completion-responses

Conversation

@gold-silver-copper

Copy link
Copy Markdown
Contributor

Summary

This PR normalizes completion responses at the provider boundary:

  • CompletionModel now returns concrete CompletionResponse and StreamingCompletionResponse values.
  • Model construction moves to CompletionClient; the model trait no longer owns response/client associated types or make.
  • Normalized responses expose provider, model, message ID, finish reason, and usage metadata.
  • Provider-native response access remains available through typed raw_completion and raw_stream escape hatches.
  • Agent APIs stay generic over the model while response-type generics disappear from the runner path.
  • Unary and streaming paths share tool-turn disambiguation, terminal-state handling, and message-ID promotion.

Breaking API changes are intentional; backwards compatibility is not a constraint for this prerequisite.

Relationship to other work

This is a prerequisite for #2252. It does not modify that PR or introduce runtime model-selection architecture. The response-shape normalization is independently useful and keeps the later hook-driven selection work focused on run-time steering.

Core API shape

pub trait CompletionModel: Clone + WasmCompatSend + WasmCompatSync {
    fn completion(
        &self,
        request: CompletionRequest,
    ) -> impl Future<Output = Result<CompletionResponse, CompletionError>> + WasmCompatSend;

    fn stream(
        &self,
        request: CompletionRequest,
    ) -> impl Future<Output = Result<StreamingCompletionResponse, CompletionError>>
           + WasmCompatSend;
}

pub trait CompletionClient<M>: Clone + WasmCompatSend + WasmCompatSync {
    type CompletionModel: CompletionModel;

    fn completion_model(&self, model: M) -> Self::CompletionModel;
}

ProviderCapabilities is public. RawStreamingChoice<R>, RawStreamingResult<R>, and normalize_stream provide the reusable typed raw-stream boundary.

Provider-extension associated types such as OpenAICompatibleProvider::{Response, StreamingUsage} intentionally remain at the native provider boundary; completion models and agents no longer propagate response associated types.

Native-response escape hatches

Direct typed raw APIs are implemented for:

  • Anthropic unary and streaming responses
  • OpenAI Completions and Responses APIs
  • ChatGPT Responses API
  • Cohere
  • Copilot chat and Responses routes
  • Gemini REST and Interactions
  • Ollama
  • xAI
  • Bedrock, Candle, Gemini gRPC, and Vertex AI companion crates

OpenAI-compatible extension types cover Azure, DeepSeek, Doubleword, Groq, Hugging Face, Hyperbolic, Llamafile, MiniMax, Mira, Mistral, Moonshot, OpenRouter, Perplexity, Together, Xiaomi MiMo, and Z.AI. Anthropic-compatible provider extensions retain their native response types as well.

Finish-reason normalization

Provider family Native terminal reason Normalized result
OpenAI-compatible chat / Copilot chat stop; length; tool_calls or function_call; content_filter Stop; Length; ToolCalls; ContentFilter
Responses / ChatGPT / Copilot Responses / xAI completed; incomplete max; content filter Stop (or ToolCalls when calls are present); Length; ContentFilter
Anthropic end_turn / stop_sequence; max_tokens; tool_use Stop; Length; ToolCalls
Cohere COMPLETE / STOP_SEQUENCE; MAX_TOKENS; TOOL_CALL Stop; Length; ToolCalls
Gemini REST / gRPC / Vertex STOP; MAX_TOKENS; safety/block reasons Stop or ToolCalls when calls are present; Length; ContentFilter
Gemini Interactions COMPLETED; REQUIRES_ACTION; BUDGET_EXCEEDED Stop; ToolCalls; Length
Ollama stop; length Stop or ToolCalls when calls are present; Length
Bedrock end turn / stop sequence; max tokens; tool use; content filter / guardrail Stop; Length; ToolCalls; ContentFilter
Candle EOS; max tokens Stop or ToolCalls when parsed calls are present; Length

Unknown native reasons are preserved as FinishReason::Other rather than discarded.

Correctness details

  • Native message IDs are preserved without confusing Responses API resp_* IDs with assistant msg_* IDs.
  • Streaming terminal IDs populate normalized stream and agent state while explicit message-ID events retain precedence.
  • Anthropic, Cohere, and Gemini streams do not synthesize successful terminal values after provider errors or truncated streams.
  • Tool-bearing turns normalize to ToolCalls consistently across unary and streaming paths.
  • Bedrock stop-sequence and guardrail aliases preserve their actual normalized meanings.

Verification

Passed locally:

  • cargo fmt --all -- --check
  • cargo clippy --workspace --all-targets --all-features
  • cargo test
  • cargo test -p rig-core --lib --all-features — 1,041 passed, 3 ignored
  • cargo test -p rig-agent --lib --all-features — 487 passed, 2 ignored
  • companion-crate tests for Bedrock, Gemini gRPC, Vertex AI, and Candle
  • strict rustdoc: RUSTDOCFLAGS='-D warnings' cargo doc -p rig-core -p rig-agent -p rig --no-deps
  • WASM check: cargo check --target wasm32-unknown-unknown -p rig-core -p rig-agent -p rig
  • cassette replay coverage for OpenAI, Anthropic, Gemini, ChatGPT, Bedrock, Doubleword, and OpenRouter

No cassette fixture files changed.

An independent reviewer inspected the complete 107-file diff against merge base 6cfae6d829da21f9dc9e775e065fee157b264f7e, reviewed all follow-up fixes, and reported no remaining P0–P3, security, unsafe-code, or regression findings.

Phase P1 of the data-oriented rearchitecture
(audit/data-oriented-rearchitecture.md §14, Revision 2.1).

The provider-typed response generics are gone from the payload layer:

- CompletionResponse<T> -> concrete CompletionResponse with normalized
  metadata: finish_reason (new FinishReason enum — closes #2090, #1886),
  provider, model, message_id. raw_response is removed; provider-typed
  wire data stays reachable through each provider's own conversion layer.
- The provider-typed streaming final becomes the normalized StreamFinal
  (kind discriminant for the untagged enum, usage as a field, finish
  reason, message_id, provider, model). The GetTokenUsage trait is
  deleted; wire usage types convert via From/Into<Usage>.
- RawStreamingChoice, StreamingResult, StreamingCompletionResponse,
  StreamedAssistantContent, MultiTurnStreamItem, DriveStream/DriveItem and
  the StreamedTurnAssembler lose their type parameters; the entire
  rig-agent run module is now generic-free.
- CompletionModel loses type Response / type StreamingResponse (the trait
  itself survives until P7); TurnSource loses type Raw entirely;
  StreamingPrompt/StreamingChat lose their response parameter.
- Every finish/stop vocabulary is mapped per provider (25 in-core + the
  four companion crates); the OpenAI-compatible shared path stamps
  Ext::PROVIDER_NAME post-conversion.
- Provider-specific streaming-final aliases (deepseek, groq, mistral,
  openrouter, gemini, copilot, candle placeholder) are removed.

Verification: cargo check --workspace --all-targets clean; clippy clean;
rig-core 1008 tests, rig-agent 500 tests, facade cassette suites 1214
tests all passing with cassette YAMLs byte-identical (no request-building
code changed). Full narrative in audit/migration-log.md.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant