Fix Mentra nested Responses endpoints
This commit is contained in:
2
Cargo.lock
generated
2
Cargo.lock
generated
@@ -4556,8 +4556,6 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "mentra-provider"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e9f644037cf45558b393becc1ac080f0d55edced3093a155a7cd3431c1180b7f"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"base64",
|
||||
|
||||
@@ -25,7 +25,7 @@ members = [
|
||||
"tools/concurrency-audit",
|
||||
"tools/performance",
|
||||
]
|
||||
exclude = ["crates/libremetaverse-openjpeg"]
|
||||
exclude = ["crates/libremetaverse-openjpeg", "vendor/mentra-provider"]
|
||||
default-members = [
|
||||
"crates/metacrate",
|
||||
"crates/libremetaverse-types",
|
||||
@@ -78,3 +78,7 @@ inherits = "dev"
|
||||
opt-level = 1
|
||||
debug = 0
|
||||
incremental = false
|
||||
|
||||
[patch.crates-io]
|
||||
# Mentra 0.5.1 duplicates `v1` for nested OpenAI-compatible base URLs.
|
||||
mentra-provider = { path = "vendor/mentra-provider" }
|
||||
|
||||
@@ -50,10 +50,19 @@ async fn luna_describes_live_renderer_evidence() -> Result<(), Box<dyn Error>> {
|
||||
ContentBlock::image_url(image),
|
||||
])
|
||||
.await?;
|
||||
for block in response.content {
|
||||
if let ContentBlock::Text { text } = block {
|
||||
println!("LIVE_LLM_VISION={text}");
|
||||
let descriptions = response
|
||||
.content
|
||||
.into_iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::Text { text } if !text.trim().is_empty() => Some(text),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if descriptions.is_empty() {
|
||||
return Err("Luna returned no snapshot description".into());
|
||||
}
|
||||
for description in descriptions {
|
||||
println!("LIVE_LLM_VISION={description}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
99
vendor/mentra-provider/Cargo.toml
vendored
Normal file
99
vendor/mentra-provider/Cargo.toml
vendored
Normal file
@@ -0,0 +1,99 @@
|
||||
# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO
|
||||
#
|
||||
# When uploading crates to the registry Cargo will automatically
|
||||
# "normalize" Cargo.toml files for maximal compatibility
|
||||
# with all versions of Cargo and also rewrite `path` dependencies
|
||||
# to registry (e.g., crates.io) dependencies.
|
||||
#
|
||||
# If you are reading this file be aware that the original Cargo.toml
|
||||
# will likely look very different (and much more reasonable).
|
||||
# See Cargo.toml.orig for the original contents.
|
||||
|
||||
[package]
|
||||
edition = "2024"
|
||||
rust-version = "1.88"
|
||||
name = "mentra-provider"
|
||||
version = "0.5.1"
|
||||
build = false
|
||||
autolib = false
|
||||
autobins = false
|
||||
autoexamples = false
|
||||
autotests = false
|
||||
autobenches = false
|
||||
description = "Shared provider core for Mentra"
|
||||
homepage = "https://github.com/oops-rs/mentra"
|
||||
documentation = "https://docs.rs/mentra-provider"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/oops-rs/mentra"
|
||||
|
||||
[features]
|
||||
default = ["responses-websocket"]
|
||||
responses-websocket = [
|
||||
"dep:tokio-tungstenite",
|
||||
"futures-util/sink",
|
||||
]
|
||||
|
||||
[lib]
|
||||
name = "mentra_provider"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies.async-trait]
|
||||
version = "0.1.89"
|
||||
|
||||
[dependencies.base64]
|
||||
version = "0.22.1"
|
||||
|
||||
[dependencies.futures-util]
|
||||
version = "0.3.31"
|
||||
|
||||
[dependencies.http]
|
||||
version = "1.3.1"
|
||||
|
||||
[dependencies.reqwest]
|
||||
version = "0.12"
|
||||
features = [
|
||||
"json",
|
||||
"rustls-tls",
|
||||
"stream",
|
||||
]
|
||||
default-features = false
|
||||
|
||||
[dependencies.serde]
|
||||
version = "1.0.228"
|
||||
features = ["derive"]
|
||||
|
||||
[dependencies.serde_json]
|
||||
version = "1.0.149"
|
||||
|
||||
[dependencies.strum]
|
||||
version = "0.27"
|
||||
features = ["derive"]
|
||||
|
||||
[dependencies.thiserror]
|
||||
version = "2.0.18"
|
||||
|
||||
[dependencies.time]
|
||||
version = "0.3"
|
||||
features = [
|
||||
"formatting",
|
||||
"parsing",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[dependencies.tokio]
|
||||
version = "1.50.0"
|
||||
features = [
|
||||
"macros",
|
||||
"sync",
|
||||
]
|
||||
|
||||
[dependencies.tokio-tungstenite]
|
||||
version = "0.28"
|
||||
optional = true
|
||||
|
||||
[dependencies.url]
|
||||
version = "2.5"
|
||||
|
||||
[dependencies.zstd]
|
||||
version = "0.13"
|
||||
21
vendor/mentra-provider/LICENSE
vendored
Normal file
21
vendor/mentra-provider/LICENSE
vendored
Normal file
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 Wendell Wang
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
11
vendor/mentra-provider/README.md
vendored
Normal file
11
vendor/mentra-provider/README.md
vendored
Normal file
@@ -0,0 +1,11 @@
|
||||
# mentra-provider
|
||||
|
||||
`mentra-provider` is Mentra's publishable provider-core crate.
|
||||
|
||||
It contains provider-neutral request, response, model, streaming, and tool
|
||||
schema types that can be reused without depending on the full Mentra runtime.
|
||||
|
||||
For most application code, depend on
|
||||
[`mentra`](https://crates.io/crates/mentra) instead. The `mentra` crate
|
||||
re-exports these provider-core types and adds the runtime, tooling,
|
||||
persistence, and collaboration layers.
|
||||
271
vendor/mentra-provider/src/anthropic.rs
vendored
Normal file
271
vendor/mentra-provider/src/anthropic.rs
vendored
Normal file
@@ -0,0 +1,271 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub(crate) mod model;
|
||||
pub(crate) mod sse;
|
||||
pub(crate) mod stream_model;
|
||||
|
||||
use crate::AuthScheme;
|
||||
use crate::BuiltinProvider;
|
||||
use crate::CompactionRequest;
|
||||
use crate::CompactionResponse;
|
||||
use crate::CredentialSource;
|
||||
use crate::ModelCatalog;
|
||||
use crate::ModelInfo;
|
||||
use crate::ProviderCapabilities;
|
||||
use crate::ProviderDefinition;
|
||||
use crate::ProviderError;
|
||||
use crate::ProviderEventStream;
|
||||
use crate::ProviderSession;
|
||||
use crate::ProviderSessionFactory;
|
||||
use crate::RegisteredProvider;
|
||||
use crate::Request;
|
||||
use crate::StaticCredentialSource;
|
||||
use crate::WireApi;
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
|
||||
const ANTHROPIC_VERSION: &str = "2023-06-01";
|
||||
|
||||
/// Returns the default Anthropic-compatible provider definition.
|
||||
pub fn definition() -> ProviderDefinition {
|
||||
let mut definition = ProviderDefinition::new(BuiltinProvider::Anthropic);
|
||||
definition.descriptor.display_name = Some("Anthropic".to_string());
|
||||
definition.descriptor.description = Some("Anthropic Messages API provider".to_string());
|
||||
definition.wire_api = WireApi::AnthropicMessages;
|
||||
definition.auth_scheme = AuthScheme::Header {
|
||||
name: "x-api-key".to_string(),
|
||||
};
|
||||
definition.capabilities = ProviderCapabilities {
|
||||
supports_model_listing: true,
|
||||
supports_streaming: true,
|
||||
supports_websockets: false,
|
||||
supports_tool_calls: true,
|
||||
supports_images: true,
|
||||
supports_history_compaction: true,
|
||||
supports_memory_summarization: true,
|
||||
supports_deferred_tools: true,
|
||||
supports_hosted_tool_search: true,
|
||||
supports_hosted_web_search: false,
|
||||
supports_image_generation: false,
|
||||
supports_reasoning_effort: true,
|
||||
reports_reasoning_tokens: false,
|
||||
reports_thoughts_tokens: false,
|
||||
supports_structured_tool_results: false,
|
||||
supports_embeddings: false,
|
||||
};
|
||||
definition.base_url = Some(DEFAULT_BASE_URL.to_string());
|
||||
definition.headers = Some(HashMap::from([(
|
||||
"anthropic-version".to_string(),
|
||||
ANTHROPIC_VERSION.to_string(),
|
||||
)]));
|
||||
definition
|
||||
}
|
||||
|
||||
pub struct AnthropicProvider<C = StaticCredentialSource> {
|
||||
client: reqwest::Client,
|
||||
credential_source: Arc<C>,
|
||||
definition: ProviderDefinition,
|
||||
}
|
||||
|
||||
impl<C> Clone for AnthropicProvider<C> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
client: self.client.clone(),
|
||||
credential_source: Arc::clone(&self.credential_source),
|
||||
definition: self.definition.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AnthropicProvider<StaticCredentialSource> {
|
||||
pub fn new(api_key: impl Into<String>) -> Self {
|
||||
Self::with_credential_source(StaticCredentialSource::new(api_key))
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> AnthropicProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
pub fn with_credential_source(credential_source: C) -> Self {
|
||||
Self::with_shared_credential_source(Arc::new(credential_source))
|
||||
}
|
||||
|
||||
pub fn with_shared_credential_source(credential_source: Arc<C>) -> Self {
|
||||
Self::with_definition_and_shared_credential_source(definition(), credential_source)
|
||||
}
|
||||
|
||||
pub fn with_definition_and_credential_source(
|
||||
definition: ProviderDefinition,
|
||||
credential_source: C,
|
||||
) -> Self {
|
||||
Self::with_definition_and_shared_credential_source(definition, Arc::new(credential_source))
|
||||
}
|
||||
|
||||
pub fn with_definition_and_shared_credential_source(
|
||||
definition: ProviderDefinition,
|
||||
credential_source: Arc<C>,
|
||||
) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.build()
|
||||
.expect("Failed to build client");
|
||||
|
||||
Self {
|
||||
client,
|
||||
credential_source,
|
||||
definition,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ModelCatalog for AnthropicProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
let mut models = Vec::new();
|
||||
let mut after_id = None;
|
||||
|
||||
loop {
|
||||
let credentials = self.credential_source.credentials().await?;
|
||||
let request = self
|
||||
.client
|
||||
.get(
|
||||
self.definition
|
||||
.request_url_with_auth_for_path("v1/models", &credentials)?,
|
||||
)
|
||||
.headers(self.definition.build_headers(&credentials)?)
|
||||
.query(&[
|
||||
("limit", "1000"),
|
||||
("after_id", after_id.as_deref().unwrap_or("")),
|
||||
]);
|
||||
|
||||
let response = request.send().await.map_err(ProviderError::Transport)?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::from_http_response(response).await);
|
||||
}
|
||||
|
||||
let page = response
|
||||
.json::<model::AnthropicModelsPage>()
|
||||
.await
|
||||
.map_err(ProviderError::Decode)?;
|
||||
|
||||
after_id = page.last_id.clone();
|
||||
models.extend(page.data.into_iter().map(|model| model.into()));
|
||||
|
||||
if !page.has_more {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(models)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ProviderSessionFactory for AnthropicProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn create_session(&self) -> Result<Box<dyn ProviderSession>, ProviderError> {
|
||||
Ok(Box::new((*self).clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ProviderSession for AnthropicProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
|
||||
let requested_model = request.model.to_string();
|
||||
let provider = self.definition.provider_id().clone();
|
||||
let response = self.send_message(request, true).await?;
|
||||
Ok(sse::spawn_event_stream(response, provider, requested_model))
|
||||
}
|
||||
|
||||
async fn compact(
|
||||
&self,
|
||||
request: CompactionRequest<'_>,
|
||||
) -> Result<CompactionResponse, ProviderError> {
|
||||
let request = request.into_model_request()?;
|
||||
let response = ProviderSession::send(self, request).await?;
|
||||
Ok(response.into_compaction_response())
|
||||
}
|
||||
|
||||
async fn summarize_memories(
|
||||
&self,
|
||||
request: crate::MemorySummarizeRequest<'_>,
|
||||
) -> Result<crate::MemorySummarizeResponse, ProviderError> {
|
||||
let request = request.into_model_request()?;
|
||||
let response = ProviderSession::send(self, request).await?;
|
||||
response.into_memory_summarize_response()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> RegisteredProvider for AnthropicProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
fn definition(&self) -> ProviderDefinition {
|
||||
self.definition.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> AnthropicProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn send_message(
|
||||
&self,
|
||||
request: Request<'_>,
|
||||
stream: bool,
|
||||
) -> Result<reqwest::Response, ProviderError> {
|
||||
let session = request.provider_request_options.session.clone();
|
||||
let request = model::AnthropicRequest::try_from_with_provider(
|
||||
request,
|
||||
self.definition.provider_id(),
|
||||
)?;
|
||||
let mut body = serde_json::to_value(request).map_err(ProviderError::Serialize)?;
|
||||
if stream {
|
||||
body["stream"] = Value::Bool(true);
|
||||
}
|
||||
let credentials = self.credential_source.credentials().await?;
|
||||
let response = self
|
||||
.client
|
||||
.post(
|
||||
self.definition
|
||||
.request_url_with_auth_for_path("v1/messages", &credentials)?,
|
||||
)
|
||||
.headers(self.definition.build_headers_for_session(
|
||||
&credentials,
|
||||
Some(&session),
|
||||
None,
|
||||
)?)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(ProviderError::Transport)?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::from_http_response(response).await);
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn definition_advertises_history_compaction_support() {
|
||||
assert!(definition().capabilities.supports_history_compaction);
|
||||
}
|
||||
}
|
||||
1397
vendor/mentra-provider/src/anthropic/model.rs
vendored
Normal file
1397
vendor/mentra-provider/src/anthropic/model.rs
vendored
Normal file
File diff suppressed because it is too large
Load Diff
460
vendor/mentra-provider/src/anthropic/sse.rs
vendored
Normal file
460
vendor/mentra-provider/src/anthropic/sse.rs
vendored
Normal file
@@ -0,0 +1,460 @@
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::{
|
||||
ProviderError, ProviderEvent, ProviderEventStream, ProviderId, ReasoningFormat,
|
||||
ReasoningProvenance, TokenUsage,
|
||||
};
|
||||
|
||||
use super::stream_model::AnthropicContentBlockDelta;
|
||||
use super::stream_model::AnthropicStreamContentBlock;
|
||||
use super::stream_model::AnthropicStreamEvent;
|
||||
|
||||
pub(crate) fn spawn_event_stream(
|
||||
response: reqwest::Response,
|
||||
provider: ProviderId,
|
||||
requested_model: String,
|
||||
) -> ProviderEventStream {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = forward_events(response, tx.clone(), provider, requested_model).await {
|
||||
let _ = tx.send(Err(error));
|
||||
}
|
||||
});
|
||||
|
||||
rx
|
||||
}
|
||||
|
||||
async fn forward_events(
|
||||
response: reqwest::Response,
|
||||
tx: mpsc::UnboundedSender<Result<ProviderEvent, ProviderError>>,
|
||||
provider: ProviderId,
|
||||
requested_model: String,
|
||||
) -> Result<(), ProviderError> {
|
||||
let mut bytes_stream = response.bytes_stream();
|
||||
let mut buffer = Vec::new();
|
||||
let mut state = StreamState::new(provider, requested_model);
|
||||
|
||||
while let Some(chunk) = bytes_stream.next().await {
|
||||
let chunk = chunk.map_err(ProviderError::Transport)?;
|
||||
buffer.extend_from_slice(&chunk);
|
||||
|
||||
while let Some((frame_end, delimiter_len)) = find_frame_boundary(&buffer) {
|
||||
let frame = buffer.drain(..frame_end).collect::<Vec<_>>();
|
||||
buffer.drain(..delimiter_len);
|
||||
|
||||
for event in parse_frame(&frame, &mut state)? {
|
||||
if tx.send(Ok(event)).is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !buffer.is_empty() {
|
||||
for event in parse_frame(&buffer, &mut state)? {
|
||||
let _ = tx.send(Ok(event));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct StreamState {
|
||||
ignored_blocks: HashSet<usize>,
|
||||
latest_usage: Option<TokenUsage>,
|
||||
block_kinds: HashMap<usize, StreamingBlockKind>,
|
||||
reasoning_provenance: ReasoningProvenance,
|
||||
}
|
||||
|
||||
impl StreamState {
|
||||
fn new(provider: ProviderId, requested_model: String) -> Self {
|
||||
Self {
|
||||
ignored_blocks: HashSet::new(),
|
||||
latest_usage: None,
|
||||
block_kinds: HashMap::new(),
|
||||
reasoning_provenance: ReasoningProvenance {
|
||||
provider,
|
||||
model: requested_model,
|
||||
format: ReasoningFormat::AnthropicSigned,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum StreamingBlockKind {
|
||||
ToolUse,
|
||||
HostedToolSearch,
|
||||
}
|
||||
|
||||
fn parse_frame(frame: &[u8], state: &mut StreamState) -> Result<Vec<ProviderEvent>, ProviderError> {
|
||||
let frame = std::str::from_utf8(frame)
|
||||
.map_err(|error| ProviderError::MalformedStream(error.to_string()))?;
|
||||
let mut data_lines = Vec::new();
|
||||
|
||||
for raw_line in frame.lines() {
|
||||
let line = raw_line.strip_suffix('\r').unwrap_or(raw_line);
|
||||
if line.is_empty() || line.starts_with(':') {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(rest) = line.strip_prefix("data:") {
|
||||
data_lines.push(rest.trim_start().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
if data_lines.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let data = data_lines.join("\n");
|
||||
let event: AnthropicStreamEvent =
|
||||
serde_json::from_str(&data).map_err(ProviderError::Deserialize)?;
|
||||
|
||||
if let AnthropicStreamEvent::ContentBlockStart {
|
||||
index,
|
||||
content_block,
|
||||
} = &event
|
||||
{
|
||||
match content_block {
|
||||
AnthropicStreamContentBlock::ToolUse { .. } => {
|
||||
state
|
||||
.block_kinds
|
||||
.insert(*index, StreamingBlockKind::ToolUse);
|
||||
}
|
||||
AnthropicStreamContentBlock::ServerToolUse { name, .. }
|
||||
if name.starts_with("tool_search") =>
|
||||
{
|
||||
state
|
||||
.block_kinds
|
||||
.insert(*index, StreamingBlockKind::HostedToolSearch);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
match &event {
|
||||
AnthropicStreamEvent::ContentBlockStart {
|
||||
index,
|
||||
content_block,
|
||||
} if !content_block.is_supported() => {
|
||||
state.ignored_blocks.insert(*index);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
AnthropicStreamEvent::ContentBlockDelta { index, .. }
|
||||
| AnthropicStreamEvent::ContentBlockStop { index }
|
||||
if state.ignored_blocks.contains(index) =>
|
||||
{
|
||||
if matches!(event, AnthropicStreamEvent::ContentBlockStop { .. }) {
|
||||
state.ignored_blocks.remove(index);
|
||||
}
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
AnthropicStreamEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: AnthropicContentBlockDelta::InputJsonDelta { partial_json },
|
||||
} => {
|
||||
if matches!(
|
||||
state.block_kinds.get(index),
|
||||
Some(StreamingBlockKind::HostedToolSearch)
|
||||
) {
|
||||
return Ok(vec![ProviderEvent::ContentBlockDelta {
|
||||
index: *index,
|
||||
delta: ProviderEventDeltaExt::hosted_tool_search_delta(partial_json),
|
||||
}]);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let events = event
|
||||
.into_provider_events(&state.reasoning_provenance)
|
||||
.map_err(|error| {
|
||||
ProviderError::MalformedStream(format!(
|
||||
"anthropic stream error ({}): {}",
|
||||
error.kind, error.message
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(events
|
||||
.into_iter()
|
||||
.map(|event| match event {
|
||||
ProviderEvent::ContentBlockStopped { index } => {
|
||||
state.block_kinds.remove(&index);
|
||||
ProviderEvent::ContentBlockStopped { index }
|
||||
}
|
||||
ProviderEvent::MessageDelta { stop_reason, usage } => {
|
||||
let usage = merge_usage(state.latest_usage.clone(), usage);
|
||||
state.latest_usage = usage.clone();
|
||||
ProviderEvent::MessageDelta { stop_reason, usage }
|
||||
}
|
||||
other => other,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
struct ProviderEventDeltaExt;
|
||||
|
||||
impl ProviderEventDeltaExt {
|
||||
fn hosted_tool_search_delta(partial_json: &str) -> crate::ContentBlockDelta {
|
||||
crate::ContentBlockDelta::HostedToolSearchQuery(
|
||||
extract_tool_search_query(partial_json).unwrap_or_else(|| partial_json.to_string()),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_tool_search_query(partial_json: &str) -> Option<String> {
|
||||
serde_json::from_str::<serde_json::Value>(partial_json)
|
||||
.ok()
|
||||
.and_then(|value| {
|
||||
value
|
||||
.get("query")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_string)
|
||||
})
|
||||
}
|
||||
|
||||
fn merge_usage(base: Option<TokenUsage>, update: Option<TokenUsage>) -> Option<TokenUsage> {
|
||||
match (base, update) {
|
||||
(Some(base), Some(update)) => {
|
||||
let merged = TokenUsage {
|
||||
input_tokens: update.input_tokens.or(base.input_tokens),
|
||||
output_tokens: update.output_tokens.or(base.output_tokens),
|
||||
total_tokens: update.total_tokens.or(base.total_tokens),
|
||||
cache_read_input_tokens: update
|
||||
.cache_read_input_tokens
|
||||
.or(base.cache_read_input_tokens),
|
||||
cache_creation_input_tokens: update
|
||||
.cache_creation_input_tokens
|
||||
.or(base.cache_creation_input_tokens),
|
||||
reasoning_tokens: update.reasoning_tokens.or(base.reasoning_tokens),
|
||||
thoughts_tokens: update.thoughts_tokens.or(base.thoughts_tokens),
|
||||
tool_input_tokens: update.tool_input_tokens.or(base.tool_input_tokens),
|
||||
};
|
||||
Some(merged)
|
||||
}
|
||||
(Some(base), None) => Some(base),
|
||||
(None, Some(update)) => Some(update),
|
||||
(None, None) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn find_frame_boundary(buffer: &[u8]) -> Option<(usize, usize)> {
|
||||
for (index, window) in buffer.windows(2).enumerate() {
|
||||
if window == b"\n\n" {
|
||||
return Some((index, 2));
|
||||
}
|
||||
}
|
||||
|
||||
for (index, window) in buffer.windows(4).enumerate() {
|
||||
if window == b"\r\n\r\n" {
|
||||
return Some((index, 4));
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StreamState, parse_frame};
|
||||
use crate::{ProviderEvent, ProviderId, Role, TokenUsage};
|
||||
|
||||
fn stream_state() -> StreamState {
|
||||
StreamState::new(
|
||||
ProviderId::new("anthropic-edge"),
|
||||
"claude-requested".to_string(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_anthropic_usage_updates_into_cumulative_totals() {
|
||||
let mut state = stream_state();
|
||||
|
||||
let started = parse_frame(
|
||||
br#"data: {"type":"message_start","message":{"id":"msg_1","model":"claude-sonnet","role":"assistant","content":[],"usage":{"input_tokens":10,"cache_read_input_tokens":2}}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("message start should parse");
|
||||
assert_eq!(
|
||||
started,
|
||||
vec![
|
||||
ProviderEvent::MessageStarted {
|
||||
id: "msg_1".to_string(),
|
||||
model: "claude-sonnet".to_string(),
|
||||
role: Role::Assistant,
|
||||
},
|
||||
ProviderEvent::MessageDelta {
|
||||
stop_reason: None,
|
||||
usage: Some(TokenUsage {
|
||||
input_tokens: Some(10),
|
||||
output_tokens: None,
|
||||
total_tokens: None,
|
||||
cache_read_input_tokens: Some(2),
|
||||
cache_creation_input_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
thoughts_tokens: None,
|
||||
tool_input_tokens: None,
|
||||
}),
|
||||
},
|
||||
]
|
||||
);
|
||||
|
||||
let delta = parse_frame(
|
||||
br#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn","usage":{"output_tokens":3}}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("message delta should parse");
|
||||
assert_eq!(
|
||||
delta,
|
||||
vec![ProviderEvent::MessageDelta {
|
||||
stop_reason: Some("end_turn".to_string()),
|
||||
usage: Some(TokenUsage {
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(3),
|
||||
total_tokens: None,
|
||||
cache_read_input_tokens: Some(2),
|
||||
cache_creation_input_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
thoughts_tokens: None,
|
||||
tool_input_tokens: None,
|
||||
}),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_hosted_tool_search_bookkeeping_blocks() {
|
||||
let mut state = stream_state();
|
||||
|
||||
let started = parse_frame(
|
||||
br#"data: {"type":"content_block_start","index":1,"content_block":{"type":"server_tool_use","id":"srvtoolu_1","name":"tool_search_tool_bm25"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("server tool use should parse");
|
||||
assert_eq!(
|
||||
started,
|
||||
vec![ProviderEvent::ContentBlockStarted {
|
||||
index: 1,
|
||||
kind: crate::ContentBlockStart::HostedToolSearch {
|
||||
call: crate::HostedToolSearchCall {
|
||||
id: "srvtoolu_1".to_string(),
|
||||
status: Some("in_progress".to_string()),
|
||||
query: None,
|
||||
},
|
||||
},
|
||||
}]
|
||||
);
|
||||
|
||||
let delta = parse_frame(
|
||||
br#"data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"query\":\"weather\"}"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("hosted search delta should parse");
|
||||
assert_eq!(
|
||||
delta,
|
||||
vec![ProviderEvent::ContentBlockDelta {
|
||||
index: 1,
|
||||
delta: crate::ContentBlockDelta::HostedToolSearchQuery("weather".to_string()),
|
||||
}]
|
||||
);
|
||||
|
||||
let stopped = parse_frame(
|
||||
br#"data: {"type":"content_block_stop","index":1}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("hosted search stop should parse");
|
||||
assert_eq!(
|
||||
stopped,
|
||||
vec![ProviderEvent::ContentBlockStopped { index: 1 }]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_signed_and_redacted_thinking_with_requested_provenance() {
|
||||
let mut state = stream_state();
|
||||
|
||||
let started = parse_frame(
|
||||
br#"data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("thinking start should parse");
|
||||
assert_eq!(
|
||||
started,
|
||||
vec![ProviderEvent::ContentBlockStarted {
|
||||
index: 0,
|
||||
kind: crate::ContentBlockStart::Thinking {
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(crate::ReasoningProvenance {
|
||||
provider: ProviderId::new("anthropic-edge"),
|
||||
model: "claude-requested".to_string(),
|
||||
format: crate::ReasoningFormat::AnthropicSigned,
|
||||
}),
|
||||
redacted: false,
|
||||
},
|
||||
}]
|
||||
);
|
||||
|
||||
let thinking = parse_frame(
|
||||
br#"data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"private chain"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("thinking delta should parse");
|
||||
assert_eq!(
|
||||
thinking,
|
||||
vec![ProviderEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: crate::ContentBlockDelta::ThinkingText("private chain".to_string()),
|
||||
}]
|
||||
);
|
||||
|
||||
let signature = parse_frame(
|
||||
br#"data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"opaque-signature"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("signature delta should parse");
|
||||
assert_eq!(
|
||||
signature,
|
||||
vec![ProviderEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: crate::ContentBlockDelta::ThinkingSignature("opaque-signature".to_string()),
|
||||
}]
|
||||
);
|
||||
|
||||
let redacted = parse_frame(
|
||||
br#"data: {"type":"content_block_start","index":1,"content_block":{"type":"redacted_thinking","data":"opaque-redacted-data"}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("redacted thinking should parse");
|
||||
assert_eq!(
|
||||
redacted,
|
||||
vec![
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index: 1,
|
||||
kind: crate::ContentBlockStart::Thinking {
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(crate::ReasoningProvenance {
|
||||
provider: ProviderId::new("anthropic-edge"),
|
||||
model: "claude-requested".to_string(),
|
||||
format: crate::ReasoningFormat::AnthropicSigned,
|
||||
}),
|
||||
redacted: true,
|
||||
},
|
||||
},
|
||||
ProviderEvent::ContentBlockDelta {
|
||||
index: 1,
|
||||
delta: crate::ContentBlockDelta::ThinkingSignature(
|
||||
"opaque-redacted-data".to_string()
|
||||
),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
232
vendor/mentra-provider/src/anthropic/stream_model.rs
vendored
Normal file
232
vendor/mentra-provider/src/anthropic/stream_model.rs
vendored
Normal file
@@ -0,0 +1,232 @@
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::{
|
||||
ContentBlockDelta, ContentBlockStart, HostedToolSearchCall, ProviderEvent, ReasoningProvenance,
|
||||
Role,
|
||||
};
|
||||
|
||||
use super::model::{AnthropicResponse, AnthropicUsage};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum AnthropicStreamEvent {
|
||||
MessageStart {
|
||||
message: AnthropicResponse,
|
||||
},
|
||||
ContentBlockStart {
|
||||
index: usize,
|
||||
content_block: AnthropicStreamContentBlock,
|
||||
},
|
||||
ContentBlockDelta {
|
||||
index: usize,
|
||||
delta: AnthropicContentBlockDelta,
|
||||
},
|
||||
ContentBlockStop {
|
||||
index: usize,
|
||||
},
|
||||
MessageDelta {
|
||||
delta: AnthropicMessageDelta,
|
||||
},
|
||||
MessageStop,
|
||||
Ping,
|
||||
Error {
|
||||
error: AnthropicStreamError,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum AnthropicStreamContentBlock {
|
||||
Text {},
|
||||
Thinking {
|
||||
#[serde(default)]
|
||||
thinking: String,
|
||||
},
|
||||
RedactedThinking {
|
||||
data: String,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
},
|
||||
ServerToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
},
|
||||
#[serde(other)]
|
||||
Unsupported,
|
||||
}
|
||||
|
||||
impl AnthropicStreamContentBlock {
|
||||
pub(crate) fn into_provider_events(
|
||||
self,
|
||||
index: usize,
|
||||
provenance: &ReasoningProvenance,
|
||||
) -> Vec<ProviderEvent> {
|
||||
match self {
|
||||
AnthropicStreamContentBlock::Text {} => vec![ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::Text,
|
||||
}],
|
||||
AnthropicStreamContentBlock::Thinking { thinking } => {
|
||||
let mut events = vec![ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::Thinking {
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(provenance.clone()),
|
||||
redacted: false,
|
||||
},
|
||||
}];
|
||||
if !thinking.is_empty() {
|
||||
events.push(ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::ThinkingText(thinking),
|
||||
});
|
||||
}
|
||||
events
|
||||
}
|
||||
AnthropicStreamContentBlock::RedactedThinking { data } => vec![
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::Thinking {
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(provenance.clone()),
|
||||
redacted: true,
|
||||
},
|
||||
},
|
||||
ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::ThinkingSignature(data),
|
||||
},
|
||||
],
|
||||
AnthropicStreamContentBlock::ToolUse { id, name } => {
|
||||
vec![ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::ToolUse { id, name },
|
||||
}]
|
||||
}
|
||||
AnthropicStreamContentBlock::ServerToolUse { id, name } => name
|
||||
.starts_with("tool_search")
|
||||
.then(|| ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::HostedToolSearch {
|
||||
call: HostedToolSearchCall {
|
||||
id,
|
||||
status: Some("in_progress".to_string()),
|
||||
query: None,
|
||||
},
|
||||
},
|
||||
})
|
||||
.into_iter()
|
||||
.collect(),
|
||||
AnthropicStreamContentBlock::Unsupported => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_supported(&self) -> bool {
|
||||
!matches!(self, AnthropicStreamContentBlock::Unsupported)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub(crate) enum AnthropicContentBlockDelta {
|
||||
TextDelta {
|
||||
text: String,
|
||||
},
|
||||
InputJsonDelta {
|
||||
partial_json: String,
|
||||
},
|
||||
ThinkingDelta {
|
||||
thinking: String,
|
||||
},
|
||||
SignatureDelta {
|
||||
signature: String,
|
||||
},
|
||||
#[serde(other)]
|
||||
Unsupported,
|
||||
}
|
||||
|
||||
impl AnthropicContentBlockDelta {
|
||||
pub(crate) fn into_provider_delta(self) -> Option<ContentBlockDelta> {
|
||||
match self {
|
||||
AnthropicContentBlockDelta::TextDelta { text } => Some(ContentBlockDelta::Text(text)),
|
||||
AnthropicContentBlockDelta::InputJsonDelta { partial_json } => {
|
||||
Some(ContentBlockDelta::ToolUseInputJson(partial_json))
|
||||
}
|
||||
AnthropicContentBlockDelta::ThinkingDelta { thinking } => {
|
||||
Some(ContentBlockDelta::ThinkingText(thinking))
|
||||
}
|
||||
AnthropicContentBlockDelta::SignatureDelta { signature } => {
|
||||
Some(ContentBlockDelta::ThinkingSignature(signature))
|
||||
}
|
||||
AnthropicContentBlockDelta::Unsupported => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(crate) struct AnthropicMessageDelta {
|
||||
pub(crate) stop_reason: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) usage: Option<AnthropicUsage>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(crate) struct AnthropicStreamError {
|
||||
#[serde(rename = "type")]
|
||||
pub(crate) kind: String,
|
||||
pub(crate) message: String,
|
||||
}
|
||||
|
||||
impl AnthropicStreamEvent {
|
||||
pub(crate) fn into_provider_events(
|
||||
self,
|
||||
reasoning_provenance: &ReasoningProvenance,
|
||||
) -> Result<Vec<ProviderEvent>, AnthropicStreamError> {
|
||||
match self {
|
||||
AnthropicStreamEvent::MessageStart { message } => {
|
||||
let usage = message
|
||||
.usage
|
||||
.clone()
|
||||
.and_then(AnthropicUsage::into_token_usage);
|
||||
let mut events = vec![ProviderEvent::MessageStarted {
|
||||
id: message.id,
|
||||
model: message.model,
|
||||
role: match message.role.as_str() {
|
||||
"user" => Role::User,
|
||||
"assistant" => Role::Assistant,
|
||||
_ => Role::Unknown(message.role),
|
||||
},
|
||||
}];
|
||||
if let Some(usage) = usage {
|
||||
events.push(ProviderEvent::MessageDelta {
|
||||
stop_reason: None,
|
||||
usage: Some(usage),
|
||||
});
|
||||
}
|
||||
Ok(events)
|
||||
}
|
||||
AnthropicStreamEvent::ContentBlockStart {
|
||||
index,
|
||||
content_block,
|
||||
} => Ok(content_block.into_provider_events(index, reasoning_provenance)),
|
||||
AnthropicStreamEvent::ContentBlockDelta { index, delta } => Ok(delta
|
||||
.into_provider_delta()
|
||||
.map(|delta| vec![ProviderEvent::ContentBlockDelta { index, delta }])
|
||||
.unwrap_or_default()),
|
||||
AnthropicStreamEvent::ContentBlockStop { index } => {
|
||||
Ok(vec![ProviderEvent::ContentBlockStopped { index }])
|
||||
}
|
||||
AnthropicStreamEvent::MessageDelta { delta } => Ok(vec![ProviderEvent::MessageDelta {
|
||||
stop_reason: delta.stop_reason,
|
||||
usage: delta.usage.and_then(AnthropicUsage::into_token_usage),
|
||||
}]),
|
||||
AnthropicStreamEvent::MessageStop => Ok(vec![ProviderEvent::MessageStopped]),
|
||||
AnthropicStreamEvent::Ping => Ok(Vec::new()),
|
||||
AnthropicStreamEvent::Error { error } => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
72
vendor/mentra-provider/src/auth.rs
vendored
Normal file
72
vendor/mentra-provider/src/auth.rs
vendored
Normal file
@@ -0,0 +1,72 @@
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::error::ProviderError;
|
||||
|
||||
/// What kind of auth material a provider expects.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub enum AuthScheme {
|
||||
#[default]
|
||||
None,
|
||||
BearerToken,
|
||||
Header {
|
||||
name: String,
|
||||
},
|
||||
QueryParam {
|
||||
name: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Provider credentials resolved at runtime.
|
||||
///
|
||||
/// `bearer_token` is the provider's primary auth secret and is applied according
|
||||
/// to the provider definition's [`AuthScheme`]. For example, Responses-family
|
||||
/// providers send it as `Authorization: Bearer ...`, while header/query auth
|
||||
/// providers can reuse the same resolved secret in a different location.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct ProviderCredentials {
|
||||
pub bearer_token: Option<String>,
|
||||
pub account_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub headers: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// Supplies credentials on demand for a provider instance.
|
||||
#[async_trait]
|
||||
pub trait CredentialSource: Send + Sync {
|
||||
async fn credentials(&self) -> Result<ProviderCredentials, ProviderError>;
|
||||
|
||||
async fn bearer_token(&self) -> Result<String, ProviderError> {
|
||||
self.credentials()
|
||||
.await?
|
||||
.bearer_token
|
||||
.ok_or_else(|| ProviderError::InvalidRequest("missing bearer token".to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Supplies a fixed auth secret to the provider.
|
||||
#[derive(Clone)]
|
||||
pub struct StaticCredentialSource {
|
||||
secret: Arc<str>,
|
||||
}
|
||||
|
||||
impl StaticCredentialSource {
|
||||
pub fn new(secret: impl Into<String>) -> Self {
|
||||
Self {
|
||||
secret: Arc::from(secret.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CredentialSource for StaticCredentialSource {
|
||||
async fn credentials(&self) -> Result<ProviderCredentials, ProviderError> {
|
||||
Ok(ProviderCredentials {
|
||||
bearer_token: Some(self.secret.to_string()),
|
||||
account_id: None,
|
||||
headers: HashMap::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
483
vendor/mentra-provider/src/definition.rs
vendored
Normal file
483
vendor/mentra-provider/src/definition.rs
vendored
Normal file
@@ -0,0 +1,483 @@
|
||||
use http::HeaderMap;
|
||||
use http::HeaderName;
|
||||
use http::HeaderValue;
|
||||
use http::header;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::borrow::Cow;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Display;
|
||||
use std::time::Duration;
|
||||
use strum::Display as StrumDisplay;
|
||||
use strum::IntoStaticStr;
|
||||
use url::Url;
|
||||
|
||||
use crate::request::SessionRequestOptions;
|
||||
|
||||
/// Builtin provider families Mentra can construct from presets.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, StrumDisplay, IntoStaticStr)]
|
||||
#[strum(serialize_all = "lowercase")]
|
||||
pub enum BuiltinProvider {
|
||||
Anthropic,
|
||||
Gemini,
|
||||
OpenAI,
|
||||
OpenRouter,
|
||||
Ollama,
|
||||
LmStudio,
|
||||
}
|
||||
|
||||
impl From<BuiltinProvider> for ProviderId {
|
||||
fn from(value: BuiltinProvider) -> Self {
|
||||
Self(Cow::Borrowed(value.into()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Stable identifier for a registered provider implementation.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
|
||||
pub struct ProviderId(Cow<'static, str>);
|
||||
|
||||
impl ProviderId {
|
||||
pub fn new(id: impl Into<String>) -> Self {
|
||||
Self(Cow::Owned(id.into()))
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
self.0.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for ProviderId {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for ProviderId {
|
||||
fn from(value: &str) -> Self {
|
||||
Self::new(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for ProviderId {
|
||||
fn from(value: String) -> Self {
|
||||
Self(Cow::Owned(value))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&String> for ProviderId {
|
||||
fn from(value: &String) -> Self {
|
||||
Self::new(value.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
/// Human-facing metadata about a provider.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ProviderDescriptor {
|
||||
pub id: ProviderId,
|
||||
pub display_name: Option<String>,
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
impl ProviderDescriptor {
|
||||
pub fn new(id: impl Into<ProviderId>) -> Self {
|
||||
Self {
|
||||
id: id.into(),
|
||||
display_name: None,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Capabilities advertised by a provider instance.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct ProviderCapabilities {
|
||||
pub supports_model_listing: bool,
|
||||
pub supports_streaming: bool,
|
||||
pub supports_websockets: bool,
|
||||
pub supports_tool_calls: bool,
|
||||
pub supports_images: bool,
|
||||
pub supports_history_compaction: bool,
|
||||
pub supports_memory_summarization: bool,
|
||||
pub supports_deferred_tools: bool,
|
||||
pub supports_hosted_tool_search: bool,
|
||||
pub supports_hosted_web_search: bool,
|
||||
pub supports_image_generation: bool,
|
||||
pub supports_reasoning_effort: bool,
|
||||
pub reports_reasoning_tokens: bool,
|
||||
pub reports_thoughts_tokens: bool,
|
||||
pub supports_structured_tool_results: bool,
|
||||
pub supports_embeddings: bool,
|
||||
}
|
||||
|
||||
/// Wire protocol supported by a provider.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum WireApi {
|
||||
#[default]
|
||||
Responses,
|
||||
AnthropicMessages,
|
||||
GeminiGenerateContent,
|
||||
}
|
||||
|
||||
impl Display for WireApi {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let value = match self {
|
||||
Self::Responses => "responses",
|
||||
Self::AnthropicMessages => "anthropic_messages",
|
||||
Self::GeminiGenerateContent => "gemini_generate_content",
|
||||
};
|
||||
f.write_str(value)
|
||||
}
|
||||
}
|
||||
|
||||
/// Retry configuration for provider transport calls.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RetryPolicy {
|
||||
pub max_attempts: u64,
|
||||
pub base_delay: Duration,
|
||||
pub retry_429: bool,
|
||||
pub retry_5xx: bool,
|
||||
pub retry_transport: bool,
|
||||
}
|
||||
|
||||
impl Default for RetryPolicy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_attempts: 5,
|
||||
base_delay: Duration::from_millis(200),
|
||||
retry_429: false,
|
||||
retry_5xx: true,
|
||||
retry_transport: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Serializable provider definition used by runtime and adapter layers.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ProviderDefinition {
|
||||
pub descriptor: ProviderDescriptor,
|
||||
#[serde(default)]
|
||||
pub wire_api: WireApi,
|
||||
#[serde(default)]
|
||||
pub auth_scheme: crate::AuthScheme,
|
||||
#[serde(default)]
|
||||
pub capabilities: ProviderCapabilities,
|
||||
pub base_url: Option<String>,
|
||||
#[serde(default)]
|
||||
pub query_params: Option<HashMap<String, String>>,
|
||||
#[serde(default)]
|
||||
pub headers: Option<HashMap<String, String>>,
|
||||
#[serde(default)]
|
||||
pub retry: RetryPolicy,
|
||||
#[serde(default = "default_stream_idle_timeout")]
|
||||
pub stream_idle_timeout: Duration,
|
||||
#[serde(default = "default_websocket_connect_timeout")]
|
||||
pub websocket_connect_timeout: Duration,
|
||||
}
|
||||
|
||||
fn default_stream_idle_timeout() -> Duration {
|
||||
Duration::from_millis(300_000)
|
||||
}
|
||||
|
||||
fn default_websocket_connect_timeout() -> Duration {
|
||||
Duration::from_millis(15_000)
|
||||
}
|
||||
|
||||
impl ProviderDefinition {
|
||||
pub fn new(id: impl Into<ProviderId>) -> Self {
|
||||
Self {
|
||||
descriptor: ProviderDescriptor::new(id),
|
||||
wire_api: WireApi::default(),
|
||||
auth_scheme: crate::AuthScheme::default(),
|
||||
capabilities: ProviderCapabilities {
|
||||
supports_model_listing: true,
|
||||
supports_streaming: true,
|
||||
supports_websockets: false,
|
||||
supports_tool_calls: true,
|
||||
supports_images: true,
|
||||
supports_history_compaction: false,
|
||||
supports_memory_summarization: false,
|
||||
supports_deferred_tools: false,
|
||||
supports_hosted_tool_search: false,
|
||||
supports_hosted_web_search: false,
|
||||
supports_image_generation: false,
|
||||
supports_reasoning_effort: false,
|
||||
reports_reasoning_tokens: false,
|
||||
reports_thoughts_tokens: false,
|
||||
supports_structured_tool_results: false,
|
||||
supports_embeddings: false,
|
||||
},
|
||||
base_url: None,
|
||||
query_params: None,
|
||||
headers: None,
|
||||
retry: RetryPolicy::default(),
|
||||
stream_idle_timeout: default_stream_idle_timeout(),
|
||||
websocket_connect_timeout: default_websocket_connect_timeout(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn descriptor(&self) -> ProviderDescriptor {
|
||||
self.descriptor.clone()
|
||||
}
|
||||
|
||||
pub fn provider_id(&self) -> &ProviderId {
|
||||
&self.descriptor.id
|
||||
}
|
||||
|
||||
pub fn url_for_path(&self, path: &str) -> String {
|
||||
let base = self
|
||||
.base_url
|
||||
.as_deref()
|
||||
.unwrap_or_default()
|
||||
.trim_end_matches('/');
|
||||
let path = path.trim_start_matches('/');
|
||||
let mut url = if path.is_empty() {
|
||||
base.to_string()
|
||||
} else {
|
||||
format!("{base}/{path}")
|
||||
};
|
||||
|
||||
if let Some(params) = self
|
||||
.query_params
|
||||
.as_ref()
|
||||
.filter(|params| !params.is_empty())
|
||||
{
|
||||
let qs = params
|
||||
.iter()
|
||||
.map(|(key, value)| format!("{key}={value}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("&");
|
||||
url.push('?');
|
||||
url.push_str(&qs);
|
||||
}
|
||||
|
||||
url
|
||||
}
|
||||
|
||||
pub fn build_headers(
|
||||
&self,
|
||||
credentials: &crate::ProviderCredentials,
|
||||
) -> Result<HeaderMap, crate::ProviderError> {
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
if let Some(configured_headers) = &self.headers {
|
||||
for (name, value) in configured_headers {
|
||||
insert_header(&mut headers, name, value)?;
|
||||
}
|
||||
}
|
||||
|
||||
for (name, value) in &credentials.headers {
|
||||
insert_header(&mut headers, name, value)?;
|
||||
}
|
||||
|
||||
match &self.auth_scheme {
|
||||
crate::AuthScheme::None | crate::AuthScheme::QueryParam { .. } => {}
|
||||
crate::AuthScheme::BearerToken => {
|
||||
let token = required_auth_value(credentials)?;
|
||||
let auth_value =
|
||||
HeaderValue::from_str(&format!("Bearer {token}")).map_err(|error| {
|
||||
crate::ProviderError::InvalidRequest(format!(
|
||||
"invalid bearer token header: {error}"
|
||||
))
|
||||
})?;
|
||||
headers.insert(header::AUTHORIZATION, auth_value);
|
||||
}
|
||||
crate::AuthScheme::Header { name } => {
|
||||
let token = required_auth_value(credentials)?;
|
||||
insert_header(&mut headers, name, token)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
pub fn build_headers_for_session(
|
||||
&self,
|
||||
credentials: &crate::ProviderCredentials,
|
||||
session: Option<&SessionRequestOptions>,
|
||||
fallback_turn_state: Option<&str>,
|
||||
) -> Result<HeaderMap, crate::ProviderError> {
|
||||
let mut headers = self.build_headers(credentials)?;
|
||||
|
||||
if let Some(value) = session
|
||||
.and_then(|session| session.sticky_turn_state.as_deref())
|
||||
.or(fallback_turn_state)
|
||||
.and_then(|turn_state| HeaderValue::from_str(turn_state).ok())
|
||||
{
|
||||
headers.insert("x-mentra-turn-state", value.clone());
|
||||
headers.insert("x-codex-turn-state", value);
|
||||
}
|
||||
if let Some(value) = session
|
||||
.and_then(|session| session.turn_metadata.as_deref())
|
||||
.and_then(|value| HeaderValue::from_str(value).ok())
|
||||
{
|
||||
headers.insert("x-mentra-turn-metadata", value.clone());
|
||||
headers.insert("x-codex-turn-metadata", value);
|
||||
}
|
||||
if let Some(value) = session
|
||||
.and_then(|session| session.session_affinity.as_deref())
|
||||
.and_then(|value| HeaderValue::from_str(value).ok())
|
||||
{
|
||||
headers.insert("x-mentra-session-affinity", value);
|
||||
}
|
||||
if let Some(prefer_connection_reuse) =
|
||||
session.and_then(|session| session.prefer_connection_reuse)
|
||||
{
|
||||
headers.insert(
|
||||
"x-mentra-connection-reuse",
|
||||
HeaderValue::from_static(if prefer_connection_reuse {
|
||||
"prefer-reuse"
|
||||
} else {
|
||||
"prefer-fresh"
|
||||
}),
|
||||
);
|
||||
}
|
||||
if let Some(value) = session
|
||||
.and_then(|session| session.subagent.as_deref())
|
||||
.and_then(|value| HeaderValue::from_str(value).ok())
|
||||
{
|
||||
headers.insert("x-openai-subagent", value);
|
||||
}
|
||||
if let Some(extra_headers) = session.map(|session| &session.extra_headers) {
|
||||
for (name, value) in extra_headers {
|
||||
if let (Ok(name), Ok(value)) = (
|
||||
name.parse::<http::HeaderName>(),
|
||||
HeaderValue::from_str(value),
|
||||
) {
|
||||
headers.insert(name, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
pub fn request_url_with_auth_for_path(
|
||||
&self,
|
||||
path: &str,
|
||||
credentials: &crate::ProviderCredentials,
|
||||
) -> Result<Url, crate::ProviderError> {
|
||||
let mut url = Url::parse(&self.url_for_path(path))
|
||||
.map_err(|error| crate::ProviderError::InvalidRequest(error.to_string()))?;
|
||||
|
||||
if let crate::AuthScheme::QueryParam { name } = &self.auth_scheme {
|
||||
let token = required_auth_value(credentials)?;
|
||||
url.query_pairs_mut().append_pair(name, token);
|
||||
}
|
||||
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
pub fn websocket_url_for_path(&self, path: &str) -> Result<Url, url::ParseError> {
|
||||
let mut url = Url::parse(&self.url_for_path(path))?;
|
||||
|
||||
let scheme = match url.scheme() {
|
||||
"http" => "ws",
|
||||
"https" => "wss",
|
||||
"ws" | "wss" => return Ok(url),
|
||||
_ => return Ok(url),
|
||||
};
|
||||
let _ = url.set_scheme(scheme);
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
pub fn websocket_url_with_auth_for_path(
|
||||
&self,
|
||||
path: &str,
|
||||
credentials: &crate::ProviderCredentials,
|
||||
) -> Result<Url, crate::ProviderError> {
|
||||
let mut url = self
|
||||
.websocket_url_for_path(path)
|
||||
.map_err(|error| crate::ProviderError::InvalidRequest(error.to_string()))?;
|
||||
|
||||
if let crate::AuthScheme::QueryParam { name } = &self.auth_scheme {
|
||||
let token = required_auth_value(credentials)?;
|
||||
url.query_pairs_mut().append_pair(name, token);
|
||||
}
|
||||
|
||||
Ok(url)
|
||||
}
|
||||
}
|
||||
|
||||
fn insert_header(
|
||||
headers: &mut HeaderMap,
|
||||
name: &str,
|
||||
value: &str,
|
||||
) -> Result<(), crate::ProviderError> {
|
||||
let header_name = HeaderName::from_bytes(name.as_bytes()).map_err(|error| {
|
||||
crate::ProviderError::InvalidRequest(format!(
|
||||
"invalid provider header name {name:?}: {error}"
|
||||
))
|
||||
})?;
|
||||
let header_value = HeaderValue::from_str(value).map_err(|error| {
|
||||
crate::ProviderError::InvalidRequest(format!(
|
||||
"invalid provider header value for {name:?}: {error}"
|
||||
))
|
||||
})?;
|
||||
headers.insert(header_name, header_value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn required_auth_value(
|
||||
credentials: &crate::ProviderCredentials,
|
||||
) -> Result<&str, crate::ProviderError> {
|
||||
credentials.bearer_token.as_deref().ok_or_else(|| {
|
||||
crate::ProviderError::InvalidRequest("missing provider auth credential".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn build_headers_applies_bearer_auth_and_static_headers() {
|
||||
let mut definition = ProviderDefinition::new("test");
|
||||
definition.auth_scheme = crate::AuthScheme::BearerToken;
|
||||
definition.headers = Some(HashMap::from([(
|
||||
"x-provider-header".to_string(),
|
||||
"static".to_string(),
|
||||
)]));
|
||||
|
||||
let headers = definition
|
||||
.build_headers(&crate::ProviderCredentials {
|
||||
bearer_token: Some("secret".to_string()),
|
||||
account_id: None,
|
||||
headers: HashMap::from([("x-runtime-header".to_string(), "dynamic".to_string())]),
|
||||
})
|
||||
.expect("headers should build");
|
||||
|
||||
assert_eq!(headers["x-provider-header"], "static");
|
||||
assert_eq!(headers["x-runtime-header"], "dynamic");
|
||||
assert_eq!(headers[header::AUTHORIZATION], "Bearer secret");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_url_with_auth_appends_query_param_auth() {
|
||||
let mut definition = ProviderDefinition::new("test");
|
||||
definition.base_url = Some("https://example.com/v1".to_string());
|
||||
definition.query_params = Some(HashMap::from([(
|
||||
"api-version".to_string(),
|
||||
"2026".to_string(),
|
||||
)]));
|
||||
definition.auth_scheme = crate::AuthScheme::QueryParam {
|
||||
name: "api-key".to_string(),
|
||||
};
|
||||
|
||||
let url = definition
|
||||
.request_url_with_auth_for_path(
|
||||
"responses",
|
||||
&crate::ProviderCredentials {
|
||||
bearer_token: Some("secret".to_string()),
|
||||
account_id: None,
|
||||
headers: HashMap::new(),
|
||||
},
|
||||
)
|
||||
.expect("url should build");
|
||||
|
||||
assert_eq!(
|
||||
url.as_str(),
|
||||
"https://example.com/v1/responses?api-version=2026&api-key=secret"
|
||||
);
|
||||
}
|
||||
}
|
||||
175
vendor/mentra-provider/src/embedding.rs
vendored
Normal file
175
vendor/mentra-provider/src/embedding.rs
vendored
Normal file
@@ -0,0 +1,175 @@
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::ProviderError;
|
||||
|
||||
/// Static metadata about an embedding model.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct EmbeddingModelInfo {
|
||||
pub id: String,
|
||||
pub dimensions: usize,
|
||||
pub max_tokens: usize,
|
||||
}
|
||||
|
||||
impl EmbeddingModelInfo {
|
||||
pub fn new(id: impl Into<String>, dimensions: usize, max_tokens: usize) -> Self {
|
||||
Self {
|
||||
id: id.into(),
|
||||
dimensions,
|
||||
max_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Input variants for an embedding request.
|
||||
///
|
||||
/// Only `Serialize` is derived — these types are sent in requests but never
|
||||
/// deserialized from responses, so `Deserialize` is not needed.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum EmbeddingInput<'a> {
|
||||
Single(&'a str),
|
||||
Batch(&'a [&'a str]),
|
||||
}
|
||||
|
||||
/// Request payload sent to the embeddings endpoint.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct EmbeddingRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub input: EmbeddingInput<'a>,
|
||||
}
|
||||
|
||||
impl<'a> EmbeddingRequest<'a> {
|
||||
pub fn single(model: &'a str, text: &'a str) -> Self {
|
||||
Self {
|
||||
model,
|
||||
input: EmbeddingInput::Single(text),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn batch(model: &'a str, texts: &'a [&'a str]) -> Self {
|
||||
Self {
|
||||
model,
|
||||
input: EmbeddingInput::Batch(texts),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A single embedding vector with its position in the batch.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct EmbeddingData {
|
||||
pub index: usize,
|
||||
pub embedding: Vec<f32>,
|
||||
}
|
||||
|
||||
/// Token usage reported by the embeddings endpoint.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct EmbeddingUsage {
|
||||
pub prompt_tokens: u32,
|
||||
pub total_tokens: u32,
|
||||
}
|
||||
|
||||
/// Response returned from the embeddings endpoint.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct EmbeddingResponse {
|
||||
pub data: Vec<EmbeddingData>,
|
||||
pub model: String,
|
||||
pub usage: EmbeddingUsage,
|
||||
}
|
||||
|
||||
/// Trait implemented by providers that support vector embeddings.
|
||||
///
|
||||
/// `EmbeddingProvider` is intentionally separate from `Provider` because not all
|
||||
/// LLM providers expose an embeddings endpoint (e.g. Anthropic and Gemini do not).
|
||||
#[async_trait]
|
||||
pub trait EmbeddingProvider: Send + Sync {
|
||||
/// Embed a single piece of text.
|
||||
async fn embed(&self, model: &str, text: &str) -> Result<Vec<f32>, ProviderError> {
|
||||
let texts = [text];
|
||||
let mut response = self.embed_batch(model, &texts).await?;
|
||||
response
|
||||
.data
|
||||
.pop()
|
||||
.map(|d| d.embedding)
|
||||
.ok_or_else(|| ProviderError::InvalidResponse("empty embedding response".to_string()))
|
||||
}
|
||||
|
||||
/// Embed a batch of texts, returning one vector per input in order.
|
||||
async fn embed_batch(
|
||||
&self,
|
||||
model: &str,
|
||||
texts: &[&str],
|
||||
) -> Result<EmbeddingResponse, ProviderError>;
|
||||
|
||||
/// Returns metadata about the embedding models available from this provider.
|
||||
fn embedding_models(&self) -> Vec<EmbeddingModelInfo>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn embedding_request_single_serializes_as_string() {
|
||||
let req = EmbeddingRequest::single("text-embedding-3-small", "hello world");
|
||||
let json = serde_json::to_value(&req).unwrap();
|
||||
|
||||
assert_eq!(json["model"], "text-embedding-3-small");
|
||||
assert_eq!(json["input"], "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embedding_request_batch_serializes_as_array() {
|
||||
let texts = ["hello", "world"];
|
||||
let req = EmbeddingRequest::batch("text-embedding-3-small", &texts);
|
||||
let json = serde_json::to_value(&req).unwrap();
|
||||
|
||||
assert_eq!(json["model"], "text-embedding-3-small");
|
||||
assert_eq!(json["input"], serde_json::json!(["hello", "world"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embedding_response_deserializes_correctly() {
|
||||
let raw = serde_json::json!({
|
||||
"data": [
|
||||
{ "index": 0, "embedding": [0.1, 0.2, 0.3] },
|
||||
{ "index": 1, "embedding": [0.4, 0.5, 0.6] }
|
||||
],
|
||||
"model": "text-embedding-3-small",
|
||||
"usage": { "prompt_tokens": 5, "total_tokens": 5 }
|
||||
});
|
||||
|
||||
let response: EmbeddingResponse = serde_json::from_value(raw).unwrap();
|
||||
|
||||
assert_eq!(response.model, "text-embedding-3-small");
|
||||
assert_eq!(response.data.len(), 2);
|
||||
assert_eq!(response.data[0].index, 0);
|
||||
assert_eq!(response.data[0].embedding, vec![0.1f32, 0.2, 0.3]);
|
||||
assert_eq!(response.data[1].index, 1);
|
||||
assert_eq!(response.usage.prompt_tokens, 5);
|
||||
assert_eq!(response.usage.total_tokens, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embedding_model_info_stores_fields() {
|
||||
let info = EmbeddingModelInfo::new("text-embedding-3-large", 3072, 8191);
|
||||
assert_eq!(info.id, "text-embedding-3-large");
|
||||
assert_eq!(info.dimensions, 3072);
|
||||
assert_eq!(info.max_tokens, 8191);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embedding_input_single_serializes_as_bare_string() {
|
||||
let input = EmbeddingInput::Single("test text");
|
||||
let serialized = serde_json::to_string(&input).unwrap();
|
||||
assert_eq!(serialized, r#""test text""#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embedding_input_batch_serializes_as_array() {
|
||||
let texts = ["a", "b"];
|
||||
let input = EmbeddingInput::Batch(&texts);
|
||||
let json = serde_json::to_value(&input).unwrap();
|
||||
assert_eq!(json, serde_json::json!(["a", "b"]));
|
||||
}
|
||||
}
|
||||
271
vendor/mentra-provider/src/error.rs
vendored
Normal file
271
vendor/mentra-provider/src/error.rs
vendored
Normal file
@@ -0,0 +1,271 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use thiserror::Error;
|
||||
use time::{OffsetDateTime, PrimitiveDateTime};
|
||||
|
||||
/// Errors returned by provider implementations and stream adapters.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ProviderError {
|
||||
#[error("provider transport error: {0}")]
|
||||
Transport(#[source] reqwest::Error),
|
||||
#[error("retryable provider error: {message}")]
|
||||
Retryable {
|
||||
message: String,
|
||||
delay: Option<Duration>,
|
||||
},
|
||||
#[error("provider does not support capability: {0}")]
|
||||
UnsupportedCapability(String),
|
||||
#[error("{message}", message = provider_http_error(.status, .body))]
|
||||
Http {
|
||||
status: reqwest::StatusCode,
|
||||
body: String,
|
||||
/// How long the server asked the caller to wait, from the response's
|
||||
/// `Retry-After` header, or `None` when it sent none.
|
||||
///
|
||||
/// A rate limit is the one failure whose recovery time the server
|
||||
/// knows and the client cannot guess: an exponential backoff shaped
|
||||
/// for a connection blip retries five times inside the window and
|
||||
/// then gives up while the limit is still in force. Capturing the
|
||||
/// header here, where the response is turned into an error, is what
|
||||
/// lets a caller wait the interval the server named instead —
|
||||
/// nothing further up the stack ever sees the headers.
|
||||
///
|
||||
/// Read it back through [`retry_after`](ProviderError::retry_after),
|
||||
/// which answers the same question for
|
||||
/// [`Retryable`](ProviderError::Retryable) too. Build one of these
|
||||
/// from a live response with
|
||||
/// [`from_http_response`](ProviderError::from_http_response) rather
|
||||
/// than filling the field in by hand.
|
||||
retry_after: Option<Duration>,
|
||||
},
|
||||
#[error("failed to decode provider response: {0}")]
|
||||
Decode(#[source] reqwest::Error),
|
||||
#[error("failed to serialize provider request: {0}")]
|
||||
Serialize(#[source] serde_json::Error),
|
||||
#[error("failed to deserialize provider payload: {0}")]
|
||||
Deserialize(#[source] serde_json::Error),
|
||||
#[error("invalid provider request: {0}")]
|
||||
InvalidRequest(String),
|
||||
#[error("invalid provider response: {0}")]
|
||||
InvalidResponse(String),
|
||||
#[error("malformed provider stream: {0}")]
|
||||
MalformedStream(String),
|
||||
}
|
||||
|
||||
impl ProviderError {
|
||||
/// Turns an unsuccessful HTTP response into an [`Http`](ProviderError::Http)
|
||||
/// error, reading `Retry-After` before the body is consumed.
|
||||
///
|
||||
/// This is the constructor every provider in this crate uses, and the one
|
||||
/// a custom provider should use: the status and the retry hint both come
|
||||
/// off the response, so neither can be forgotten at a call site.
|
||||
pub async fn from_http_response(response: reqwest::Response) -> Self {
|
||||
let status = response.status();
|
||||
let retry_after = retry_after_from_headers(response.headers());
|
||||
Self::Http {
|
||||
status,
|
||||
body: response.text().await.unwrap_or_default(),
|
||||
retry_after,
|
||||
}
|
||||
}
|
||||
|
||||
/// How long the provider asked the caller to wait before trying again, or
|
||||
/// `None` when it asked for nothing.
|
||||
///
|
||||
/// This is a request from the server, not a promise about the schedule: a
|
||||
/// caller decides what to do with it, including refusing an interval long
|
||||
/// enough to be an outage rather than a rate limit.
|
||||
pub fn retry_after(&self) -> Option<Duration> {
|
||||
match self {
|
||||
Self::Http { retry_after, .. } => *retry_after,
|
||||
Self::Retryable { delay, .. } => *delay,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn provider_http_error(status: &reqwest::StatusCode, body: &str) -> String {
|
||||
if body.trim().is_empty() {
|
||||
format!("provider returned HTTP {status}")
|
||||
} else {
|
||||
format!("provider returned HTTP {status}: {body}")
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads `Retry-After` off a response's headers, in whichever of its two forms
|
||||
/// the server chose.
|
||||
pub(crate) fn retry_after_from_headers(headers: &reqwest::header::HeaderMap) -> Option<Duration> {
|
||||
retry_after_from_header_value(headers.get(reqwest::header::RETRY_AFTER)?.to_str().ok()?)
|
||||
}
|
||||
|
||||
/// Reads one already-extracted `Retry-After` value, for transports that carry
|
||||
/// their headers as something other than a [`HeaderMap`](reqwest::header::HeaderMap).
|
||||
pub(crate) fn retry_after_from_header_value(value: &str) -> Option<Duration> {
|
||||
parse_retry_after(value, OffsetDateTime::now_utc())
|
||||
}
|
||||
|
||||
/// Parses a `Retry-After` value against a known `now`.
|
||||
///
|
||||
/// RFC 9110 allows two spellings — a count of seconds, and an HTTP-date — and
|
||||
/// providers use both, so a parser that understands only one silently ignores
|
||||
/// half the rate limits it is meant to honor. `now` is a parameter rather than
|
||||
/// read from the clock so the date form can be tested without waiting.
|
||||
///
|
||||
/// A date already in the past yields [`Duration::ZERO`] ("retry now") rather
|
||||
/// than nothing: the server did answer, and the answer was that the wait is
|
||||
/// over. Only the IMF-fixdate spelling that RFC 9110 requires senders to emit
|
||||
/// is understood; the two obsolete formats it only requires recipients to
|
||||
/// tolerate are read as no hint at all, which costs nothing but the schedule's
|
||||
/// own delay.
|
||||
fn parse_retry_after(value: &str, now: OffsetDateTime) -> Option<Duration> {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Ok(seconds) = value.parse::<u64>() {
|
||||
return Some(Duration::from_secs(seconds));
|
||||
}
|
||||
|
||||
let deadline = parse_http_date(value)?;
|
||||
Some((deadline - now).try_into().unwrap_or(Duration::ZERO))
|
||||
}
|
||||
|
||||
/// Parses an IMF-fixdate, the `Sun, 06 Nov 1994 08:49:37 GMT` spelling.
|
||||
fn parse_http_date(value: &str) -> Option<OffsetDateTime> {
|
||||
// `parse_borrowed` rather than `parse`: the latter is deprecated from
|
||||
// time 0.3.55 and a downstream resolving a newer `time` would see the
|
||||
// warning in this crate. Version 2 of the description syntax, which the
|
||||
// spelling below is already written in.
|
||||
let format = time::format_description::parse_borrowed::<2>(
|
||||
"[weekday repr:short], [day] [month repr:short] [year] [hour]:[minute]:[second] GMT",
|
||||
)
|
||||
.ok()?;
|
||||
PrimitiveDateTime::parse(value, format.as_slice())
|
||||
.ok()
|
||||
.map(PrimitiveDateTime::assume_utc)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The instant in RFC 9110's own `Retry-After` example, built without the
|
||||
/// `time` macros feature so the crate's dependency set stays as it is.
|
||||
fn now() -> OffsetDateTime {
|
||||
time::Date::from_calendar_date(1994, time::Month::November, 6)
|
||||
.expect("a real date")
|
||||
.with_hms(8, 49, 37)
|
||||
.expect("a real time")
|
||||
.assume_utc()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_delay_in_seconds_is_read_as_seconds() {
|
||||
assert_eq!(
|
||||
parse_retry_after("30", now()),
|
||||
Some(Duration::from_secs(30))
|
||||
);
|
||||
assert_eq!(
|
||||
parse_retry_after(" 30 ", now()),
|
||||
Some(Duration::from_secs(30)),
|
||||
"surrounding whitespace is not part of the value"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_http_date_is_read_as_the_wait_until_it() {
|
||||
// The other spelling RFC 9110 allows. A parser that understood only
|
||||
// seconds would return None here and fall back to its own schedule,
|
||||
// which is exactly the rate limit it was supposed to wait out.
|
||||
assert_eq!(
|
||||
parse_retry_after("Sun, 06 Nov 1994 08:50:37 GMT", now()),
|
||||
Some(Duration::from_secs(60))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_http_date_that_has_passed_means_retry_now() {
|
||||
assert_eq!(
|
||||
parse_retry_after("Sun, 06 Nov 1994 08:49:00 GMT", now()),
|
||||
Some(Duration::ZERO),
|
||||
"the server answered; the answer is that the wait is over"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unparseable_value_is_no_hint_at_all() {
|
||||
assert_eq!(parse_retry_after("", now()), None);
|
||||
assert_eq!(parse_retry_after("soon", now()), None);
|
||||
assert_eq!(parse_retry_after("-5", now()), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_rate_limited_response_reports_the_interval_it_named() {
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
headers.insert(reqwest::header::RETRY_AFTER, "42".parse().expect("valid"));
|
||||
|
||||
assert_eq!(
|
||||
retry_after_from_headers(&headers),
|
||||
Some(Duration::from_secs(42))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_429_response_becomes_an_error_that_still_knows_the_interval() {
|
||||
// The whole point of reading the header at construction time: nothing
|
||||
// above this call ever sees the response, so a hint not captured here
|
||||
// is a hint lost.
|
||||
let response = http::Response::builder()
|
||||
.status(429)
|
||||
.header("retry-after", "60")
|
||||
.body("rate limit exceeded")
|
||||
.expect("a response");
|
||||
|
||||
let error = ProviderError::from_http_response(reqwest::Response::from(response)).await;
|
||||
|
||||
let ProviderError::Http {
|
||||
status,
|
||||
body,
|
||||
retry_after,
|
||||
} = error
|
||||
else {
|
||||
panic!("an unsuccessful response is an Http error");
|
||||
};
|
||||
assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(body, "rate limit exceeded");
|
||||
assert_eq!(retry_after, Some(Duration::from_secs(60)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_response_without_the_header_asks_for_nothing() {
|
||||
let response = http::Response::builder()
|
||||
.status(503)
|
||||
.body("upstream is restarting")
|
||||
.expect("a response");
|
||||
|
||||
let error = ProviderError::from_http_response(reqwest::Response::from(response)).await;
|
||||
|
||||
assert_eq!(error.retry_after(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn both_retryable_shapes_answer_the_same_question() {
|
||||
// A caller shaping a backoff asks one question and must not have to
|
||||
// know which variant carried the answer.
|
||||
let http = ProviderError::Http {
|
||||
status: reqwest::StatusCode::TOO_MANY_REQUESTS,
|
||||
body: String::new(),
|
||||
retry_after: Some(Duration::from_secs(20)),
|
||||
};
|
||||
let retryable = ProviderError::Retryable {
|
||||
message: "connection closed".to_string(),
|
||||
delay: Some(Duration::from_millis(750)),
|
||||
};
|
||||
let silent = ProviderError::InvalidRequest("bad model".to_string());
|
||||
|
||||
assert_eq!(http.retry_after(), Some(Duration::from_secs(20)));
|
||||
assert_eq!(retryable.retry_after(), Some(Duration::from_millis(750)));
|
||||
assert_eq!(silent.retry_after(), None);
|
||||
}
|
||||
}
|
||||
266
vendor/mentra-provider/src/gemini.rs
vendored
Normal file
266
vendor/mentra-provider/src/gemini.rs
vendored
Normal file
@@ -0,0 +1,266 @@
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub(crate) mod model;
|
||||
pub(crate) mod sse;
|
||||
|
||||
use crate::AuthScheme;
|
||||
use crate::BuiltinProvider;
|
||||
use crate::CompactionRequest;
|
||||
use crate::CompactionResponse;
|
||||
use crate::CredentialSource;
|
||||
use crate::ModelCatalog;
|
||||
use crate::ModelInfo;
|
||||
use crate::ProviderCapabilities;
|
||||
use crate::ProviderDefinition;
|
||||
use crate::ProviderError;
|
||||
use crate::ProviderEventStream;
|
||||
use crate::ProviderSession;
|
||||
use crate::ProviderSessionFactory;
|
||||
use crate::RegisteredProvider;
|
||||
use crate::Request;
|
||||
use crate::StaticCredentialSource;
|
||||
use crate::WireApi;
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/";
|
||||
|
||||
pub struct GeminiProvider<C = StaticCredentialSource> {
|
||||
client: reqwest::Client,
|
||||
credential_source: Arc<C>,
|
||||
definition: ProviderDefinition,
|
||||
}
|
||||
|
||||
impl<C> Clone for GeminiProvider<C> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
client: self.client.clone(),
|
||||
credential_source: Arc::clone(&self.credential_source),
|
||||
definition: self.definition.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GeminiProvider<StaticCredentialSource> {
|
||||
pub fn new(api_key: impl Into<String>) -> Self {
|
||||
Self::with_credential_source(StaticCredentialSource::new(api_key))
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> GeminiProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
pub fn with_credential_source(credential_source: C) -> Self {
|
||||
Self::with_shared_credential_source(Arc::new(credential_source))
|
||||
}
|
||||
|
||||
pub fn with_shared_credential_source(credential_source: Arc<C>) -> Self {
|
||||
Self::with_definition_and_shared_credential_source(Self::definition(), credential_source)
|
||||
}
|
||||
|
||||
pub fn with_definition_and_credential_source(
|
||||
definition: ProviderDefinition,
|
||||
credential_source: C,
|
||||
) -> Self {
|
||||
Self::with_definition_and_shared_credential_source(definition, Arc::new(credential_source))
|
||||
}
|
||||
|
||||
pub fn with_definition_and_shared_credential_source(
|
||||
definition: ProviderDefinition,
|
||||
credential_source: Arc<C>,
|
||||
) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.build()
|
||||
.expect("Failed to build client");
|
||||
|
||||
Self {
|
||||
client,
|
||||
credential_source,
|
||||
definition,
|
||||
}
|
||||
}
|
||||
|
||||
fn definition() -> ProviderDefinition {
|
||||
let mut definition = ProviderDefinition::new(BuiltinProvider::Gemini);
|
||||
definition.descriptor.display_name = Some("Gemini".to_string());
|
||||
definition.descriptor.description =
|
||||
Some("Google Gemini Developer API provider".to_string());
|
||||
definition.wire_api = WireApi::GeminiGenerateContent;
|
||||
definition.auth_scheme = AuthScheme::Header {
|
||||
name: "x-goog-api-key".to_string(),
|
||||
};
|
||||
definition.capabilities = ProviderCapabilities {
|
||||
supports_model_listing: true,
|
||||
supports_streaming: true,
|
||||
supports_websockets: false,
|
||||
supports_tool_calls: true,
|
||||
supports_images: true,
|
||||
supports_history_compaction: true,
|
||||
supports_memory_summarization: true,
|
||||
supports_deferred_tools: false,
|
||||
supports_hosted_tool_search: false,
|
||||
supports_hosted_web_search: false,
|
||||
supports_image_generation: false,
|
||||
supports_reasoning_effort: true,
|
||||
reports_reasoning_tokens: false,
|
||||
reports_thoughts_tokens: true,
|
||||
supports_structured_tool_results: false,
|
||||
supports_embeddings: false,
|
||||
};
|
||||
definition.base_url = Some(DEFAULT_BASE_URL.to_string());
|
||||
definition
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ModelCatalog for GeminiProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
let mut models = Vec::new();
|
||||
let mut page_token = None::<String>;
|
||||
|
||||
loop {
|
||||
let credentials = self.credential_source.credentials().await?;
|
||||
let mut request = self
|
||||
.client
|
||||
.get(
|
||||
self.definition
|
||||
.request_url_with_auth_for_path("v1beta/models", &credentials)?,
|
||||
)
|
||||
.headers(self.definition.build_headers(&credentials)?)
|
||||
.query(&[("pageSize", "1000")]);
|
||||
|
||||
if let Some(token) = page_token.as_deref() {
|
||||
request = request.query(&[("pageToken", token)]);
|
||||
}
|
||||
|
||||
let response = request.send().await.map_err(ProviderError::Transport)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::from_http_response(response).await);
|
||||
}
|
||||
|
||||
let page = response
|
||||
.json::<model::GeminiModelsPage>()
|
||||
.await
|
||||
.map_err(ProviderError::Decode)?;
|
||||
|
||||
models.extend(
|
||||
page.models
|
||||
.into_iter()
|
||||
.filter(|model| model.supports_generate_content())
|
||||
.map(ModelInfo::from),
|
||||
);
|
||||
|
||||
page_token = page.next_page_token;
|
||||
if page_token.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(models)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ProviderSessionFactory for GeminiProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn create_session(&self) -> Result<Box<dyn ProviderSession>, ProviderError> {
|
||||
Ok(Box::new((*self).clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ProviderSession for GeminiProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
|
||||
let session = request.provider_request_options.session.clone();
|
||||
let model_name = request.model.to_string();
|
||||
let request = model::GeminiGenerateContentRequest::try_from(request)?;
|
||||
let credentials = self.credential_source.credentials().await?;
|
||||
let response = self
|
||||
.client
|
||||
.post(self.definition.request_url_with_auth_for_path(
|
||||
&format!(
|
||||
"v1beta/{}:streamGenerateContent",
|
||||
normalize_model_name(&model_name)
|
||||
),
|
||||
&credentials,
|
||||
)?)
|
||||
.headers(self.definition.build_headers_for_session(
|
||||
&credentials,
|
||||
Some(&session),
|
||||
None,
|
||||
)?)
|
||||
.query(&[("alt", "sse")])
|
||||
.json(&request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(ProviderError::Transport)?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::from_http_response(response).await);
|
||||
}
|
||||
|
||||
Ok(sse::spawn_event_stream(response, model_name))
|
||||
}
|
||||
|
||||
async fn compact(
|
||||
&self,
|
||||
request: CompactionRequest<'_>,
|
||||
) -> Result<CompactionResponse, ProviderError> {
|
||||
let request = request.into_model_request()?;
|
||||
let response = ProviderSession::send(self, request).await?;
|
||||
Ok(response.into_compaction_response())
|
||||
}
|
||||
|
||||
async fn summarize_memories(
|
||||
&self,
|
||||
request: crate::MemorySummarizeRequest<'_>,
|
||||
) -> Result<crate::MemorySummarizeResponse, ProviderError> {
|
||||
let request = request.into_model_request()?;
|
||||
let response = ProviderSession::send(self, request).await?;
|
||||
response.into_memory_summarize_response()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> RegisteredProvider for GeminiProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
fn definition(&self) -> ProviderDefinition {
|
||||
self.definition.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_model_name(model: &str) -> String {
|
||||
if model.starts_with("models/") {
|
||||
model.to_string()
|
||||
} else {
|
||||
format!("models/{model}")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::RegisteredProvider;
|
||||
|
||||
#[test]
|
||||
fn definition_advertises_history_compaction_support() {
|
||||
let provider = GeminiProvider::new("test-key");
|
||||
|
||||
assert!(
|
||||
provider
|
||||
.definition()
|
||||
.capabilities
|
||||
.supports_history_compaction
|
||||
);
|
||||
}
|
||||
}
|
||||
945
vendor/mentra-provider/src/gemini/model.rs
vendored
Normal file
945
vendor/mentra-provider/src/gemini/model.rs
vendored
Normal file
@@ -0,0 +1,945 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::{
|
||||
BuiltinProvider, ContentBlock, ImageSource, Message, ModelInfo, ProviderError,
|
||||
ProviderToolKind, ReasoningEffort, Request, Role, ToolChoice, ToolLoadingPolicy,
|
||||
ToolSearchMode, ToolSpec,
|
||||
};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(crate) struct GeminiModelsPage {
|
||||
#[serde(default)]
|
||||
pub(crate) models: Vec<GeminiModel>,
|
||||
#[serde(default, rename = "nextPageToken", alias = "next_page_token")]
|
||||
pub(crate) next_page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub(crate) struct GeminiModel {
|
||||
pub(crate) name: String,
|
||||
#[serde(default, rename = "baseModelId", alias = "base_model_id")]
|
||||
pub(crate) base_model_id: Option<String>,
|
||||
#[serde(default, rename = "displayName", alias = "display_name")]
|
||||
pub(crate) display_name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) description: Option<String>,
|
||||
#[serde(
|
||||
default,
|
||||
rename = "supportedGenerationMethods",
|
||||
alias = "supported_generation_methods"
|
||||
)]
|
||||
supported_generation_methods: Vec<String>,
|
||||
}
|
||||
|
||||
impl GeminiModel {
|
||||
pub(crate) fn supports_generate_content(&self) -> bool {
|
||||
self.supported_generation_methods
|
||||
.iter()
|
||||
.any(|method| matches!(method.as_str(), "generateContent" | "streamGenerateContent"))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<GeminiModel> for ModelInfo {
|
||||
fn from(model: GeminiModel) -> Self {
|
||||
let id = model.base_model_id.unwrap_or_else(|| {
|
||||
model
|
||||
.name
|
||||
.strip_prefix("models/")
|
||||
.unwrap_or(&model.name)
|
||||
.to_string()
|
||||
});
|
||||
|
||||
ModelInfo {
|
||||
id,
|
||||
provider: BuiltinProvider::Gemini.into(),
|
||||
display_name: model.display_name,
|
||||
description: model.description,
|
||||
created_at: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(crate) struct GeminiGenerateContentRequest {
|
||||
#[serde(rename = "systemInstruction", skip_serializing_if = "Option::is_none")]
|
||||
system_instruction: Option<GeminiInstruction>,
|
||||
contents: Vec<GeminiContent>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
tools: Vec<GeminiTool>,
|
||||
#[serde(rename = "toolConfig", skip_serializing_if = "Option::is_none")]
|
||||
tool_config: Option<GeminiToolConfig>,
|
||||
#[serde(rename = "generationConfig", skip_serializing_if = "Option::is_none")]
|
||||
generation_config: Option<GeminiGenerationConfig>,
|
||||
}
|
||||
|
||||
impl<'a> TryFrom<Request<'a>> for GeminiGenerateContentRequest {
|
||||
type Error = ProviderError;
|
||||
|
||||
fn try_from(value: Request<'a>) -> Result<Self, Self::Error> {
|
||||
let generation_config = GeminiGenerationConfig::from_request(&value)?;
|
||||
let tool_name_by_id = collect_tool_name_by_id(value.messages.as_ref());
|
||||
let contents = value
|
||||
.messages
|
||||
.iter()
|
||||
.map(|message| GeminiContent::try_from_message(message, &tool_name_by_id))
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
.into_iter()
|
||||
.filter(|content| !content.parts.is_empty())
|
||||
.collect::<Vec<_>>();
|
||||
validate_gemini_tools(
|
||||
value.tools.as_ref(),
|
||||
value.tool_choice.as_ref(),
|
||||
value.provider_request_options.tool_search_mode,
|
||||
)?;
|
||||
let tools = if value.tools.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
vec![GeminiTool {
|
||||
function_declarations: value
|
||||
.tools
|
||||
.iter()
|
||||
.map(GeminiFunctionDeclaration::from)
|
||||
.collect(),
|
||||
}]
|
||||
};
|
||||
|
||||
Ok(GeminiGenerateContentRequest {
|
||||
system_instruction: value.system.map(|system| GeminiInstruction {
|
||||
parts: vec![GeminiPart::Text {
|
||||
text: system.into_owned(),
|
||||
}],
|
||||
}),
|
||||
contents,
|
||||
tool_config: value
|
||||
.tool_choice
|
||||
.filter(|_| !tools.is_empty())
|
||||
.map(Into::into),
|
||||
tools,
|
||||
generation_config,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_gemini_tools(
|
||||
tools: &[ToolSpec],
|
||||
tool_choice: Option<&ToolChoice>,
|
||||
tool_search_mode: ToolSearchMode,
|
||||
) -> Result<(), ProviderError> {
|
||||
if let Some(tool) = tools
|
||||
.iter()
|
||||
.find(|tool| tool.kind != ProviderToolKind::Function)
|
||||
{
|
||||
return Err(ProviderError::InvalidRequest(format!(
|
||||
"Gemini does not support provider tool kind {:?} for '{}'",
|
||||
tool.kind, tool.name
|
||||
)));
|
||||
}
|
||||
|
||||
let forced_tool_name = match tool_choice {
|
||||
Some(ToolChoice::Tool { name }) => Some(name.as_str()),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let has_deferred_tools = tools.iter().any(|tool| {
|
||||
tool.loading_policy == ToolLoadingPolicy::Deferred
|
||||
&& forced_tool_name != Some(tool.name.as_str())
|
||||
});
|
||||
|
||||
if !has_deferred_tools {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let message = match tool_search_mode {
|
||||
ToolSearchMode::Hosted => {
|
||||
"Gemini does not support hosted tool search for deferred custom tools"
|
||||
}
|
||||
ToolSearchMode::Disabled => {
|
||||
"Gemini does not support deferred custom tools without hosted tool search"
|
||||
}
|
||||
};
|
||||
|
||||
Err(ProviderError::InvalidRequest(message.to_string()))
|
||||
}
|
||||
|
||||
fn collect_tool_name_by_id(messages: &[Message]) -> BTreeMap<String, String> {
|
||||
let mut names = BTreeMap::new();
|
||||
|
||||
for message in messages {
|
||||
for block in &message.content {
|
||||
if let ContentBlock::ToolUse { id, name, .. } = block {
|
||||
names.insert(id.clone(), name.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
names
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiInstruction {
|
||||
parts: Vec<GeminiPart>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiContent {
|
||||
role: String,
|
||||
parts: Vec<GeminiPart>,
|
||||
}
|
||||
|
||||
impl GeminiContent {
|
||||
fn try_from_message(
|
||||
message: &Message,
|
||||
tool_name_by_id: &BTreeMap<String, String>,
|
||||
) -> Result<Self, ProviderError> {
|
||||
let role = match &message.role {
|
||||
Role::User | Role::Assistant => message.role.to_string(),
|
||||
Role::Unknown(role) => {
|
||||
return Err(ProviderError::InvalidRequest(format!(
|
||||
"Gemini message role '{role}' is not supported"
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
let mut parts = Vec::with_capacity(message.content.len());
|
||||
for block in &message.content {
|
||||
parts.push(GeminiPart::try_from_block(
|
||||
block,
|
||||
&message.role,
|
||||
tool_name_by_id,
|
||||
)?);
|
||||
}
|
||||
|
||||
Ok(GeminiContent { role, parts })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(untagged)]
|
||||
enum GeminiPart {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
InlineData {
|
||||
#[serde(rename = "inlineData")]
|
||||
inline_data: GeminiInlineData,
|
||||
},
|
||||
FunctionCall {
|
||||
#[serde(rename = "functionCall")]
|
||||
function_call: GeminiFunctionCall,
|
||||
},
|
||||
FunctionResponse {
|
||||
#[serde(rename = "functionResponse")]
|
||||
function_response: GeminiFunctionResponse,
|
||||
},
|
||||
}
|
||||
|
||||
impl GeminiPart {
|
||||
fn try_from_block(
|
||||
block: &ContentBlock,
|
||||
role: &Role,
|
||||
tool_name_by_id: &BTreeMap<String, String>,
|
||||
) -> Result<Self, ProviderError> {
|
||||
match block {
|
||||
ContentBlock::Text { text } => Ok(GeminiPart::Text { text: text.clone() }),
|
||||
ContentBlock::Thinking { .. } => Ok(GeminiPart::Text {
|
||||
text: block
|
||||
.thinking_fallback_text()
|
||||
.expect("thinking block has fallback text"),
|
||||
}),
|
||||
ContentBlock::Image { source } => {
|
||||
if !matches!(role, Role::User) {
|
||||
return Err(ProviderError::InvalidRequest(
|
||||
"Gemini image inputs are only supported in user messages".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
match source {
|
||||
ImageSource::Bytes { media_type, data } => Ok(GeminiPart::InlineData {
|
||||
inline_data: GeminiInlineData {
|
||||
mime_type: media_type.clone(),
|
||||
data: STANDARD.encode(data),
|
||||
},
|
||||
}),
|
||||
ImageSource::Url { .. } => Err(ProviderError::InvalidRequest(
|
||||
"Gemini image URL inputs are not supported without a file upload flow"
|
||||
.to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
ContentBlock::ToolUse { name, input, .. } => Ok(GeminiPart::FunctionCall {
|
||||
function_call: GeminiFunctionCall {
|
||||
name: name.clone(),
|
||||
args: input.clone(),
|
||||
},
|
||||
}),
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
} => {
|
||||
let name = tool_name_by_id.get(tool_use_id).cloned().ok_or_else(|| {
|
||||
ProviderError::InvalidRequest(format!(
|
||||
"Gemini tool result references unknown tool_use_id '{tool_use_id}'"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(GeminiPart::FunctionResponse {
|
||||
function_response: GeminiFunctionResponse {
|
||||
name,
|
||||
response: json!({
|
||||
"content": content.to_display_string(),
|
||||
"is_error": is_error,
|
||||
}),
|
||||
},
|
||||
})
|
||||
}
|
||||
ContentBlock::HostedToolSearch { call } => Ok(GeminiPart::FunctionCall {
|
||||
function_call: GeminiFunctionCall {
|
||||
name: "tool_search".to_string(),
|
||||
args: json!({ "query": call.query }),
|
||||
},
|
||||
}),
|
||||
ContentBlock::HostedWebSearch { call } => Ok(GeminiPart::FunctionCall {
|
||||
function_call: GeminiFunctionCall {
|
||||
name: "web_search".to_string(),
|
||||
args: serde_json::to_value(call.action.clone()).unwrap_or(Value::Null),
|
||||
},
|
||||
}),
|
||||
ContentBlock::ImageGeneration { call } => Ok(GeminiPart::FunctionCall {
|
||||
function_call: GeminiFunctionCall {
|
||||
name: "image_generation".to_string(),
|
||||
args: json!({
|
||||
"status": call.status,
|
||||
"revised_prompt": call.revised_prompt,
|
||||
}),
|
||||
},
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiInlineData {
|
||||
#[serde(rename = "mimeType")]
|
||||
mime_type: String,
|
||||
data: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiFunctionCall {
|
||||
name: String,
|
||||
args: Value,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiFunctionResponse {
|
||||
name: String,
|
||||
response: Value,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiTool {
|
||||
#[serde(rename = "functionDeclarations")]
|
||||
function_declarations: Vec<GeminiFunctionDeclaration>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiFunctionDeclaration {
|
||||
name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
description: Option<String>,
|
||||
parameters: Value,
|
||||
}
|
||||
|
||||
impl From<&ToolSpec> for GeminiFunctionDeclaration {
|
||||
fn from(tool: &ToolSpec) -> Self {
|
||||
GeminiFunctionDeclaration {
|
||||
name: tool.name.clone(),
|
||||
description: tool.description.clone(),
|
||||
parameters: tool.input_schema.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiToolConfig {
|
||||
#[serde(rename = "functionCallingConfig")]
|
||||
function_calling_config: GeminiFunctionCallingConfig,
|
||||
}
|
||||
|
||||
impl From<ToolChoice> for GeminiToolConfig {
|
||||
fn from(choice: ToolChoice) -> Self {
|
||||
let function_calling_config = match choice {
|
||||
ToolChoice::Auto => GeminiFunctionCallingConfig {
|
||||
mode: GeminiFunctionCallingMode::Auto,
|
||||
allowed_function_names: Vec::new(),
|
||||
},
|
||||
ToolChoice::Any => GeminiFunctionCallingConfig {
|
||||
mode: GeminiFunctionCallingMode::Any,
|
||||
allowed_function_names: Vec::new(),
|
||||
},
|
||||
ToolChoice::Tool { name } => GeminiFunctionCallingConfig {
|
||||
mode: GeminiFunctionCallingMode::Any,
|
||||
allowed_function_names: vec![name],
|
||||
},
|
||||
};
|
||||
|
||||
GeminiToolConfig {
|
||||
function_calling_config,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiFunctionCallingConfig {
|
||||
mode: GeminiFunctionCallingMode,
|
||||
#[serde(rename = "allowedFunctionNames", skip_serializing_if = "Vec::is_empty")]
|
||||
allowed_function_names: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
enum GeminiFunctionCallingMode {
|
||||
#[serde(rename = "AUTO")]
|
||||
Auto,
|
||||
#[serde(rename = "ANY")]
|
||||
Any,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiGenerationConfig {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
temperature: Option<f32>,
|
||||
#[serde(rename = "maxOutputTokens", skip_serializing_if = "Option::is_none")]
|
||||
max_output_tokens: Option<u32>,
|
||||
#[serde(rename = "thinkingConfig", skip_serializing_if = "Option::is_none")]
|
||||
thinking_config: Option<GeminiThinkingConfig>,
|
||||
}
|
||||
|
||||
impl GeminiGenerationConfig {
|
||||
fn from_request(request: &Request<'_>) -> Result<Option<Self>, ProviderError> {
|
||||
let thinking_config =
|
||||
if let Some(reasoning) = request.provider_request_options.reasoning.as_ref() {
|
||||
let Some(effort) = reasoning.effort else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !supports_gemini_thinking_level(&request.model) {
|
||||
return Err(ProviderError::InvalidRequest(format!(
|
||||
"Gemini reasoning effort requires a Gemini 3 model, got '{}'",
|
||||
request.model
|
||||
)));
|
||||
}
|
||||
|
||||
Some(GeminiThinkingConfig {
|
||||
thinking_level: effort.try_into()?,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let config = GeminiGenerationConfig {
|
||||
temperature: request.temperature,
|
||||
max_output_tokens: request.max_output_tokens,
|
||||
thinking_config,
|
||||
};
|
||||
|
||||
Ok((!config.is_empty()).then_some(config))
|
||||
}
|
||||
|
||||
fn is_empty(&self) -> bool {
|
||||
self.temperature.is_none()
|
||||
&& self.max_output_tokens.is_none()
|
||||
&& self.thinking_config.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GeminiThinkingConfig {
|
||||
#[serde(rename = "thinkingLevel")]
|
||||
thinking_level: GeminiThinkingLevel,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum GeminiThinkingLevel {
|
||||
Low,
|
||||
Medium,
|
||||
High,
|
||||
}
|
||||
|
||||
impl TryFrom<ReasoningEffort> for GeminiThinkingLevel {
|
||||
type Error = ProviderError;
|
||||
|
||||
fn try_from(value: ReasoningEffort) -> Result<Self, Self::Error> {
|
||||
match value {
|
||||
ReasoningEffort::Low => Ok(Self::Low),
|
||||
ReasoningEffort::Medium => Ok(Self::Medium),
|
||||
ReasoningEffort::High => Ok(Self::High),
|
||||
ReasoningEffort::XHigh => Err(ProviderError::InvalidRequest(
|
||||
"Gemini does not support reasoning effort 'xhigh'".to_string(),
|
||||
)),
|
||||
ReasoningEffort::Max => Err(ProviderError::InvalidRequest(
|
||||
"Gemini does not support reasoning effort 'max'".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn supports_gemini_thinking_level(model: &str) -> bool {
|
||||
let model = model.strip_prefix("models/").unwrap_or(model);
|
||||
model.starts_with("gemini-3")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{borrow::Cow, collections::BTreeMap};
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use crate::{
|
||||
BuiltinProvider, ContentBlock, Message, ModelInfo, ProviderError, ProviderRequestOptions,
|
||||
ReasoningEffort, ReasoningOptions, Request, Role, ToolChoice, ToolLoadingPolicy,
|
||||
ToolResultContent, ToolSearchMode, ToolSpec,
|
||||
};
|
||||
|
||||
use super::{GeminiGenerateContentRequest, GeminiModel};
|
||||
|
||||
#[test]
|
||||
fn converts_model_name_to_base_model_id() {
|
||||
let model = GeminiModel {
|
||||
name: "models/gemini-3-flash".to_string(),
|
||||
base_model_id: Some("gemini-3-flash".to_string()),
|
||||
display_name: Some("Gemini 3 Flash".to_string()),
|
||||
description: Some("Test".to_string()),
|
||||
supported_generation_methods: vec!["generateContent".to_string()],
|
||||
};
|
||||
|
||||
let info = ModelInfo::from(model);
|
||||
|
||||
assert_eq!(info.id, "gemini-3-flash");
|
||||
assert_eq!(info.provider, BuiltinProvider::Gemini.into());
|
||||
assert_eq!(info.display_name.as_deref(), Some("Gemini 3 Flash"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_request_to_gemini_payload() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: Some(Cow::Borrowed("Be helpful.")),
|
||||
messages: Cow::Owned(vec![
|
||||
Message::user(ContentBlock::text("What files changed?")),
|
||||
Message::assistant(ContentBlock::ToolUse {
|
||||
id: "call_1".to_string(),
|
||||
name: "files".to_string(),
|
||||
input: json!({ "operations": [{ "op": "read", "path": "README.md" }] }),
|
||||
}),
|
||||
Message::user(ContentBlock::ToolResult {
|
||||
tool_use_id: "call_1".to_string(),
|
||||
content: ToolResultContent::text("README contents"),
|
||||
is_error: false,
|
||||
}),
|
||||
]),
|
||||
tools: Cow::Owned(vec![ToolSpec {
|
||||
name: "files".to_string(),
|
||||
description: Some("Read and edit files".to_string()),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"operations": { "type": "array" }
|
||||
}
|
||||
}),
|
||||
output_schema: None,
|
||||
kind: crate::ProviderToolKind::Function,
|
||||
loading_policy: ToolLoadingPolicy::Immediate,
|
||||
strict: None,
|
||||
options: None,
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Tool {
|
||||
name: "files".to_string(),
|
||||
}),
|
||||
temperature: Some(0.2),
|
||||
max_output_tokens: Some(256),
|
||||
metadata: Cow::Owned(BTreeMap::from([(
|
||||
"agent".to_string(),
|
||||
"mentra".to_string(),
|
||||
)])),
|
||||
provider_request_options: ProviderRequestOptions::default(),
|
||||
};
|
||||
|
||||
let payload =
|
||||
serde_json::to_value(GeminiGenerateContentRequest::try_from(request).unwrap())
|
||||
.expect("request should serialize");
|
||||
|
||||
assert_eq!(
|
||||
payload["systemInstruction"]["parts"][0]["text"],
|
||||
"Be helpful."
|
||||
);
|
||||
assert_eq!(payload["contents"][0]["role"], "user");
|
||||
assert_eq!(
|
||||
payload["contents"][0]["parts"][0]["text"],
|
||||
"What files changed?"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["contents"][1]["parts"][0]["functionCall"]["name"],
|
||||
"files"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["contents"][2]["parts"][0]["functionResponse"]["name"],
|
||||
"files"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["contents"][2]["parts"][0]["functionResponse"]["response"]["content"],
|
||||
"README contents"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["tools"][0]["functionDeclarations"][0]["name"],
|
||||
"files"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"ANY"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["toolConfig"]["functionCallingConfig"]["allowedFunctionNames"][0],
|
||||
"files"
|
||||
);
|
||||
let temperature = payload["generationConfig"]["temperature"]
|
||||
.as_f64()
|
||||
.expect("temperature should be numeric");
|
||||
assert!((temperature - 0.2).abs() < 1e-6);
|
||||
assert_eq!(payload["generationConfig"]["maxOutputTokens"], 256);
|
||||
assert!(payload.get("metadata").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serializes_inline_images_into_inline_data_parts() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message {
|
||||
role: Role::User,
|
||||
content: vec![
|
||||
ContentBlock::text("Describe this"),
|
||||
ContentBlock::image_bytes("image/png", [1_u8, 2, 3]),
|
||||
],
|
||||
}]),
|
||||
tools: Cow::Owned(vec![]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions::default(),
|
||||
};
|
||||
|
||||
let payload =
|
||||
serde_json::to_value(GeminiGenerateContentRequest::try_from(request).unwrap())
|
||||
.expect("request should serialize");
|
||||
|
||||
assert_eq!(payload["contents"][0]["parts"][0]["text"], "Describe this");
|
||||
assert_eq!(
|
||||
payload["contents"][0]["parts"][1]["inlineData"]["mimeType"],
|
||||
"image/png"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["contents"][0]["parts"][1]["inlineData"]["data"],
|
||||
"AQID"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_url_images() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::image_url(
|
||||
"https://example.com/image.png",
|
||||
))]),
|
||||
tools: Cow::Owned(vec![]),
|
||||
tool_choice: None,
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions::default(),
|
||||
};
|
||||
|
||||
let error = GeminiGenerateContentRequest::try_from(request)
|
||||
.err()
|
||||
.expect("request should fail");
|
||||
match error {
|
||||
ProviderError::InvalidRequest(message) => {
|
||||
assert!(message.contains("image URL inputs are not supported"));
|
||||
}
|
||||
other => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serializes_tool_choice_modes() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![ToolSpec {
|
||||
name: "echo".to_string(),
|
||||
description: None,
|
||||
input_schema: json!({"type":"object"}),
|
||||
output_schema: None,
|
||||
kind: crate::ProviderToolKind::Function,
|
||||
loading_policy: ToolLoadingPolicy::Immediate,
|
||||
strict: None,
|
||||
options: None,
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Any),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions::default(),
|
||||
};
|
||||
let any_payload =
|
||||
serde_json::to_value(GeminiGenerateContentRequest::try_from(request).unwrap())
|
||||
.expect("request should serialize");
|
||||
assert_eq!(
|
||||
any_payload["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"ANY"
|
||||
);
|
||||
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![ToolSpec {
|
||||
name: "echo".to_string(),
|
||||
description: None,
|
||||
input_schema: json!({"type":"object"}),
|
||||
output_schema: None,
|
||||
kind: crate::ProviderToolKind::Function,
|
||||
loading_policy: ToolLoadingPolicy::Immediate,
|
||||
strict: None,
|
||||
options: None,
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions::default(),
|
||||
};
|
||||
let auto_payload =
|
||||
serde_json::to_value(GeminiGenerateContentRequest::try_from(request).unwrap())
|
||||
.expect("request should serialize");
|
||||
assert_eq!(
|
||||
auto_payload["toolConfig"]["functionCallingConfig"]["mode"],
|
||||
"AUTO"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn omits_tool_config_when_tool_choice_is_unset() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![ToolSpec {
|
||||
name: "echo".to_string(),
|
||||
description: None,
|
||||
input_schema: json!({"type":"object"}),
|
||||
output_schema: None,
|
||||
kind: crate::ProviderToolKind::Function,
|
||||
loading_policy: ToolLoadingPolicy::Immediate,
|
||||
strict: None,
|
||||
options: None,
|
||||
}]),
|
||||
tool_choice: None,
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions::default(),
|
||||
};
|
||||
|
||||
let payload =
|
||||
serde_json::to_value(GeminiGenerateContentRequest::try_from(request).unwrap())
|
||||
.expect("request should serialize");
|
||||
|
||||
assert!(payload.get("toolConfig").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serializes_shared_reasoning_effort_for_gemini_3_models() {
|
||||
for (effort, expected) in [
|
||||
(ReasoningEffort::Low, "low"),
|
||||
(ReasoningEffort::Medium, "medium"),
|
||||
(ReasoningEffort::High, "high"),
|
||||
] {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-3-flash-preview"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
reasoning: Some(ReasoningOptions {
|
||||
effort: Some(effort),
|
||||
summary: None,
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
|
||||
let payload =
|
||||
serde_json::to_value(GeminiGenerateContentRequest::try_from(request).unwrap())
|
||||
.expect("request should serialize");
|
||||
|
||||
assert_eq!(
|
||||
payload["generationConfig"]["thinkingConfig"]["thinkingLevel"],
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_reasoning_effort_for_gemini_2_5_models() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.5-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
reasoning: Some(ReasoningOptions {
|
||||
effort: Some(ReasoningEffort::Low),
|
||||
summary: None,
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
|
||||
let error = GeminiGenerateContentRequest::try_from(request)
|
||||
.err()
|
||||
.expect("request should fail");
|
||||
match error {
|
||||
ProviderError::InvalidRequest(message) => {
|
||||
assert!(message.contains("Gemini 3"));
|
||||
}
|
||||
other => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_extended_reasoning_effort_for_gemini() {
|
||||
for (effort, expected) in [
|
||||
(ReasoningEffort::XHigh, "xhigh"),
|
||||
(ReasoningEffort::Max, "max"),
|
||||
] {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-3-flash-preview"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
reasoning: Some(ReasoningOptions {
|
||||
effort: Some(effort),
|
||||
summary: None,
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
|
||||
let error = GeminiGenerateContentRequest::try_from(request)
|
||||
.err()
|
||||
.expect("extended Gemini effort should fail");
|
||||
match error {
|
||||
ProviderError::InvalidRequest(message) => {
|
||||
assert!(message.contains(expected));
|
||||
}
|
||||
other => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_hosted_tool_search_with_deferred_tools() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![ToolSpec {
|
||||
name: "echo".to_string(),
|
||||
description: None,
|
||||
input_schema: json!({"type":"object"}),
|
||||
output_schema: None,
|
||||
kind: crate::ProviderToolKind::Function,
|
||||
loading_policy: ToolLoadingPolicy::Deferred,
|
||||
strict: None,
|
||||
options: None,
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
tool_search_mode: ToolSearchMode::Hosted,
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
|
||||
let error = GeminiGenerateContentRequest::try_from(request)
|
||||
.err()
|
||||
.expect("request should fail");
|
||||
match error {
|
||||
ProviderError::InvalidRequest(message) => {
|
||||
assert!(message.contains("does not support hosted tool search"));
|
||||
}
|
||||
other => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forced_deferred_tool_still_serializes_as_function_declaration() {
|
||||
let request = Request {
|
||||
model: Cow::Borrowed("gemini-2.0-flash"),
|
||||
system: None,
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text("hi"))]),
|
||||
tools: Cow::Owned(vec![ToolSpec {
|
||||
name: "echo".to_string(),
|
||||
description: None,
|
||||
input_schema: json!({"type":"object"}),
|
||||
output_schema: None,
|
||||
kind: crate::ProviderToolKind::Function,
|
||||
loading_policy: ToolLoadingPolicy::Deferred,
|
||||
strict: None,
|
||||
options: None,
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Tool {
|
||||
name: "echo".to_string(),
|
||||
}),
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(BTreeMap::new()),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
tool_search_mode: ToolSearchMode::Hosted,
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
|
||||
let payload =
|
||||
serde_json::to_value(GeminiGenerateContentRequest::try_from(request).unwrap())
|
||||
.expect("request should serialize");
|
||||
|
||||
assert_eq!(
|
||||
payload["tools"][0]["functionDeclarations"][0]["name"],
|
||||
"echo"
|
||||
);
|
||||
}
|
||||
}
|
||||
714
vendor/mentra-provider/src/gemini/sse.rs
vendored
Normal file
714
vendor/mentra-provider/src/gemini/sse.rs
vendored
Normal file
@@ -0,0 +1,714 @@
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::{
|
||||
ContentBlockDelta, ContentBlockStart, ProviderError, ProviderEvent, ProviderEventStream, Role,
|
||||
TokenUsage,
|
||||
};
|
||||
|
||||
pub(crate) fn spawn_event_stream(
|
||||
response: reqwest::Response,
|
||||
request_model: String,
|
||||
) -> ProviderEventStream {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = forward_events(response, request_model, tx.clone()).await {
|
||||
let _ = tx.send(Err(error));
|
||||
}
|
||||
});
|
||||
|
||||
rx
|
||||
}
|
||||
|
||||
async fn forward_events(
|
||||
response: reqwest::Response,
|
||||
request_model: String,
|
||||
tx: mpsc::UnboundedSender<Result<ProviderEvent, ProviderError>>,
|
||||
) -> Result<(), ProviderError> {
|
||||
let mut bytes_stream = response.bytes_stream();
|
||||
let mut buffer = Vec::new();
|
||||
let mut state = StreamState::new(request_model);
|
||||
|
||||
while let Some(chunk) = bytes_stream.next().await {
|
||||
let chunk = chunk.map_err(ProviderError::Transport)?;
|
||||
buffer.extend_from_slice(&chunk);
|
||||
|
||||
while let Some((frame_end, delimiter_len)) = find_frame_boundary(&buffer) {
|
||||
let frame = buffer.drain(..frame_end).collect::<Vec<_>>();
|
||||
buffer.drain(..delimiter_len);
|
||||
|
||||
for event in parse_frame(&frame, &mut state)? {
|
||||
if tx.send(Ok(event)).is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !buffer.is_empty() {
|
||||
for event in parse_frame(&buffer, &mut state)? {
|
||||
let _ = tx.send(Ok(event));
|
||||
}
|
||||
}
|
||||
|
||||
if state.started && !state.stopped {
|
||||
return Err(ProviderError::MalformedStream(
|
||||
"Gemini stream ended before MessageStopped".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct StreamState {
|
||||
request_model: String,
|
||||
response_id: Option<String>,
|
||||
model_version: Option<String>,
|
||||
started: bool,
|
||||
stopped: bool,
|
||||
latest_usage: Option<TokenUsage>,
|
||||
open_blocks: BTreeSet<usize>,
|
||||
text_snapshots: HashMap<usize, String>,
|
||||
tool_snapshots: HashMap<usize, String>,
|
||||
tool_call_ids: HashMap<usize, String>,
|
||||
}
|
||||
|
||||
impl StreamState {
|
||||
fn new(request_model: String) -> Self {
|
||||
Self {
|
||||
request_model,
|
||||
response_id: None,
|
||||
model_version: None,
|
||||
started: false,
|
||||
stopped: false,
|
||||
latest_usage: None,
|
||||
open_blocks: BTreeSet::new(),
|
||||
text_snapshots: HashMap::new(),
|
||||
tool_snapshots: HashMap::new(),
|
||||
tool_call_ids: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_message_started(&mut self, chunk: &GeminiStreamChunk) -> Option<ProviderEvent> {
|
||||
if self.started {
|
||||
return None;
|
||||
}
|
||||
|
||||
self.started = true;
|
||||
self.response_id = chunk
|
||||
.response_id
|
||||
.clone()
|
||||
.or_else(|| Some(format!("gemini-{}", self.request_model)));
|
||||
self.model_version = chunk.model_version.clone();
|
||||
|
||||
Some(ProviderEvent::MessageStarted {
|
||||
id: self
|
||||
.response_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| "gemini".to_string()),
|
||||
model: self
|
||||
.model_version
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.request_model.clone()),
|
||||
role: Role::Assistant,
|
||||
})
|
||||
}
|
||||
|
||||
fn ensure_text_block_started(&mut self, index: usize) -> Option<ProviderEvent> {
|
||||
if self.open_blocks.insert(index) {
|
||||
Some(ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::Text,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_tool_block_started(
|
||||
&mut self,
|
||||
index: usize,
|
||||
function_call: &GeminiFunctionCall,
|
||||
) -> Option<ProviderEvent> {
|
||||
if self.open_blocks.insert(index) {
|
||||
let response_id = self
|
||||
.response_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("gemini-{}", self.request_model));
|
||||
let id = format!("{response_id}-{index}-{}", function_call.name);
|
||||
self.tool_call_ids.insert(index, id.clone());
|
||||
Some(ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::ToolUse {
|
||||
id,
|
||||
name: function_call.name.clone(),
|
||||
},
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn close_all_blocks(&mut self) -> Vec<ProviderEvent> {
|
||||
let indices = self.open_blocks.iter().copied().collect::<Vec<_>>();
|
||||
self.open_blocks.clear();
|
||||
self.text_snapshots.clear();
|
||||
self.tool_snapshots.clear();
|
||||
self.tool_call_ids.clear();
|
||||
|
||||
indices
|
||||
.into_iter()
|
||||
.map(|index| ProviderEvent::ContentBlockStopped { index })
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn update_usage(&mut self, usage: Option<TokenUsage>) -> Option<TokenUsage> {
|
||||
match usage {
|
||||
Some(usage) if self.latest_usage.as_ref() != Some(&usage) => {
|
||||
self.latest_usage = Some(usage.clone());
|
||||
Some(usage)
|
||||
}
|
||||
Some(usage) => {
|
||||
self.latest_usage = Some(usage);
|
||||
None
|
||||
}
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_frame(frame: &[u8], state: &mut StreamState) -> Result<Vec<ProviderEvent>, ProviderError> {
|
||||
let frame = std::str::from_utf8(frame)
|
||||
.map_err(|error| ProviderError::MalformedStream(error.to_string()))?;
|
||||
let mut data_lines = Vec::new();
|
||||
|
||||
for raw_line in frame.lines() {
|
||||
let line = raw_line.strip_suffix('\r').unwrap_or(raw_line);
|
||||
if line.is_empty() || line.starts_with(':') || line.starts_with("event:") {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(rest) = line.strip_prefix("data:") {
|
||||
data_lines.push(rest.trim_start().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
if data_lines.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let data = data_lines.join("\n");
|
||||
let chunk: GeminiStreamChunk =
|
||||
serde_json::from_str(&data).map_err(ProviderError::Deserialize)?;
|
||||
|
||||
if let Some(error) = chunk.error {
|
||||
return Err(ProviderError::MalformedStream(
|
||||
error
|
||||
.message
|
||||
.unwrap_or_else(|| "gemini stream error".to_string()),
|
||||
));
|
||||
}
|
||||
|
||||
let mut events = Vec::new();
|
||||
let latest_usage = chunk
|
||||
.usage_metadata
|
||||
.as_ref()
|
||||
.and_then(GeminiUsageMetadata::to_token_usage);
|
||||
let usage_changed = state.update_usage(latest_usage.clone());
|
||||
|
||||
if let Some(candidate) = chunk.candidates.first() {
|
||||
if let Some(event) = state.ensure_message_started(&chunk) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
if let Some(content) = candidate.content.as_ref() {
|
||||
for (index, part) in content.parts.iter().enumerate() {
|
||||
if let Some(text) = part.text.as_ref() {
|
||||
if let Some(event) = state.ensure_text_block_started(index) {
|
||||
events.push(event);
|
||||
}
|
||||
if let Some(delta) = merge_chunk(
|
||||
state.text_snapshots.entry(index).or_default(),
|
||||
text.as_str(),
|
||||
) {
|
||||
events.push(ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::Text(delta),
|
||||
});
|
||||
}
|
||||
} else if let Some(function_call) = part.function_call.as_ref() {
|
||||
if let Some(event) = state.ensure_tool_block_started(index, function_call) {
|
||||
events.push(event);
|
||||
}
|
||||
if let Some(delta) = merge_chunk(
|
||||
state.tool_snapshots.entry(index).or_default(),
|
||||
&serde_json::to_string(&function_call.args)
|
||||
.expect("function call args should serialize"),
|
||||
) {
|
||||
events.push(ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::ToolUseInputJson(delta),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(stop_reason) = candidate.finish_reason.clone() {
|
||||
events.extend(state.close_all_blocks());
|
||||
events.push(ProviderEvent::MessageDelta {
|
||||
stop_reason: Some(stop_reason),
|
||||
usage: latest_usage.or_else(|| state.latest_usage.clone()),
|
||||
});
|
||||
events.push(ProviderEvent::MessageStopped);
|
||||
state.stopped = true;
|
||||
} else if let Some(usage) = usage_changed {
|
||||
events.push(ProviderEvent::MessageDelta {
|
||||
stop_reason: None,
|
||||
usage: Some(usage),
|
||||
});
|
||||
}
|
||||
} else if let Some(prompt_feedback) = chunk.prompt_feedback.as_ref() {
|
||||
if let Some(event) = state.ensure_message_started(&chunk) {
|
||||
events.push(event);
|
||||
}
|
||||
let _ = state.update_usage(latest_usage.clone());
|
||||
events.extend(state.close_all_blocks());
|
||||
events.push(ProviderEvent::MessageDelta {
|
||||
stop_reason: Some(prompt_feedback.stop_reason()),
|
||||
usage: latest_usage.or_else(|| state.latest_usage.clone()),
|
||||
});
|
||||
events.push(ProviderEvent::MessageStopped);
|
||||
state.stopped = true;
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
fn merge_chunk(previous: &mut String, current: &str) -> Option<String> {
|
||||
if current.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if previous.is_empty() {
|
||||
*previous = current.to_string();
|
||||
return Some(current.to_string());
|
||||
}
|
||||
|
||||
if current == previous {
|
||||
return None;
|
||||
}
|
||||
|
||||
if current.starts_with(previous.as_str()) {
|
||||
let delta = current[previous.len()..].to_string();
|
||||
*previous = current.to_string();
|
||||
return (!delta.is_empty()).then_some(delta);
|
||||
}
|
||||
|
||||
previous.push_str(current);
|
||||
Some(current.to_string())
|
||||
}
|
||||
|
||||
fn find_frame_boundary(buffer: &[u8]) -> Option<(usize, usize)> {
|
||||
for (index, window) in buffer.windows(2).enumerate() {
|
||||
if window == b"\n\n" {
|
||||
return Some((index, 2));
|
||||
}
|
||||
}
|
||||
|
||||
for (index, window) in buffer.windows(4).enumerate() {
|
||||
if window == b"\r\n\r\n" {
|
||||
return Some((index, 4));
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiStreamChunk {
|
||||
#[serde(default)]
|
||||
candidates: Vec<GeminiCandidate>,
|
||||
#[serde(default, rename = "promptFeedback", alias = "prompt_feedback")]
|
||||
prompt_feedback: Option<GeminiPromptFeedback>,
|
||||
#[serde(default, rename = "usageMetadata", alias = "usage_metadata")]
|
||||
usage_metadata: Option<GeminiUsageMetadata>,
|
||||
#[serde(default, rename = "responseId", alias = "response_id")]
|
||||
response_id: Option<String>,
|
||||
#[serde(default, rename = "modelVersion", alias = "model_version")]
|
||||
model_version: Option<String>,
|
||||
#[serde(default)]
|
||||
error: Option<GeminiErrorBody>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiCandidate {
|
||||
#[serde(default)]
|
||||
content: Option<GeminiContent>,
|
||||
#[serde(default, rename = "finishReason", alias = "finish_reason")]
|
||||
finish_reason: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiContent {
|
||||
#[allow(dead_code)]
|
||||
#[serde(default)]
|
||||
role: Option<String>,
|
||||
#[serde(default)]
|
||||
parts: Vec<GeminiPart>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiPart {
|
||||
#[serde(default)]
|
||||
text: Option<String>,
|
||||
#[serde(default, rename = "functionCall", alias = "function_call")]
|
||||
function_call: Option<GeminiFunctionCall>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiFunctionCall {
|
||||
name: String,
|
||||
#[serde(default)]
|
||||
args: Value,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiErrorBody {
|
||||
#[serde(default)]
|
||||
message: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiPromptFeedback {
|
||||
#[serde(default, rename = "blockReason", alias = "block_reason")]
|
||||
block_reason: Option<String>,
|
||||
}
|
||||
|
||||
impl GeminiPromptFeedback {
|
||||
fn stop_reason(&self) -> String {
|
||||
self.block_reason
|
||||
.clone()
|
||||
.unwrap_or_else(|| "BLOCKED".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GeminiUsageMetadata {
|
||||
#[serde(default, rename = "promptTokenCount", alias = "prompt_token_count")]
|
||||
prompt_token_count: Option<u64>,
|
||||
#[serde(
|
||||
default,
|
||||
rename = "candidatesTokenCount",
|
||||
alias = "candidates_token_count"
|
||||
)]
|
||||
candidates_token_count: Option<u64>,
|
||||
#[serde(default, rename = "totalTokenCount", alias = "total_token_count")]
|
||||
total_token_count: Option<u64>,
|
||||
#[serde(
|
||||
default,
|
||||
rename = "cachedContentTokenCount",
|
||||
alias = "cached_content_token_count"
|
||||
)]
|
||||
cached_content_token_count: Option<u64>,
|
||||
#[serde(default, rename = "thoughtsTokenCount", alias = "thoughts_token_count")]
|
||||
thoughts_token_count: Option<u64>,
|
||||
#[serde(
|
||||
default,
|
||||
rename = "toolUsePromptTokenCount",
|
||||
alias = "tool_use_prompt_token_count"
|
||||
)]
|
||||
tool_use_prompt_token_count: Option<u64>,
|
||||
}
|
||||
|
||||
impl GeminiUsageMetadata {
|
||||
fn to_token_usage(&self) -> Option<TokenUsage> {
|
||||
let usage = TokenUsage {
|
||||
input_tokens: self.prompt_token_count,
|
||||
output_tokens: self.candidates_token_count,
|
||||
total_tokens: self.total_token_count,
|
||||
cache_read_input_tokens: self.cached_content_token_count,
|
||||
cache_creation_input_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
thoughts_tokens: self.thoughts_token_count,
|
||||
tool_input_tokens: self.tool_use_prompt_token_count,
|
||||
};
|
||||
|
||||
(!usage.is_empty()).then_some(usage)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::{ContentBlockDelta, ContentBlockStart, ProviderEvent, TokenUsage};
|
||||
|
||||
use super::{StreamState, parse_frame};
|
||||
|
||||
#[test]
|
||||
fn streams_text_and_completion_events() {
|
||||
let mut state = StreamState::new("gemini-2.0-flash".to_string());
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"responseId":"resp-1","modelVersion":"gemini-2.0-flash-001","candidates":[{"content":{"role":"model","parts":[{"text":"Hel"}]}}]}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ProviderEvent::MessageStarted {
|
||||
id: "resp-1".to_string(),
|
||||
model: "gemini-2.0-flash-001".to_string(),
|
||||
role: crate::Role::Assistant,
|
||||
},
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index: 0,
|
||||
kind: ContentBlockStart::Text,
|
||||
},
|
||||
ProviderEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::Text("Hel".to_string()),
|
||||
},
|
||||
]
|
||||
);
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"candidates":[{"content":{"parts":[{"text":"lo"}]}}]}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![ProviderEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::Text("lo".to_string()),
|
||||
}]
|
||||
);
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"candidates":[{"finishReason":"STOP"}]}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ProviderEvent::ContentBlockStopped { index: 0 },
|
||||
ProviderEvent::MessageDelta {
|
||||
stop_reason: Some("STOP".to_string()),
|
||||
usage: None,
|
||||
},
|
||||
ProviderEvent::MessageStopped,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streams_function_calls() {
|
||||
let mut state = StreamState::new("gemini-2.0-flash".to_string());
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"responseId":"resp-1","candidates":[{"content":{"parts":[{"functionCall":{"name":"read_file","args":{"path":"README.md"}}}]}}]}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ProviderEvent::MessageStarted {
|
||||
id: "resp-1".to_string(),
|
||||
model: "gemini-2.0-flash".to_string(),
|
||||
role: crate::Role::Assistant,
|
||||
},
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index: 0,
|
||||
kind: ContentBlockStart::ToolUse {
|
||||
id: "resp-1-0-read_file".to_string(),
|
||||
name: "read_file".to_string(),
|
||||
},
|
||||
},
|
||||
ProviderEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::ToolUseInputJson(
|
||||
"{\"path\":\"README.md\"}".to_string()
|
||||
),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_duplicate_full_function_call_payloads() {
|
||||
let mut state = StreamState::new("gemini-2.0-flash".to_string());
|
||||
parse_frame(
|
||||
br#"data: {"responseId":"resp-1","candidates":[{"content":{"parts":[{"functionCall":{"name":"read_file","args":{"path":"README.md"}}}]}}]}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("first frame should parse");
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"read_file","args":{"path":"README.md"}}}]}}]}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("second frame should parse");
|
||||
|
||||
assert!(events.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_unsupported_parts_without_breaking_indexes() {
|
||||
let mut state = StreamState::new("gemini-2.0-flash".to_string());
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"responseId":"resp-1","candidates":[{"content":{"parts":[{"fileData":{"mimeType":"image/png","fileUri":"files/1"}},{"text":"Done"}]}}]}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ProviderEvent::MessageStarted {
|
||||
id: "resp-1".to_string(),
|
||||
model: "gemini-2.0-flash".to_string(),
|
||||
role: crate::Role::Assistant,
|
||||
},
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index: 1,
|
||||
kind: ContentBlockStart::Text,
|
||||
},
|
||||
ProviderEvent::ContentBlockDelta {
|
||||
index: 1,
|
||||
delta: ContentBlockDelta::Text("Done".to_string()),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn surfaces_stream_errors() {
|
||||
let mut state = StreamState::new("gemini-2.0-flash".to_string());
|
||||
let error = parse_frame(br#"data: {"error":{"message":"boom"}}"#, &mut state)
|
||||
.expect_err("frame should fail");
|
||||
|
||||
match error {
|
||||
crate::ProviderError::MalformedStream(message) => {
|
||||
assert_eq!(message, "boom");
|
||||
}
|
||||
other => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn treats_prompt_feedback_only_chunks_as_terminal() {
|
||||
let mut state = StreamState::new("gemini-2.0-flash".to_string());
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"responseId":"resp-2","promptFeedback":{"blockReason":"SAFETY"},"usageMetadata":{"promptTokenCount":11,"totalTokenCount":11,"cachedContentTokenCount":4}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ProviderEvent::MessageStarted {
|
||||
id: "resp-2".to_string(),
|
||||
model: "gemini-2.0-flash".to_string(),
|
||||
role: crate::Role::Assistant,
|
||||
},
|
||||
ProviderEvent::MessageDelta {
|
||||
stop_reason: Some("SAFETY".to_string()),
|
||||
usage: Some(TokenUsage {
|
||||
input_tokens: Some(11),
|
||||
output_tokens: None,
|
||||
total_tokens: Some(11),
|
||||
cache_read_input_tokens: Some(4),
|
||||
cache_creation_input_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
thoughts_tokens: None,
|
||||
tool_input_tokens: None,
|
||||
}),
|
||||
},
|
||||
ProviderEvent::MessageStopped,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emits_usage_metadata_updates_and_final_usage() {
|
||||
let mut state = StreamState::new("gemini-2.0-flash".to_string());
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"responseId":"resp-3","candidates":[{"content":{"parts":[{"text":"Hi"}]}}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":1,"totalTokenCount":9,"thoughtsTokenCount":2,"toolUsePromptTokenCount":1}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ProviderEvent::MessageStarted {
|
||||
id: "resp-3".to_string(),
|
||||
model: "gemini-2.0-flash".to_string(),
|
||||
role: crate::Role::Assistant,
|
||||
},
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index: 0,
|
||||
kind: ContentBlockStart::Text,
|
||||
},
|
||||
ProviderEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::Text("Hi".to_string()),
|
||||
},
|
||||
ProviderEvent::MessageDelta {
|
||||
stop_reason: None,
|
||||
usage: Some(TokenUsage {
|
||||
input_tokens: Some(8),
|
||||
output_tokens: Some(1),
|
||||
total_tokens: Some(9),
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
thoughts_tokens: Some(2),
|
||||
tool_input_tokens: Some(1),
|
||||
}),
|
||||
},
|
||||
]
|
||||
);
|
||||
|
||||
let events = parse_frame(
|
||||
br#"data: {"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":2,"totalTokenCount":10,"thoughtsTokenCount":2,"toolUsePromptTokenCount":1}}"#,
|
||||
&mut state,
|
||||
)
|
||||
.expect("frame should parse");
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
ProviderEvent::ContentBlockStopped { index: 0 },
|
||||
ProviderEvent::MessageDelta {
|
||||
stop_reason: Some("STOP".to_string()),
|
||||
usage: Some(TokenUsage {
|
||||
input_tokens: Some(8),
|
||||
output_tokens: Some(2),
|
||||
total_tokens: Some(10),
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
thoughts_tokens: Some(2),
|
||||
tool_input_tokens: Some(1),
|
||||
}),
|
||||
},
|
||||
ProviderEvent::MessageStopped,
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
154
vendor/mentra-provider/src/lib.rs
vendored
Normal file
154
vendor/mentra-provider/src/lib.rs
vendored
Normal file
@@ -0,0 +1,154 @@
|
||||
pub mod anthropic;
|
||||
mod auth;
|
||||
mod definition;
|
||||
pub mod embedding;
|
||||
mod error;
|
||||
pub mod gemini;
|
||||
mod model;
|
||||
mod registry;
|
||||
mod request;
|
||||
mod response;
|
||||
pub mod responses;
|
||||
mod stream;
|
||||
mod tool;
|
||||
|
||||
pub use auth::AuthScheme;
|
||||
pub use auth::CredentialSource;
|
||||
pub use auth::ProviderCredentials;
|
||||
pub use auth::StaticCredentialSource;
|
||||
pub use definition::BuiltinProvider;
|
||||
pub use definition::ProviderCapabilities;
|
||||
pub use definition::ProviderDefinition;
|
||||
pub use definition::ProviderDescriptor;
|
||||
pub use definition::ProviderId;
|
||||
pub use definition::RetryPolicy;
|
||||
pub use definition::WireApi;
|
||||
pub use embedding::EmbeddingData;
|
||||
pub use embedding::EmbeddingModelInfo;
|
||||
pub use embedding::EmbeddingProvider;
|
||||
pub use embedding::EmbeddingRequest;
|
||||
pub use embedding::EmbeddingResponse;
|
||||
pub use embedding::EmbeddingUsage;
|
||||
pub use error::ProviderError;
|
||||
pub use model::ContentBlock;
|
||||
pub use model::HostedToolSearchCall;
|
||||
pub use model::HostedWebSearchCall;
|
||||
pub use model::ImageGenerationCall;
|
||||
pub use model::ImageGenerationResult;
|
||||
pub use model::ImageSource;
|
||||
pub use model::Message;
|
||||
pub use model::ModelInfo;
|
||||
pub use model::ModelSelector;
|
||||
pub use model::ReasoningFormat;
|
||||
pub use model::ReasoningProvenance;
|
||||
pub use model::Role;
|
||||
pub use model::TokenUsage;
|
||||
pub use model::ToolChoice;
|
||||
pub use model::ToolResultContent;
|
||||
pub use model::WebSearchAction;
|
||||
pub use registry::ModelCatalog;
|
||||
pub use registry::Provider;
|
||||
pub use registry::ProviderRegistry;
|
||||
pub use registry::ProviderSession;
|
||||
pub use registry::ProviderSessionFactory;
|
||||
pub use registry::RegisteredProvider;
|
||||
pub use request::AnthropicRequestOptions;
|
||||
pub use request::CompactionInputItem;
|
||||
pub use request::CompactionRequest;
|
||||
pub use request::GeminiRequestOptions;
|
||||
pub use request::MemorySummarizeRequest;
|
||||
pub use request::ProviderRequestOptions;
|
||||
pub use request::RawMemory;
|
||||
pub use request::RawMemoryMetadata;
|
||||
pub use request::ReasoningEffort;
|
||||
pub use request::ReasoningOptions;
|
||||
pub use request::ReasoningSummary;
|
||||
pub use request::Request;
|
||||
pub use request::ResponsesRequestCompression;
|
||||
pub use request::ResponsesRequestOptions;
|
||||
pub use request::ResponsesStateMode;
|
||||
pub use request::ResponsesTextControls;
|
||||
pub use request::ResponsesTextFormat;
|
||||
pub use request::ResponsesTextFormatType;
|
||||
pub use request::ResponsesTransport;
|
||||
pub use request::ResponsesVerbosity;
|
||||
pub use request::SessionRequestOptions;
|
||||
pub use request::ToolSearchMode;
|
||||
pub use response::CompactionResponse;
|
||||
pub use response::MemorySummarizeOutput;
|
||||
pub use response::MemorySummarizeResponse;
|
||||
pub use response::Response;
|
||||
pub use response::collect_response_from_stream;
|
||||
pub use response::provider_event_stream_from_response;
|
||||
pub use stream::ContentBlockDelta;
|
||||
pub use stream::ContentBlockStart;
|
||||
pub use stream::ProviderEvent;
|
||||
pub use stream::ProviderEventStream;
|
||||
pub use stream::ResponseHeaders;
|
||||
pub use tool::ProviderToolKind;
|
||||
pub use tool::ToolLoadingPolicy;
|
||||
pub use tool::ToolSpec;
|
||||
pub use tool::ToolSpecBuilder;
|
||||
|
||||
pub type OpenAIRequestOptions = ResponsesRequestOptions;
|
||||
|
||||
pub mod provider {
|
||||
pub use crate::Provider;
|
||||
|
||||
pub mod model {
|
||||
pub use crate::AnthropicRequestOptions;
|
||||
pub use crate::ContentBlock;
|
||||
pub use crate::ContentBlockDelta;
|
||||
pub use crate::ContentBlockStart;
|
||||
pub use crate::HostedToolSearchCall;
|
||||
pub use crate::HostedWebSearchCall;
|
||||
pub use crate::ImageGenerationCall;
|
||||
pub use crate::ImageGenerationResult;
|
||||
pub use crate::ImageSource;
|
||||
pub use crate::Message;
|
||||
pub use crate::ModelInfo;
|
||||
pub use crate::OpenAIRequestOptions;
|
||||
pub use crate::ProviderError;
|
||||
pub use crate::ProviderEvent;
|
||||
pub use crate::ProviderEventStream;
|
||||
pub use crate::ProviderId;
|
||||
pub use crate::ProviderRequestOptions;
|
||||
pub use crate::ReasoningEffort;
|
||||
pub use crate::ReasoningFormat;
|
||||
pub use crate::ReasoningOptions;
|
||||
pub use crate::ReasoningProvenance;
|
||||
pub use crate::ReasoningSummary;
|
||||
pub use crate::Request;
|
||||
pub use crate::Response;
|
||||
pub use crate::ResponsesTextControls;
|
||||
pub use crate::ResponsesTextFormat;
|
||||
pub use crate::ResponsesTextFormatType;
|
||||
pub use crate::ResponsesVerbosity;
|
||||
pub use crate::Role;
|
||||
pub use crate::SessionRequestOptions;
|
||||
pub use crate::TokenUsage;
|
||||
pub use crate::ToolChoice;
|
||||
pub use crate::ToolResultContent;
|
||||
pub use crate::ToolSearchMode;
|
||||
pub use crate::WebSearchAction;
|
||||
pub use crate::collect_response_from_stream;
|
||||
pub use crate::provider_event_stream_from_response;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn provider_definition_defaults_to_responses_wire_api_and_websockets_disabled() {
|
||||
let definition = ProviderDefinition::new(BuiltinProvider::OpenAI);
|
||||
|
||||
assert_eq!(
|
||||
definition.descriptor.id,
|
||||
ProviderId::from(BuiltinProvider::OpenAI)
|
||||
);
|
||||
assert_eq!(definition.wire_api, WireApi::Responses);
|
||||
assert!(!definition.capabilities.supports_websockets);
|
||||
}
|
||||
}
|
||||
509
vendor/mentra-provider/src/model.rs
vendored
Normal file
509
vendor/mentra-provider/src/model.rs
vendored
Normal file
@@ -0,0 +1,509 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::fmt::Display;
|
||||
use time::OffsetDateTime;
|
||||
|
||||
/// Metadata describing a model available from a provider.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ModelInfo {
|
||||
pub id: String,
|
||||
pub provider: crate::ProviderId,
|
||||
pub display_name: Option<String>,
|
||||
pub description: Option<String>,
|
||||
pub created_at: Option<OffsetDateTime>,
|
||||
}
|
||||
|
||||
impl ModelInfo {
|
||||
pub fn new(id: impl Into<String>, provider: impl Into<crate::ProviderId>) -> Self {
|
||||
Self {
|
||||
id: id.into(),
|
||||
provider: provider.into(),
|
||||
display_name: None,
|
||||
description: None,
|
||||
created_at: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Selection strategy used when resolving a model from a provider.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ModelSelector {
|
||||
Id(String),
|
||||
NewestAvailable,
|
||||
}
|
||||
|
||||
/// Provider-neutral token usage metadata for a completed or in-progress response.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct TokenUsage {
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
pub total_tokens: Option<u64>,
|
||||
pub cache_read_input_tokens: Option<u64>,
|
||||
pub cache_creation_input_tokens: Option<u64>,
|
||||
pub reasoning_tokens: Option<u64>,
|
||||
pub thoughts_tokens: Option<u64>,
|
||||
pub tool_input_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
impl TokenUsage {
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.input_tokens.is_none()
|
||||
&& self.output_tokens.is_none()
|
||||
&& self.total_tokens.is_none()
|
||||
&& self.cache_read_input_tokens.is_none()
|
||||
&& self.cache_creation_input_tokens.is_none()
|
||||
&& self.reasoning_tokens.is_none()
|
||||
&& self.thoughts_tokens.is_none()
|
||||
&& self.tool_input_tokens.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider-neutral chat role labels.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum Role {
|
||||
User,
|
||||
Assistant,
|
||||
Unknown(String),
|
||||
}
|
||||
|
||||
impl Display for Role {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let value = match self {
|
||||
Self::User => "user",
|
||||
Self::Assistant => "assistant",
|
||||
Self::Unknown(role) => role.as_str(),
|
||||
};
|
||||
f.write_str(value)
|
||||
}
|
||||
}
|
||||
|
||||
/// Image payload supported by model providers.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum ImageSource {
|
||||
Bytes { media_type: String, data: Vec<u8> },
|
||||
Url { url: String },
|
||||
}
|
||||
|
||||
impl ImageSource {
|
||||
pub fn bytes(media_type: impl Into<String>, data: impl Into<Vec<u8>>) -> Self {
|
||||
Self::Bytes {
|
||||
media_type: media_type.into(),
|
||||
data: data.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn url(url: impl Into<String>) -> Self {
|
||||
Self::Url { url: url.into() }
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool result payloads supported by provider streams and history replay.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ToolResultContent {
|
||||
Text(String),
|
||||
Structured(Value),
|
||||
}
|
||||
|
||||
impl ToolResultContent {
|
||||
pub fn text(value: impl Into<String>) -> Self {
|
||||
Self::Text(value.into())
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
match self {
|
||||
Self::Text(text) => text.len(),
|
||||
Self::Structured(value) => value.to_string().len(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
|
||||
pub fn clear(&mut self) {
|
||||
*self = Self::Text(String::new());
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
Self::Text(text) => text.as_str(),
|
||||
Self::Structured(_) => panic!("ToolResultContent::as_str requires text content"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn contains(&self, pattern: &str) -> bool {
|
||||
match self {
|
||||
Self::Text(text) => text.contains(pattern),
|
||||
Self::Structured(value) => value.to_string().contains(pattern),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn starts_with(&self, pattern: &str) -> bool {
|
||||
match self {
|
||||
Self::Text(text) => text.starts_with(pattern),
|
||||
Self::Structured(value) => value.to_string().starts_with(pattern),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push_str(&mut self, value: &str) {
|
||||
match self {
|
||||
Self::Text(text) => text.push_str(value),
|
||||
Self::Structured(existing) => {
|
||||
let mut text = existing.to_string();
|
||||
text.push_str(value);
|
||||
*self = Self::Text(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_display_string(&self) -> String {
|
||||
match self {
|
||||
Self::Text(text) => text.clone(),
|
||||
Self::Structured(value) => value.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ToolResultContent {
|
||||
fn default() -> Self {
|
||||
Self::Text(String::new())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for ToolResultContent {
|
||||
fn from(value: String) -> Self {
|
||||
Self::Text(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for ToolResultContent {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(&self.to_display_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<&str> for ToolResultContent {
|
||||
fn eq(&self, other: &&str) -> bool {
|
||||
self.to_display_string() == *other
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<str> for ToolResultContent {
|
||||
fn eq(&self, other: &str) -> bool {
|
||||
self.to_display_string() == other
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<ToolResultContent> for &str {
|
||||
fn eq(&self, other: &ToolResultContent) -> bool {
|
||||
*self == other.to_display_string()
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<ToolResultContent> for str {
|
||||
fn eq(&self, other: &ToolResultContent) -> bool {
|
||||
self == other.to_display_string()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for ToolResultContent {
|
||||
fn from(value: &str) -> Self {
|
||||
Self::Text(value.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider-neutral hosted tool search action.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct HostedToolSearchCall {
|
||||
pub id: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub status: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub query: Option<String>,
|
||||
}
|
||||
|
||||
/// Provider-neutral hosted web search actions.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum WebSearchAction {
|
||||
Search {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
query: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
queries: Option<Vec<String>>,
|
||||
},
|
||||
OpenPage {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
url: Option<String>,
|
||||
},
|
||||
FindInPage {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
url: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pattern: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Provider-neutral hosted web search call.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct HostedWebSearchCall {
|
||||
pub id: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub status: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub action: Option<WebSearchAction>,
|
||||
}
|
||||
|
||||
/// Provider-neutral image generation result.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ImageGenerationResult {
|
||||
Image { source: ImageSource },
|
||||
ArtifactRef { artifact_id: String },
|
||||
}
|
||||
|
||||
/// Provider-neutral image generation call.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ImageGenerationCall {
|
||||
pub id: String,
|
||||
pub status: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub revised_prompt: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<ImageGenerationResult>,
|
||||
}
|
||||
|
||||
/// Provider-specific format carried by a provider-neutral reasoning block.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ReasoningFormat {
|
||||
AnthropicSigned,
|
||||
OpenAiEncrypted,
|
||||
GeminiThought,
|
||||
}
|
||||
|
||||
/// Origin required to decide whether opaque reasoning metadata is safe to replay.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ReasoningProvenance {
|
||||
pub provider: crate::ProviderId,
|
||||
pub model: String,
|
||||
pub format: ReasoningFormat,
|
||||
}
|
||||
|
||||
/// A provider-neutral content block exchanged with models.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum ContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
Thinking {
|
||||
#[serde(default, skip_serializing_if = "String::is_empty")]
|
||||
thinking: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
signature: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
encrypted_content: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
provenance: Option<ReasoningProvenance>,
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
redacted: bool,
|
||||
},
|
||||
Image {
|
||||
source: ImageSource,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
input: Value,
|
||||
},
|
||||
ToolResult {
|
||||
tool_use_id: String,
|
||||
content: ToolResultContent,
|
||||
is_error: bool,
|
||||
},
|
||||
HostedToolSearch {
|
||||
call: HostedToolSearchCall,
|
||||
},
|
||||
HostedWebSearch {
|
||||
call: HostedWebSearchCall,
|
||||
},
|
||||
ImageGeneration {
|
||||
call: ImageGenerationCall,
|
||||
},
|
||||
}
|
||||
|
||||
impl ContentBlock {
|
||||
pub fn text(text: impl Into<String>) -> Self {
|
||||
Self::Text { text: text.into() }
|
||||
}
|
||||
|
||||
pub fn thinking(thinking: impl Into<String>) -> Self {
|
||||
Self::Thinking {
|
||||
thinking: thinking.into(),
|
||||
signature: None,
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: None,
|
||||
redacted: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn thinking_fallback_text(&self) -> Option<String> {
|
||||
let Self::Thinking {
|
||||
thinking, redacted, ..
|
||||
} = self
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
|
||||
if !thinking.is_empty() {
|
||||
Some(thinking.clone())
|
||||
} else if *redacted {
|
||||
Some("[redacted reasoning]".to_string())
|
||||
} else {
|
||||
Some("[reasoning unavailable]".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn image_bytes(media_type: impl Into<String>, data: impl Into<Vec<u8>>) -> Self {
|
||||
Self::Image {
|
||||
source: ImageSource::bytes(media_type, data),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn image_url(url: impl Into<String>) -> Self {
|
||||
Self::Image {
|
||||
source: ImageSource::url(url),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn thinking_serde_is_externally_tagged_and_omits_empty_optional_fields() {
|
||||
let block = ContentBlock::Thinking {
|
||||
thinking: "private chain".to_string(),
|
||||
signature: Some("opaque-signature".to_string()),
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(ReasoningProvenance {
|
||||
provider: crate::ProviderId::new("anthropic-edge"),
|
||||
model: "claude-test".to_string(),
|
||||
format: ReasoningFormat::AnthropicSigned,
|
||||
}),
|
||||
redacted: false,
|
||||
};
|
||||
|
||||
let json = serde_json::to_value(&block).expect("thinking block should serialize");
|
||||
assert_eq!(json["Thinking"]["thinking"], "private chain");
|
||||
assert_eq!(json["Thinking"]["signature"], "opaque-signature");
|
||||
assert_eq!(json["Thinking"]["provenance"]["provider"], "anthropic-edge");
|
||||
assert_eq!(json["Thinking"]["provenance"]["format"], "anthropic_signed");
|
||||
assert!(json["Thinking"].get("encrypted_content").is_none());
|
||||
assert!(json["Thinking"].get("id").is_none());
|
||||
assert!(json["Thinking"].get("redacted").is_none());
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ContentBlock>(json).expect("thinking block should load"),
|
||||
block
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn thinking_serde_defaults_omitted_payload_fields() {
|
||||
let block: ContentBlock = serde_json::from_value(serde_json::json!({
|
||||
"Thinking": {
|
||||
"provenance": {
|
||||
"provider": "anthropic",
|
||||
"model": "claude-test",
|
||||
"format": "anthropic_signed"
|
||||
}
|
||||
}
|
||||
}))
|
||||
.expect("omitted thinking payload fields should default");
|
||||
|
||||
assert_eq!(
|
||||
block,
|
||||
ContentBlock::Thinking {
|
||||
thinking: String::new(),
|
||||
signature: None,
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(ReasoningProvenance {
|
||||
provider: crate::ProviderId::new("anthropic"),
|
||||
model: "claude-test".to_string(),
|
||||
format: ReasoningFormat::AnthropicSigned,
|
||||
}),
|
||||
redacted: false,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_thinking_content_block_json_still_deserializes_unchanged() {
|
||||
let json = serde_json::json!({"Text":{"text":"legacy"}});
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ContentBlock>(json).expect("legacy block should load"),
|
||||
ContentBlock::text("legacy")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider-neutral chat message content.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct Message {
|
||||
pub role: Role,
|
||||
pub content: Vec<ContentBlock>,
|
||||
}
|
||||
|
||||
impl Message {
|
||||
pub fn user(content: ContentBlock) -> Self {
|
||||
Self {
|
||||
role: Role::User,
|
||||
content: vec![content],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn assistant(content: ContentBlock) -> Self {
|
||||
Self {
|
||||
role: Role::Assistant,
|
||||
content: vec![content],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn unknown(role: impl Into<String>, content: ContentBlock) -> Self {
|
||||
Self {
|
||||
role: Role::Unknown(role.into()),
|
||||
content: vec![content],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn text(&self) -> String {
|
||||
self.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider-neutral tool choice hint passed to model APIs.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub enum ToolChoice {
|
||||
#[default]
|
||||
Auto,
|
||||
Any,
|
||||
Tool {
|
||||
name: String,
|
||||
},
|
||||
}
|
||||
215
vendor/mentra-provider/src/registry.rs
vendored
Normal file
215
vendor/mentra-provider/src/registry.rs
vendored
Normal file
@@ -0,0 +1,215 @@
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::definition::ProviderDefinition;
|
||||
use crate::definition::ProviderDescriptor;
|
||||
use crate::definition::ProviderId;
|
||||
use crate::error::ProviderError;
|
||||
use crate::model::ModelInfo;
|
||||
use crate::request::CompactionRequest;
|
||||
use crate::request::MemorySummarizeRequest;
|
||||
use crate::request::Request;
|
||||
use crate::response::CompactionResponse;
|
||||
use crate::response::MemorySummarizeResponse;
|
||||
use crate::response::Response;
|
||||
use crate::response::collect_response_from_stream;
|
||||
use crate::stream::ProviderEventStream;
|
||||
|
||||
/// Lists models available from a provider.
|
||||
#[async_trait]
|
||||
pub trait ModelCatalog: Send + Sync {
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError>;
|
||||
}
|
||||
|
||||
/// Creates a provider session on demand.
|
||||
#[async_trait]
|
||||
pub trait ProviderSessionFactory: Send + Sync {
|
||||
async fn create_session(&self) -> Result<Box<dyn ProviderSession>, ProviderError>;
|
||||
}
|
||||
|
||||
/// Transport-neutral session used to stream model responses.
|
||||
#[async_trait]
|
||||
pub trait ProviderSession: Send + Sync {
|
||||
async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError>;
|
||||
|
||||
async fn send(&self, request: Request<'_>) -> Result<Response, ProviderError> {
|
||||
collect_response_from_stream(self.stream(request).await?).await
|
||||
}
|
||||
|
||||
async fn compact(
|
||||
&self,
|
||||
_request: CompactionRequest<'_>,
|
||||
) -> Result<CompactionResponse, ProviderError> {
|
||||
Err(ProviderError::UnsupportedCapability(
|
||||
"history_compaction".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn summarize_memories(
|
||||
&self,
|
||||
_request: MemorySummarizeRequest<'_>,
|
||||
) -> Result<MemorySummarizeResponse, ProviderError> {
|
||||
Err(ProviderError::UnsupportedCapability(
|
||||
"memory_summarization".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Transport-neutral provider registration interface.
|
||||
#[async_trait]
|
||||
pub trait Provider: ModelCatalog + ProviderSessionFactory {
|
||||
fn definition(&self) -> ProviderDefinition;
|
||||
|
||||
fn descriptor(&self) -> ProviderDescriptor {
|
||||
self.definition().descriptor
|
||||
}
|
||||
|
||||
async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
|
||||
self.create_session().await?.stream(request).await
|
||||
}
|
||||
|
||||
async fn send(&self, request: Request<'_>) -> Result<Response, ProviderError> {
|
||||
collect_response_from_stream(self.stream(request).await?).await
|
||||
}
|
||||
|
||||
async fn compact(
|
||||
&self,
|
||||
request: CompactionRequest<'_>,
|
||||
) -> Result<CompactionResponse, ProviderError> {
|
||||
self.create_session().await?.compact(request).await
|
||||
}
|
||||
|
||||
async fn summarize_memories(
|
||||
&self,
|
||||
request: MemorySummarizeRequest<'_>,
|
||||
) -> Result<MemorySummarizeResponse, ProviderError> {
|
||||
self.create_session()
|
||||
.await?
|
||||
.summarize_memories(request)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub use Provider as RegisteredProvider;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct ProviderRegistry {
|
||||
default_provider: Option<ProviderId>,
|
||||
providers: HashMap<ProviderId, Arc<dyn Provider>>,
|
||||
}
|
||||
|
||||
impl ProviderRegistry {
|
||||
pub fn register_provider_instance<P>(&mut self, provider: P)
|
||||
where
|
||||
P: Provider + 'static,
|
||||
{
|
||||
let definition = provider.definition();
|
||||
let id = definition.descriptor.id.clone();
|
||||
|
||||
if self.default_provider.is_none() {
|
||||
self.default_provider = Some(id.clone());
|
||||
}
|
||||
|
||||
self.providers.insert(id, Arc::new(provider));
|
||||
}
|
||||
|
||||
pub fn get_provider(&self, id: Option<&ProviderId>) -> Option<Arc<dyn Provider>> {
|
||||
match id {
|
||||
Some(id) => self.providers.get(id).cloned(),
|
||||
None => self
|
||||
.default_provider
|
||||
.as_ref()
|
||||
.and_then(|id| self.providers.get(id).cloned()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn definitions(&self) -> Vec<ProviderDefinition> {
|
||||
self.providers
|
||||
.values()
|
||||
.map(|provider| provider.definition())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn descriptors(&self) -> Vec<ProviderDescriptor> {
|
||||
self.providers
|
||||
.values()
|
||||
.map(|provider| provider.descriptor())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.providers.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestProvider {
|
||||
definition: ProviderDefinition,
|
||||
models: Vec<ModelInfo>,
|
||||
}
|
||||
|
||||
struct TestSession;
|
||||
|
||||
#[async_trait]
|
||||
impl ModelCatalog for TestProvider {
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
Ok(self.models.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderSessionFactory for TestProvider {
|
||||
async fn create_session(&self) -> Result<Box<dyn ProviderSession>, ProviderError> {
|
||||
Ok(Box::new(TestSession))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for TestProvider {
|
||||
fn definition(&self) -> ProviderDefinition {
|
||||
self.definition.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderSession for TestSession {
|
||||
async fn stream(
|
||||
&self,
|
||||
_request: Request<'_>,
|
||||
) -> Result<ProviderEventStream, ProviderError> {
|
||||
let (_tx, rx) = mpsc::unbounded_channel();
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn registry_returns_registered_provider_descriptors() {
|
||||
let mut registry = ProviderRegistry::default();
|
||||
let provider = TestProvider {
|
||||
definition: ProviderDefinition::new("test-provider"),
|
||||
models: vec![ModelInfo::new("model-1", "test-provider")],
|
||||
};
|
||||
|
||||
registry.register_provider_instance(provider);
|
||||
|
||||
assert_eq!(registry.descriptors().len(), 1);
|
||||
assert_eq!(registry.definitions().len(), 1);
|
||||
assert_eq!(
|
||||
registry
|
||||
.get_provider(None)
|
||||
.expect("provider should exist")
|
||||
.definition()
|
||||
.descriptor
|
||||
.id
|
||||
.as_str(),
|
||||
"test-provider"
|
||||
);
|
||||
}
|
||||
}
|
||||
541
vendor/mentra-provider/src/request.rs
vendored
Normal file
541
vendor/mentra-provider/src/request.rs
vendored
Normal file
@@ -0,0 +1,541 @@
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::borrow::Cow;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::ContentBlock;
|
||||
use crate::Message;
|
||||
use crate::ProviderError;
|
||||
use crate::model::ToolChoice;
|
||||
use crate::tool::ToolSpec;
|
||||
|
||||
/// Provider-neutral reasoning controls supported across multiple providers.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ReasoningOptions {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub effort: Option<ReasoningEffort>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub summary: Option<ReasoningSummary>,
|
||||
}
|
||||
|
||||
/// Shared reasoning effort levels supported by Mentra's public API.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
#[non_exhaustive]
|
||||
pub enum ReasoningEffort {
|
||||
Low,
|
||||
Medium,
|
||||
High,
|
||||
#[serde(rename = "xhigh")]
|
||||
XHigh,
|
||||
Max,
|
||||
}
|
||||
|
||||
/// Shared reasoning summary levels used by Responses-family providers.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ReasoningSummary {
|
||||
Auto,
|
||||
Concise,
|
||||
Detailed,
|
||||
}
|
||||
|
||||
/// Provider-neutral tool search behavior requested for a model call.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ToolSearchMode {
|
||||
#[default]
|
||||
Disabled,
|
||||
Hosted,
|
||||
}
|
||||
|
||||
/// Responses-compatible verbosity controls for text output.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ResponsesVerbosity {
|
||||
Low,
|
||||
#[default]
|
||||
Medium,
|
||||
High,
|
||||
}
|
||||
|
||||
/// Transport-level request compression supported by Responses-family HTTP calls.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ResponsesRequestCompression {
|
||||
#[default]
|
||||
None,
|
||||
Zstd,
|
||||
}
|
||||
|
||||
/// Provider-side conversation state strategy for Responses-family providers.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ResponsesStateMode {
|
||||
/// Send the complete local transcript and do not attach provider-side state.
|
||||
ReplayOnly,
|
||||
/// Keep local replay as the source of truth while opportunistically chaining provider state.
|
||||
///
|
||||
/// An unknown HTTP endpoint may receive one capability probe before Hybrid
|
||||
/// learns that `previous_response_id` is unsupported. Hosts that already
|
||||
/// know this can disable the probe on [`crate::responses::ResponsesProvider`].
|
||||
#[default]
|
||||
Hybrid,
|
||||
/// Require provider-side state chaining once a previous response id is available.
|
||||
Stateful,
|
||||
}
|
||||
|
||||
impl ResponsesStateMode {
|
||||
pub fn uses_provider_state(self) -> bool {
|
||||
matches!(self, Self::Hybrid | Self::Stateful)
|
||||
}
|
||||
}
|
||||
|
||||
/// Streaming transport used by Responses-family providers.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ResponsesTransport {
|
||||
/// Standard HTTP request with Server-Sent Events response streaming.
|
||||
#[default]
|
||||
HttpSse,
|
||||
/// Long-lived WebSocket connection driven by `response.create` frames.
|
||||
WebSocket,
|
||||
}
|
||||
|
||||
/// Responses-compatible format discriminator for structured text output.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ResponsesTextFormatType {
|
||||
#[default]
|
||||
JsonSchema,
|
||||
}
|
||||
|
||||
/// Structured text output format controls for Responses-family providers.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct ResponsesTextFormat {
|
||||
#[serde(default)]
|
||||
pub r#type: ResponsesTextFormatType,
|
||||
#[serde(default)]
|
||||
pub strict: bool,
|
||||
pub schema: serde_json::Value,
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
/// Responses-compatible text controls combining verbosity and output schemas.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct ResponsesTextControls {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub verbosity: Option<ResponsesVerbosity>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub format: Option<ResponsesTextFormat>,
|
||||
}
|
||||
|
||||
/// Shared Responses-family request options.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ResponsesRequestOptions {
|
||||
#[serde(default)]
|
||||
pub parallel_tool_calls: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub previous_response_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub state_mode: ResponsesStateMode,
|
||||
#[serde(default)]
|
||||
pub transport: ResponsesTransport,
|
||||
#[serde(default)]
|
||||
pub store: Option<bool>,
|
||||
#[serde(default)]
|
||||
pub stream: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub include: Vec<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub service_tier: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_cache_key: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub text: Option<ResponsesTextControls>,
|
||||
#[serde(default)]
|
||||
pub compression: ResponsesRequestCompression,
|
||||
}
|
||||
|
||||
impl Default for ResponsesRequestOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
parallel_tool_calls: None,
|
||||
previous_response_id: None,
|
||||
state_mode: ResponsesStateMode::Hybrid,
|
||||
transport: ResponsesTransport::HttpSse,
|
||||
store: None,
|
||||
stream: Some(true),
|
||||
include: Vec::new(),
|
||||
service_tier: None,
|
||||
prompt_cache_key: None,
|
||||
text: None,
|
||||
compression: ResponsesRequestCompression::None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Anthropic-specific request options.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct AnthropicRequestOptions {
|
||||
#[serde(default)]
|
||||
pub disable_parallel_tool_use: Option<bool>,
|
||||
}
|
||||
|
||||
/// Gemini-specific request options.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct GeminiRequestOptions {
|
||||
#[serde(default)]
|
||||
pub thoughts: Option<bool>,
|
||||
}
|
||||
|
||||
/// Provider-neutral session metadata and affinity hints.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct SessionRequestOptions {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub sticky_turn_state: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub turn_metadata: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub subagent: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prefer_connection_reuse: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub session_affinity: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
|
||||
pub extra_headers: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
/// Provider-specific request options that should be forwarded on the wire.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct ProviderRequestOptions {
|
||||
#[serde(default)]
|
||||
pub tool_search_mode: ToolSearchMode,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning: Option<ReasoningOptions>,
|
||||
#[serde(default)]
|
||||
pub responses: ResponsesRequestOptions,
|
||||
#[serde(default)]
|
||||
pub anthropic: AnthropicRequestOptions,
|
||||
#[serde(default)]
|
||||
pub gemini: GeminiRequestOptions,
|
||||
#[serde(default)]
|
||||
pub session: SessionRequestOptions,
|
||||
}
|
||||
|
||||
/// Provider request assembled by the runtime before dispatch.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Request<'a> {
|
||||
pub model: Cow<'a, str>,
|
||||
pub system: Option<Cow<'a, str>>,
|
||||
pub messages: Cow<'a, [Message]>,
|
||||
pub tools: Cow<'a, [ToolSpec]>,
|
||||
pub tool_choice: Option<ToolChoice>,
|
||||
pub temperature: Option<f32>,
|
||||
pub max_output_tokens: Option<u32>,
|
||||
pub metadata: Cow<'a, BTreeMap<String, String>>,
|
||||
pub provider_request_options: ProviderRequestOptions,
|
||||
}
|
||||
|
||||
impl Request<'_> {
|
||||
pub fn into_owned(self) -> Request<'static> {
|
||||
Request {
|
||||
model: Cow::Owned(self.model.into_owned()),
|
||||
system: self.system.map(|system| Cow::Owned(system.into_owned())),
|
||||
messages: Cow::Owned(self.messages.into_owned()),
|
||||
tools: Cow::Owned(self.tools.into_owned()),
|
||||
tool_choice: self.tool_choice,
|
||||
temperature: self.temperature,
|
||||
max_output_tokens: self.max_output_tokens,
|
||||
metadata: Cow::Owned(self.metadata.into_owned()),
|
||||
provider_request_options: self.provider_request_options,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider-neutral transcript item used for history compaction.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum CompactionInputItem {
|
||||
UserTurn {
|
||||
content: String,
|
||||
},
|
||||
AssistantTurn {
|
||||
content: String,
|
||||
},
|
||||
ToolExchange {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
request: Option<String>,
|
||||
result: String,
|
||||
is_error: bool,
|
||||
},
|
||||
CanonicalContext {
|
||||
content: String,
|
||||
},
|
||||
MemoryRecall {
|
||||
content: String,
|
||||
},
|
||||
DelegationResult {
|
||||
agent_id: String,
|
||||
agent_name: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
role: Option<String>,
|
||||
status: String,
|
||||
content: String,
|
||||
},
|
||||
CompactionSummary {
|
||||
content: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Provider-neutral request assembled for history compaction.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompactionRequest<'a> {
|
||||
pub model: Cow<'a, str>,
|
||||
pub instructions: Cow<'a, str>,
|
||||
pub input: Cow<'a, [CompactionInputItem]>,
|
||||
pub metadata: Cow<'a, BTreeMap<String, String>>,
|
||||
pub provider_request_options: ProviderRequestOptions,
|
||||
}
|
||||
|
||||
impl CompactionRequest<'_> {
|
||||
/// Converts a compaction request into an ordinary model request.
|
||||
pub fn into_model_request(self) -> Result<Request<'static>, ProviderError> {
|
||||
let input_json =
|
||||
serde_json::to_string(self.input.as_ref()).map_err(ProviderError::Serialize)?;
|
||||
|
||||
Ok(Request {
|
||||
model: Cow::Owned(self.model.into_owned()),
|
||||
system: Some(Cow::Owned(self.instructions.into_owned())),
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text(format!(
|
||||
"Compaction input JSON:\n{input_json}"
|
||||
)))]),
|
||||
tools: Cow::Owned(Vec::new()),
|
||||
tool_choice: None,
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(self.metadata.into_owned()),
|
||||
provider_request_options: self.provider_request_options,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonical raw memory payload used by memory summarization requests.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct RawMemory {
|
||||
pub id: String,
|
||||
pub metadata: RawMemoryMetadata,
|
||||
pub items: Vec<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RawMemoryMetadata {
|
||||
pub source_path: String,
|
||||
}
|
||||
|
||||
/// Provider-neutral request assembled for trace memory summarization.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MemorySummarizeRequest<'a> {
|
||||
pub model: Cow<'a, str>,
|
||||
pub raw_memories: Cow<'a, [RawMemory]>,
|
||||
pub reasoning: Option<ReasoningOptions>,
|
||||
pub metadata: Cow<'a, BTreeMap<String, String>>,
|
||||
pub provider_request_options: ProviderRequestOptions,
|
||||
}
|
||||
|
||||
impl MemorySummarizeRequest<'_> {
|
||||
/// Converts a memory summarize request into an ordinary model request.
|
||||
pub fn into_model_request(self) -> Result<Request<'static>, ProviderError> {
|
||||
let raw_memories_json =
|
||||
serde_json::to_string(self.raw_memories.as_ref()).map_err(ProviderError::Serialize)?;
|
||||
|
||||
Ok(Request {
|
||||
model: Cow::Owned(self.model.into_owned()),
|
||||
system: Some(Cow::Borrowed(MEMORY_SUMMARIZE_SYSTEM_PROMPT)),
|
||||
messages: Cow::Owned(vec![Message::user(ContentBlock::text(format!(
|
||||
"Memory summarize input JSON:\n{raw_memories_json}"
|
||||
)))]),
|
||||
tools: Cow::Owned(Vec::new()),
|
||||
tool_choice: None,
|
||||
temperature: None,
|
||||
max_output_tokens: None,
|
||||
metadata: Cow::Owned(self.metadata.into_owned()),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
reasoning: self.reasoning,
|
||||
..self.provider_request_options
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
const MEMORY_SUMMARIZE_SYSTEM_PROMPT: &str = concat!(
|
||||
"You summarize trace memories for Codex.\n",
|
||||
"Return valid JSON only.\n",
|
||||
"The output must be a JSON array with one object per input trace, in the same order.\n",
|
||||
"Each object must have exactly these string fields: `raw_memory` and `memory_summary`.\n",
|
||||
"`raw_memory` should be a concrete, detailed summary of the trace contents.\n",
|
||||
"`memory_summary` should be a shorter durable takeaway focused on reusable context.\n",
|
||||
"Use empty strings when information is unavailable.\n",
|
||||
"Do not include markdown fences or extra commentary.\n",
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::Value;
|
||||
|
||||
#[test]
|
||||
fn reasoning_effort_uses_the_five_public_spellings() {
|
||||
for (effort, expected) in [
|
||||
(ReasoningEffort::Low, "low"),
|
||||
(ReasoningEffort::Medium, "medium"),
|
||||
(ReasoningEffort::High, "high"),
|
||||
(ReasoningEffort::XHigh, "xhigh"),
|
||||
(ReasoningEffort::Max, "max"),
|
||||
] {
|
||||
let serialized = serde_json::to_string(&effort).expect("effort should serialize");
|
||||
assert_eq!(serialized, format!("\"{expected}\""));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<ReasoningEffort>(&serialized)
|
||||
.expect("effort should deserialize"),
|
||||
effort
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compaction_request_into_model_request_serializes_input_as_prompt_text() {
|
||||
let request = CompactionRequest {
|
||||
model: Cow::Borrowed("gpt-5"),
|
||||
instructions: Cow::Borrowed("Summarize the transcript."),
|
||||
input: Cow::Owned(vec![
|
||||
CompactionInputItem::UserTurn {
|
||||
content: "hello".to_string(),
|
||||
},
|
||||
CompactionInputItem::AssistantTurn {
|
||||
content: "world".to_string(),
|
||||
},
|
||||
]),
|
||||
metadata: Cow::Owned(BTreeMap::from([("scope".to_string(), "test".to_string())])),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
session: SessionRequestOptions {
|
||||
sticky_turn_state: Some("sticky".to_string()),
|
||||
turn_metadata: None,
|
||||
subagent: Some("compact".to_string()),
|
||||
prefer_connection_reuse: Some(true),
|
||||
session_affinity: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
},
|
||||
..ProviderRequestOptions::default()
|
||||
},
|
||||
};
|
||||
|
||||
let model_request = request
|
||||
.into_model_request()
|
||||
.expect("compaction request should convert");
|
||||
|
||||
assert_eq!(model_request.model.as_ref(), "gpt-5");
|
||||
assert_eq!(
|
||||
model_request.system.as_deref(),
|
||||
Some("Summarize the transcript.")
|
||||
);
|
||||
assert_eq!(model_request.metadata["scope"], "test");
|
||||
assert_eq!(
|
||||
model_request
|
||||
.provider_request_options
|
||||
.session
|
||||
.sticky_turn_state
|
||||
.as_deref(),
|
||||
Some("sticky")
|
||||
);
|
||||
assert_eq!(
|
||||
model_request
|
||||
.provider_request_options
|
||||
.session
|
||||
.subagent
|
||||
.as_deref(),
|
||||
Some("compact")
|
||||
);
|
||||
assert_eq!(model_request.messages.len(), 1);
|
||||
|
||||
let prompt = model_request.messages[0].text();
|
||||
assert!(prompt.starts_with("Compaction input JSON:\n"));
|
||||
let payload = prompt
|
||||
.strip_prefix("Compaction input JSON:\n")
|
||||
.expect("prompt should contain the compaction prefix");
|
||||
let input: Vec<Value> = serde_json::from_str(payload).expect("prompt should be json");
|
||||
assert_eq!(input[0]["type"], "user_turn");
|
||||
assert_eq!(input[0]["content"], "hello");
|
||||
assert_eq!(input[1]["type"], "assistant_turn");
|
||||
assert_eq!(input[1]["content"], "world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_summarize_request_into_model_request_serializes_input_as_prompt_text() {
|
||||
let request = MemorySummarizeRequest {
|
||||
model: Cow::Borrowed("gpt-5"),
|
||||
raw_memories: Cow::Owned(vec![RawMemory {
|
||||
id: "memory-1".to_string(),
|
||||
metadata: RawMemoryMetadata {
|
||||
source_path: "/tmp/trace.jsonl".to_string(),
|
||||
},
|
||||
items: vec![serde_json::json!({"type":"message","role":"user"})],
|
||||
}]),
|
||||
reasoning: Some(ReasoningOptions {
|
||||
effort: Some(ReasoningEffort::Medium),
|
||||
summary: None,
|
||||
}),
|
||||
metadata: Cow::Owned(BTreeMap::from([("scope".to_string(), "test".to_string())])),
|
||||
provider_request_options: ProviderRequestOptions {
|
||||
session: SessionRequestOptions {
|
||||
sticky_turn_state: None,
|
||||
turn_metadata: Some("{\"turn_id\":\"t1\"}".to_string()),
|
||||
subagent: None,
|
||||
prefer_connection_reuse: Some(true),
|
||||
session_affinity: Some("thread-1".to_string()),
|
||||
extra_headers: BTreeMap::new(),
|
||||
},
|
||||
..ProviderRequestOptions::default()
|
||||
},
|
||||
};
|
||||
|
||||
let model_request = request
|
||||
.into_model_request()
|
||||
.expect("memory summarize request should convert");
|
||||
|
||||
assert_eq!(model_request.model.as_ref(), "gpt-5");
|
||||
assert_eq!(
|
||||
model_request.system.as_deref(),
|
||||
Some(MEMORY_SUMMARIZE_SYSTEM_PROMPT)
|
||||
);
|
||||
assert_eq!(model_request.metadata["scope"], "test");
|
||||
assert_eq!(
|
||||
model_request
|
||||
.provider_request_options
|
||||
.session
|
||||
.turn_metadata
|
||||
.as_deref(),
|
||||
Some("{\"turn_id\":\"t1\"}")
|
||||
);
|
||||
assert_eq!(
|
||||
model_request
|
||||
.provider_request_options
|
||||
.reasoning
|
||||
.as_ref()
|
||||
.expect("reasoning options")
|
||||
.effort,
|
||||
Some(ReasoningEffort::Medium)
|
||||
);
|
||||
assert_eq!(model_request.messages.len(), 1);
|
||||
|
||||
let prompt = model_request.messages[0].text();
|
||||
assert!(prompt.starts_with("Memory summarize input JSON:\n"));
|
||||
let payload = prompt
|
||||
.strip_prefix("Memory summarize input JSON:\n")
|
||||
.expect("prompt should contain the memory summarize prefix");
|
||||
let input: Vec<RawMemory> = serde_json::from_str(payload).expect("prompt should be json");
|
||||
assert_eq!(input[0].id, "memory-1");
|
||||
assert_eq!(input[0].metadata.source_path, "/tmp/trace.jsonl");
|
||||
assert_eq!(input[0].items[0]["role"], "user");
|
||||
}
|
||||
}
|
||||
913
vendor/mentra-provider/src/response.rs
vendored
Normal file
913
vendor/mentra-provider/src/response.rs
vendored
Normal file
@@ -0,0 +1,913 @@
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::error::ProviderError;
|
||||
use crate::model::ContentBlock;
|
||||
use crate::model::HostedToolSearchCall;
|
||||
use crate::model::HostedWebSearchCall;
|
||||
use crate::model::ImageGenerationCall;
|
||||
use crate::model::ReasoningProvenance;
|
||||
use crate::model::Role;
|
||||
use crate::model::TokenUsage;
|
||||
use crate::model::ToolResultContent;
|
||||
use crate::request::CompactionInputItem;
|
||||
use crate::stream::ContentBlockDelta;
|
||||
use crate::stream::ContentBlockStart;
|
||||
use crate::stream::ProviderEvent;
|
||||
use crate::stream::ProviderEventStream;
|
||||
|
||||
/// A complete response collected from a provider stream.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct Response {
|
||||
pub id: String,
|
||||
pub model: String,
|
||||
pub role: Role,
|
||||
pub content: Vec<ContentBlock>,
|
||||
pub stop_reason: Option<String>,
|
||||
pub usage: Option<TokenUsage>,
|
||||
}
|
||||
|
||||
/// A complete history-compaction response collected from a provider.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CompactionResponse {
|
||||
pub output: Vec<CompactionInputItem>,
|
||||
}
|
||||
|
||||
/// A complete memory summarize response collected from a provider.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct MemorySummarizeResponse {
|
||||
pub output: Vec<MemorySummarizeOutput>,
|
||||
}
|
||||
|
||||
/// One summary object for a single raw memory input.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct MemorySummarizeOutput {
|
||||
#[serde(rename = "trace_summary", alias = "raw_memory")]
|
||||
pub raw_memory: String,
|
||||
pub memory_summary: String,
|
||||
}
|
||||
|
||||
/// Rebuilds a full response from a provider event stream.
|
||||
pub async fn collect_response_from_stream(
|
||||
mut stream: ProviderEventStream,
|
||||
) -> Result<Response, ProviderError> {
|
||||
let mut builder = StreamingResponseBuilder::default();
|
||||
|
||||
while let Some(event) = stream.recv().await {
|
||||
builder.apply(event?)?;
|
||||
}
|
||||
|
||||
builder.build()
|
||||
}
|
||||
|
||||
/// Converts a response into a provider event stream.
|
||||
pub fn provider_event_stream_from_response(response: Response) -> ProviderEventStream {
|
||||
let events = response.into_provider_events();
|
||||
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
|
||||
for event in events {
|
||||
if tx.send(Ok(event)).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
rx
|
||||
}
|
||||
|
||||
impl Response {
|
||||
pub fn into_provider_events(self) -> Vec<ProviderEvent> {
|
||||
let mut events = vec![ProviderEvent::MessageStarted {
|
||||
id: self.id,
|
||||
model: self.model,
|
||||
role: self.role,
|
||||
}];
|
||||
|
||||
for (index, block) in self.content.into_iter().enumerate() {
|
||||
events.extend(block.into_provider_events(index));
|
||||
}
|
||||
|
||||
events.push(ProviderEvent::MessageDelta {
|
||||
stop_reason: self.stop_reason,
|
||||
usage: self.usage,
|
||||
});
|
||||
events.push(ProviderEvent::MessageStopped);
|
||||
events
|
||||
}
|
||||
|
||||
/// Converts a normal response into a compaction response by collecting its text output.
|
||||
pub fn into_compaction_response(self) -> CompactionResponse {
|
||||
let text = self
|
||||
.content
|
||||
.into_iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::Text { text } => Some(text),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
CompactionResponse::from_text(text)
|
||||
}
|
||||
|
||||
/// Converts a normal response into a memory summarize response by parsing its text output.
|
||||
pub fn into_memory_summarize_response(self) -> Result<MemorySummarizeResponse, ProviderError> {
|
||||
let text = self
|
||||
.content
|
||||
.into_iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::Text { text } => Some(text),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
MemorySummarizeResponse::from_text(&text)
|
||||
}
|
||||
}
|
||||
|
||||
impl CompactionResponse {
|
||||
/// Wraps a text summary into a single compaction-summary item.
|
||||
pub fn from_text(text: impl Into<String>) -> Self {
|
||||
Self {
|
||||
output: vec![CompactionInputItem::CompactionSummary {
|
||||
content: text.into(),
|
||||
}],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MemorySummarizeResponse {
|
||||
pub fn from_text(text: &str) -> Result<Self, ProviderError> {
|
||||
let text = strip_markdown_code_fence(text);
|
||||
|
||||
if let Ok(output) = serde_json::from_str::<Vec<MemorySummarizeOutput>>(text) {
|
||||
return Ok(Self { output });
|
||||
}
|
||||
|
||||
if let Ok(response) = serde_json::from_str::<MemorySummarizeResponse>(text) {
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
Err(ProviderError::InvalidResponse(
|
||||
"failed to parse memory summarize response json".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_markdown_code_fence(text: &str) -> &str {
|
||||
let trimmed = text.trim();
|
||||
let Some(rest) = trimmed.strip_prefix("```") else {
|
||||
return trimmed;
|
||||
};
|
||||
let rest = rest
|
||||
.strip_prefix("json")
|
||||
.or_else(|| rest.strip_prefix("JSON"))
|
||||
.unwrap_or(rest);
|
||||
rest.trim()
|
||||
.strip_suffix("```")
|
||||
.map(str::trim)
|
||||
.unwrap_or(trimmed)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct StreamingResponseBuilder {
|
||||
id: Option<String>,
|
||||
model: Option<String>,
|
||||
role: Option<Role>,
|
||||
blocks: std::collections::BTreeMap<usize, StreamingContentBlock>,
|
||||
stop_reason: Option<String>,
|
||||
usage: Option<TokenUsage>,
|
||||
stopped: bool,
|
||||
}
|
||||
|
||||
impl StreamingResponseBuilder {
|
||||
fn apply(&mut self, event: ProviderEvent) -> Result<(), ProviderError> {
|
||||
match event {
|
||||
ProviderEvent::ResponseHeaders(_) | ProviderEvent::ResponseCreated => {}
|
||||
ProviderEvent::MessageStarted { id, model, role } => {
|
||||
self.id = Some(id);
|
||||
self.model = Some(model);
|
||||
self.role = Some(role);
|
||||
}
|
||||
ProviderEvent::ContentBlockStarted { index, kind } => {
|
||||
self.blocks.insert(index, StreamingContentBlock::from(kind));
|
||||
}
|
||||
ProviderEvent::ContentBlockDelta { index, delta } => {
|
||||
let block = self.blocks.get_mut(&index).ok_or_else(|| {
|
||||
ProviderError::MalformedStream(format!(
|
||||
"content block delta received before start for index {index}"
|
||||
))
|
||||
})?;
|
||||
block.apply_delta(delta)?;
|
||||
}
|
||||
ProviderEvent::ContentBlockStopped { index } => {
|
||||
let block = self.blocks.get_mut(&index).ok_or_else(|| {
|
||||
ProviderError::MalformedStream(format!(
|
||||
"content block stop received before start for index {index}"
|
||||
))
|
||||
})?;
|
||||
block.mark_complete();
|
||||
}
|
||||
ProviderEvent::MessageDelta { stop_reason, usage } => {
|
||||
self.stop_reason = stop_reason;
|
||||
self.usage = usage;
|
||||
}
|
||||
ProviderEvent::ReasoningSummaryDelta { .. }
|
||||
| ProviderEvent::ReasoningContentDelta { .. }
|
||||
| ProviderEvent::ReasoningSummaryPartAdded { .. } => {}
|
||||
ProviderEvent::MessageStopped => {
|
||||
self.stopped = true;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn build(self) -> Result<Response, ProviderError> {
|
||||
if !self.stopped {
|
||||
return Err(ProviderError::MalformedStream(
|
||||
"message stream ended before MessageStopped".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let id = self
|
||||
.id
|
||||
.ok_or_else(|| ProviderError::MalformedStream("missing message id".to_string()))?;
|
||||
let model = self
|
||||
.model
|
||||
.ok_or_else(|| ProviderError::MalformedStream("missing model id".to_string()))?;
|
||||
let role = self
|
||||
.role
|
||||
.ok_or_else(|| ProviderError::MalformedStream("missing message role".to_string()))?;
|
||||
let mut content = Vec::with_capacity(self.blocks.len());
|
||||
|
||||
for (index, block) in self.blocks {
|
||||
if !block.is_complete() {
|
||||
return Err(ProviderError::MalformedStream(format!(
|
||||
"content block {index} did not complete"
|
||||
)));
|
||||
}
|
||||
content.push(block.try_into_content_block()?);
|
||||
}
|
||||
|
||||
Ok(Response {
|
||||
id,
|
||||
model,
|
||||
role,
|
||||
content,
|
||||
stop_reason: self.stop_reason,
|
||||
usage: self.usage,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
enum StreamingContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
complete: bool,
|
||||
},
|
||||
Thinking {
|
||||
thinking: String,
|
||||
signature: Option<String>,
|
||||
encrypted_content: Option<String>,
|
||||
id: Option<String>,
|
||||
provenance: Option<ReasoningProvenance>,
|
||||
redacted: bool,
|
||||
complete: bool,
|
||||
},
|
||||
Image {
|
||||
source: crate::model::ImageSource,
|
||||
complete: bool,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
input_json: String,
|
||||
complete: bool,
|
||||
},
|
||||
ToolResult {
|
||||
tool_use_id: String,
|
||||
content: Option<ToolResultContent>,
|
||||
is_error: bool,
|
||||
complete: bool,
|
||||
},
|
||||
HostedToolSearch {
|
||||
call: HostedToolSearchCall,
|
||||
complete: bool,
|
||||
},
|
||||
HostedWebSearch {
|
||||
call: HostedWebSearchCall,
|
||||
complete: bool,
|
||||
},
|
||||
ImageGeneration {
|
||||
call: ImageGenerationCall,
|
||||
complete: bool,
|
||||
},
|
||||
}
|
||||
|
||||
impl StreamingContentBlock {
|
||||
fn apply_delta(&mut self, delta: ContentBlockDelta) -> Result<(), ProviderError> {
|
||||
match (self, delta) {
|
||||
(StreamingContentBlock::Text { text, .. }, ContentBlockDelta::Text(delta)) => {
|
||||
text.push_str(&delta);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::Thinking { thinking, .. },
|
||||
ContentBlockDelta::ThinkingText(delta),
|
||||
) => {
|
||||
thinking.push_str(&delta);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::Thinking { signature, .. },
|
||||
ContentBlockDelta::ThinkingSignature(delta),
|
||||
) => {
|
||||
signature.get_or_insert_with(String::new).push_str(&delta);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::Thinking {
|
||||
encrypted_content, ..
|
||||
},
|
||||
ContentBlockDelta::ThinkingEncryptedContent(value),
|
||||
) => {
|
||||
*encrypted_content = Some(value);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::ToolUse { input_json, .. },
|
||||
ContentBlockDelta::ToolUseInputJson(delta),
|
||||
) => {
|
||||
input_json.push_str(&delta);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::ToolResult { content, .. },
|
||||
ContentBlockDelta::ToolResultContent(delta),
|
||||
) => {
|
||||
*content = Some(match (content.take(), delta) {
|
||||
(_, ToolResultContent::Structured(value)) => {
|
||||
ToolResultContent::Structured(value)
|
||||
}
|
||||
(
|
||||
Some(ToolResultContent::Structured(value)),
|
||||
ToolResultContent::Text(delta),
|
||||
) => ToolResultContent::Structured(merge_structured_text(value, delta)),
|
||||
(Some(ToolResultContent::Text(existing)), ToolResultContent::Text(delta)) => {
|
||||
ToolResultContent::Text(format!("{existing}{delta}"))
|
||||
}
|
||||
(None, delta) => delta,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::HostedToolSearch { call, .. },
|
||||
ContentBlockDelta::HostedToolSearchQuery(delta),
|
||||
) => {
|
||||
let query = call.query.get_or_insert_with(String::new);
|
||||
query.push_str(&delta);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::HostedToolSearch { call, .. },
|
||||
ContentBlockDelta::HostedToolSearchStatus(status),
|
||||
) => {
|
||||
call.status = Some(status);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::HostedWebSearch { call, .. },
|
||||
ContentBlockDelta::HostedWebSearchAction(action),
|
||||
) => {
|
||||
call.action = Some(action);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::HostedWebSearch { call, .. },
|
||||
ContentBlockDelta::HostedWebSearchStatus(status),
|
||||
) => {
|
||||
call.status = Some(status);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::ImageGeneration { call, .. },
|
||||
ContentBlockDelta::ImageGenerationStatus(status),
|
||||
) => {
|
||||
call.status = status;
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::ImageGeneration { call, .. },
|
||||
ContentBlockDelta::ImageGenerationRevisedPrompt(delta),
|
||||
) => {
|
||||
let revised_prompt = call.revised_prompt.get_or_insert_with(String::new);
|
||||
revised_prompt.push_str(&delta);
|
||||
Ok(())
|
||||
}
|
||||
(
|
||||
StreamingContentBlock::ImageGeneration { call, .. },
|
||||
ContentBlockDelta::ImageGenerationResult(result),
|
||||
) => {
|
||||
call.result = Some(result);
|
||||
Ok(())
|
||||
}
|
||||
(block, delta) => Err(ProviderError::MalformedStream(format!(
|
||||
"delta {delta:?} is not valid for block {}",
|
||||
block.kind_name()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn mark_complete(&mut self) {
|
||||
match self {
|
||||
StreamingContentBlock::Text { complete, .. }
|
||||
| StreamingContentBlock::Thinking { complete, .. }
|
||||
| StreamingContentBlock::Image { complete, .. }
|
||||
| StreamingContentBlock::ToolUse { complete, .. }
|
||||
| StreamingContentBlock::ToolResult { complete, .. }
|
||||
| StreamingContentBlock::HostedToolSearch { complete, .. }
|
||||
| StreamingContentBlock::HostedWebSearch { complete, .. }
|
||||
| StreamingContentBlock::ImageGeneration { complete, .. } => *complete = true,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_complete(&self) -> bool {
|
||||
match self {
|
||||
StreamingContentBlock::Text { complete, .. }
|
||||
| StreamingContentBlock::Thinking { complete, .. }
|
||||
| StreamingContentBlock::Image { complete, .. }
|
||||
| StreamingContentBlock::ToolUse { complete, .. }
|
||||
| StreamingContentBlock::ToolResult { complete, .. }
|
||||
| StreamingContentBlock::HostedToolSearch { complete, .. }
|
||||
| StreamingContentBlock::HostedWebSearch { complete, .. }
|
||||
| StreamingContentBlock::ImageGeneration { complete, .. } => *complete,
|
||||
}
|
||||
}
|
||||
|
||||
fn try_into_content_block(self) -> Result<ContentBlock, ProviderError> {
|
||||
match self {
|
||||
StreamingContentBlock::Text { text, .. } => Ok(ContentBlock::Text { text }),
|
||||
StreamingContentBlock::Thinking {
|
||||
thinking,
|
||||
signature,
|
||||
encrypted_content,
|
||||
id,
|
||||
provenance,
|
||||
redacted,
|
||||
..
|
||||
} => Ok(ContentBlock::Thinking {
|
||||
thinking,
|
||||
signature,
|
||||
encrypted_content,
|
||||
id,
|
||||
provenance,
|
||||
redacted,
|
||||
}),
|
||||
StreamingContentBlock::Image { source, .. } => Ok(ContentBlock::Image { source }),
|
||||
StreamingContentBlock::ToolUse {
|
||||
id,
|
||||
name,
|
||||
input_json,
|
||||
..
|
||||
} => Ok(ContentBlock::ToolUse {
|
||||
id,
|
||||
name,
|
||||
input: serde_json::from_str(&input_json).map_err(ProviderError::Deserialize)?,
|
||||
}),
|
||||
StreamingContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
..
|
||||
} => Ok(ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content: content.unwrap_or_else(|| ToolResultContent::text(String::new())),
|
||||
is_error,
|
||||
}),
|
||||
StreamingContentBlock::HostedToolSearch { call, .. } => {
|
||||
Ok(ContentBlock::HostedToolSearch { call })
|
||||
}
|
||||
StreamingContentBlock::HostedWebSearch { call, .. } => {
|
||||
Ok(ContentBlock::HostedWebSearch { call })
|
||||
}
|
||||
StreamingContentBlock::ImageGeneration { call, .. } => {
|
||||
Ok(ContentBlock::ImageGeneration { call })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn kind_name(&self) -> &'static str {
|
||||
match self {
|
||||
StreamingContentBlock::Text { .. } => "text",
|
||||
StreamingContentBlock::Thinking { .. } => "thinking",
|
||||
StreamingContentBlock::Image { .. } => "image",
|
||||
StreamingContentBlock::ToolUse { .. } => "tool_use",
|
||||
StreamingContentBlock::ToolResult { .. } => "tool_result",
|
||||
StreamingContentBlock::HostedToolSearch { .. } => "hosted_tool_search",
|
||||
StreamingContentBlock::HostedWebSearch { .. } => "hosted_web_search",
|
||||
StreamingContentBlock::ImageGeneration { .. } => "image_generation",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ContentBlockStart> for StreamingContentBlock {
|
||||
fn from(value: ContentBlockStart) -> Self {
|
||||
match value {
|
||||
ContentBlockStart::Text => StreamingContentBlock::Text {
|
||||
text: String::new(),
|
||||
complete: false,
|
||||
},
|
||||
ContentBlockStart::Thinking {
|
||||
encrypted_content,
|
||||
id,
|
||||
provenance,
|
||||
redacted,
|
||||
} => StreamingContentBlock::Thinking {
|
||||
thinking: String::new(),
|
||||
signature: None,
|
||||
encrypted_content,
|
||||
id,
|
||||
provenance,
|
||||
redacted,
|
||||
complete: false,
|
||||
},
|
||||
ContentBlockStart::Image { source } => StreamingContentBlock::Image {
|
||||
source,
|
||||
complete: false,
|
||||
},
|
||||
ContentBlockStart::ToolUse { id, name } => StreamingContentBlock::ToolUse {
|
||||
id,
|
||||
name,
|
||||
input_json: String::new(),
|
||||
complete: false,
|
||||
},
|
||||
ContentBlockStart::ToolResult {
|
||||
tool_use_id,
|
||||
is_error,
|
||||
content,
|
||||
} => StreamingContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
complete: false,
|
||||
},
|
||||
ContentBlockStart::HostedToolSearch { call } => {
|
||||
StreamingContentBlock::HostedToolSearch {
|
||||
call,
|
||||
complete: false,
|
||||
}
|
||||
}
|
||||
ContentBlockStart::HostedWebSearch { call } => StreamingContentBlock::HostedWebSearch {
|
||||
call,
|
||||
complete: false,
|
||||
},
|
||||
ContentBlockStart::ImageGeneration { call } => StreamingContentBlock::ImageGeneration {
|
||||
call,
|
||||
complete: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ContentBlock {
|
||||
fn into_provider_events(self, index: usize) -> Vec<ProviderEvent> {
|
||||
match self {
|
||||
ContentBlock::Text { text } => {
|
||||
let mut events = vec![ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::Text,
|
||||
}];
|
||||
if !text.is_empty() {
|
||||
events.push(ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::Text(text),
|
||||
});
|
||||
}
|
||||
events.push(ProviderEvent::ContentBlockStopped { index });
|
||||
events
|
||||
}
|
||||
ContentBlock::Thinking {
|
||||
thinking,
|
||||
signature,
|
||||
encrypted_content,
|
||||
id,
|
||||
provenance,
|
||||
redacted,
|
||||
} => {
|
||||
let mut events = vec![ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::Thinking {
|
||||
encrypted_content,
|
||||
id,
|
||||
provenance,
|
||||
redacted,
|
||||
},
|
||||
}];
|
||||
if !thinking.is_empty() {
|
||||
events.push(ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::ThinkingText(thinking),
|
||||
});
|
||||
}
|
||||
if let Some(signature) = signature {
|
||||
events.push(ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::ThinkingSignature(signature),
|
||||
});
|
||||
}
|
||||
events.push(ProviderEvent::ContentBlockStopped { index });
|
||||
events
|
||||
}
|
||||
ContentBlock::Image { source } => vec![
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::Image { source },
|
||||
},
|
||||
ProviderEvent::ContentBlockStopped { index },
|
||||
],
|
||||
ContentBlock::ToolUse { id, name, input } => {
|
||||
let mut events = vec![ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::ToolUse { id, name },
|
||||
}];
|
||||
let input_json = input.to_string();
|
||||
if !input_json.is_empty() {
|
||||
events.push(ProviderEvent::ContentBlockDelta {
|
||||
index,
|
||||
delta: ContentBlockDelta::ToolUseInputJson(input_json),
|
||||
});
|
||||
}
|
||||
events.push(ProviderEvent::ContentBlockStopped { index });
|
||||
events
|
||||
}
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
} => {
|
||||
let mut events = vec![ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::ToolResult {
|
||||
tool_use_id,
|
||||
is_error,
|
||||
content: Some(content.clone()),
|
||||
},
|
||||
}];
|
||||
events.push(ProviderEvent::ContentBlockStopped { index });
|
||||
events
|
||||
}
|
||||
ContentBlock::HostedToolSearch { call } => {
|
||||
vec![
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::HostedToolSearch { call },
|
||||
},
|
||||
ProviderEvent::ContentBlockStopped { index },
|
||||
]
|
||||
}
|
||||
ContentBlock::HostedWebSearch { call } => {
|
||||
vec![
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::HostedWebSearch { call },
|
||||
},
|
||||
ProviderEvent::ContentBlockStopped { index },
|
||||
]
|
||||
}
|
||||
ContentBlock::ImageGeneration { call } => {
|
||||
vec![
|
||||
ProviderEvent::ContentBlockStarted {
|
||||
index,
|
||||
kind: ContentBlockStart::ImageGeneration { call },
|
||||
},
|
||||
ProviderEvent::ContentBlockStopped { index },
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_structured_text(value: serde_json::Value, delta: String) -> serde_json::Value {
|
||||
match value {
|
||||
serde_json::Value::String(existing) => {
|
||||
serde_json::Value::String(format!("{existing}{delta}"))
|
||||
}
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_round_trip_preserves_usage() {
|
||||
let response = Response {
|
||||
id: "resp-1".to_string(),
|
||||
model: "model".to_string(),
|
||||
role: Role::Assistant,
|
||||
content: vec![ContentBlock::text("hello")],
|
||||
stop_reason: Some("stop".to_string()),
|
||||
usage: Some(TokenUsage {
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(3),
|
||||
total_tokens: Some(13),
|
||||
cache_read_input_tokens: Some(2),
|
||||
cache_creation_input_tokens: None,
|
||||
reasoning_tokens: Some(1),
|
||||
thoughts_tokens: None,
|
||||
tool_input_tokens: None,
|
||||
}),
|
||||
};
|
||||
|
||||
let rebuilt =
|
||||
collect_response_from_stream(provider_event_stream_from_response(response.clone()))
|
||||
.await
|
||||
.expect("response should rebuild");
|
||||
|
||||
assert_eq!(rebuilt, response);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_round_trip_preserves_signed_and_redacted_thinking() {
|
||||
let provenance = ReasoningProvenance {
|
||||
provider: crate::ProviderId::new("anthropic-edge"),
|
||||
model: "claude-test".to_string(),
|
||||
format: crate::ReasoningFormat::AnthropicSigned,
|
||||
};
|
||||
let response = Response {
|
||||
id: "resp-thinking".to_string(),
|
||||
model: "claude-test".to_string(),
|
||||
role: Role::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Thinking {
|
||||
thinking: "private chain".to_string(),
|
||||
signature: Some("opaque-signature".to_string()),
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(provenance.clone()),
|
||||
redacted: false,
|
||||
},
|
||||
ContentBlock::Thinking {
|
||||
thinking: String::new(),
|
||||
signature: Some("opaque-redacted-data".to_string()),
|
||||
encrypted_content: None,
|
||||
id: None,
|
||||
provenance: Some(provenance),
|
||||
redacted: true,
|
||||
},
|
||||
],
|
||||
stop_reason: Some("end_turn".to_string()),
|
||||
usage: None,
|
||||
};
|
||||
|
||||
let rebuilt =
|
||||
collect_response_from_stream(provider_event_stream_from_response(response.clone()))
|
||||
.await
|
||||
.expect("thinking response should rebuild");
|
||||
|
||||
assert_eq!(rebuilt, response);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_round_trip_preserves_hosted_actions_and_structured_tool_results() {
|
||||
let response = Response {
|
||||
id: "resp-2".to_string(),
|
||||
model: "model".to_string(),
|
||||
role: Role::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id: "call-1".to_string(),
|
||||
content: ToolResultContent::Structured(serde_json::json!({"ok":true})),
|
||||
is_error: false,
|
||||
},
|
||||
ContentBlock::HostedToolSearch {
|
||||
call: HostedToolSearchCall {
|
||||
id: "search-1".to_string(),
|
||||
status: Some("completed".to_string()),
|
||||
query: Some("weather".to_string()),
|
||||
},
|
||||
},
|
||||
ContentBlock::HostedWebSearch {
|
||||
call: HostedWebSearchCall {
|
||||
id: "web-1".to_string(),
|
||||
status: Some("completed".to_string()),
|
||||
action: Some(crate::model::WebSearchAction::Search {
|
||||
query: Some("weather".to_string()),
|
||||
queries: None,
|
||||
}),
|
||||
},
|
||||
},
|
||||
ContentBlock::ImageGeneration {
|
||||
call: ImageGenerationCall {
|
||||
id: "image-1".to_string(),
|
||||
status: "completed".to_string(),
|
||||
revised_prompt: Some("A blue square".to_string()),
|
||||
result: Some(crate::model::ImageGenerationResult::ArtifactRef {
|
||||
artifact_id: "artifact-1".to_string(),
|
||||
}),
|
||||
},
|
||||
},
|
||||
],
|
||||
stop_reason: Some("stop".to_string()),
|
||||
usage: None,
|
||||
};
|
||||
|
||||
let rebuilt =
|
||||
collect_response_from_stream(provider_event_stream_from_response(response.clone()))
|
||||
.await
|
||||
.expect("response should rebuild");
|
||||
|
||||
assert_eq!(rebuilt, response);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_into_compaction_response_collects_text_blocks() {
|
||||
let response = Response {
|
||||
id: "resp-3".to_string(),
|
||||
model: "model".to_string(),
|
||||
role: Role::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::text("first"),
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id: "call-1".to_string(),
|
||||
content: ToolResultContent::text("ignored"),
|
||||
is_error: false,
|
||||
},
|
||||
ContentBlock::text("second"),
|
||||
],
|
||||
stop_reason: None,
|
||||
usage: None,
|
||||
};
|
||||
|
||||
let compaction = response.into_compaction_response();
|
||||
|
||||
assert_eq!(
|
||||
compaction.output,
|
||||
vec![CompactionInputItem::CompactionSummary {
|
||||
content: "first\nsecond".to_string(),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_summarize_response_from_text_accepts_json_array() {
|
||||
let response = MemorySummarizeResponse::from_text(
|
||||
r#"[{"raw_memory":"Detailed summary","memory_summary":"Short summary"}]"#,
|
||||
)
|
||||
.expect("memory summarize response should parse");
|
||||
|
||||
assert_eq!(
|
||||
response,
|
||||
MemorySummarizeResponse {
|
||||
output: vec![MemorySummarizeOutput {
|
||||
raw_memory: "Detailed summary".to_string(),
|
||||
memory_summary: "Short summary".to_string(),
|
||||
}],
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_summarize_response_from_text_accepts_markdown_fence_and_trace_alias() {
|
||||
let response = MemorySummarizeResponse::from_text(
|
||||
"```json\n[{\"trace_summary\":\"Detailed summary\",\"memory_summary\":\"Short summary\"}]\n```",
|
||||
)
|
||||
.expect("memory summarize response should parse");
|
||||
|
||||
assert_eq!(
|
||||
response.output[0],
|
||||
MemorySummarizeOutput {
|
||||
raw_memory: "Detailed summary".to_string(),
|
||||
memory_summary: "Short summary".to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_into_memory_summarize_response_collects_text_content() {
|
||||
let response = Response {
|
||||
id: "resp-4".to_string(),
|
||||
model: "model".to_string(),
|
||||
role: Role::Assistant,
|
||||
content: vec![ContentBlock::text(
|
||||
"[{\"raw_memory\":\"Detailed summary\",\"memory_summary\":\"Short summary\"}]",
|
||||
)],
|
||||
stop_reason: None,
|
||||
usage: None,
|
||||
};
|
||||
|
||||
let summarize = response
|
||||
.into_memory_summarize_response()
|
||||
.expect("memory summarize response should parse");
|
||||
|
||||
assert_eq!(summarize.output.len(), 1);
|
||||
assert_eq!(summarize.output[0].raw_memory, "Detailed summary");
|
||||
assert_eq!(summarize.output[0].memory_summary, "Short summary");
|
||||
}
|
||||
}
|
||||
354
vendor/mentra-provider/src/responses.rs
vendored
Normal file
354
vendor/mentra-provider/src/responses.rs
vendored
Normal file
@@ -0,0 +1,354 @@
|
||||
pub mod model;
|
||||
pub mod session;
|
||||
pub mod sse;
|
||||
/// The `response.create` websocket transport. Compiled in with the
|
||||
/// `responses-websocket` feature; without it, a request that selects
|
||||
/// [`ResponsesTransport::WebSocket`](crate::ResponsesTransport::WebSocket)
|
||||
/// fails rather than falling back to HTTP.
|
||||
#[cfg(feature = "responses-websocket")]
|
||||
pub mod websocket;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::AuthScheme;
|
||||
use crate::BuiltinProvider;
|
||||
use crate::CredentialSource;
|
||||
use crate::ModelCatalog;
|
||||
use crate::ModelInfo;
|
||||
use crate::ProviderCapabilities;
|
||||
use crate::ProviderDefinition;
|
||||
use crate::ProviderError;
|
||||
use crate::ProviderSessionFactory;
|
||||
use crate::RegisteredProvider;
|
||||
use crate::RetryPolicy;
|
||||
use crate::StaticCredentialSource;
|
||||
use crate::WireApi;
|
||||
use crate::embedding::EmbeddingModelInfo;
|
||||
use crate::embedding::EmbeddingProvider;
|
||||
use crate::embedding::EmbeddingRequest;
|
||||
use crate::embedding::EmbeddingResponse;
|
||||
|
||||
use self::session::ResponsesEndpointCapabilities;
|
||||
use self::session::ResponsesSession;
|
||||
use self::session::ResponsesSessionState;
|
||||
|
||||
pub(crate) type SharedTurnState = Arc<std::sync::Mutex<Option<String>>>;
|
||||
|
||||
const DEFAULT_OPENAI_BASE_URL: &str = "https://api.openai.com/";
|
||||
const DEFAULT_OPENROUTER_BASE_URL: &str = "https://openrouter.ai/api/";
|
||||
|
||||
pub fn openai(api_key: impl Into<String>) -> ResponsesProvider<StaticCredentialSource> {
|
||||
ResponsesProvider::openai(api_key)
|
||||
}
|
||||
|
||||
pub fn openrouter(api_key: impl Into<String>) -> ResponsesProvider<StaticCredentialSource> {
|
||||
ResponsesProvider::openrouter(api_key)
|
||||
}
|
||||
|
||||
pub fn openai_with_credential_source<C>(credential_source: C) -> ResponsesProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
ResponsesProvider::openai_with_credential_source(credential_source)
|
||||
}
|
||||
|
||||
pub fn openrouter_with_credential_source<C>(credential_source: C) -> ResponsesProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
ResponsesProvider::openrouter_with_credential_source(credential_source)
|
||||
}
|
||||
|
||||
/// Shared Responses-family provider implementation.
|
||||
///
|
||||
/// This type owns the provider definition, credential source, client, and transport state while
|
||||
/// the request mapping and SSE decoding live in the sibling modules.
|
||||
#[derive(Clone)]
|
||||
pub struct ResponsesProvider<C> {
|
||||
definition: ProviderDefinition,
|
||||
credential_source: Arc<C>,
|
||||
client: reqwest::Client,
|
||||
session_state: Arc<ResponsesSessionState>,
|
||||
endpoint_capabilities: Arc<ResponsesEndpointCapabilities>,
|
||||
hybrid_http_previous_response_id: bool,
|
||||
}
|
||||
|
||||
impl<C> ResponsesProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
pub fn new(definition: ProviderDefinition, credential_source: C) -> Self {
|
||||
Self::with_shared_credential_source(definition, Arc::new(credential_source))
|
||||
}
|
||||
|
||||
pub fn with_shared_credential_source(
|
||||
definition: ProviderDefinition,
|
||||
credential_source: Arc<C>,
|
||||
) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.build()
|
||||
.expect("failed to build reqwest client");
|
||||
Self {
|
||||
definition,
|
||||
credential_source,
|
||||
client,
|
||||
session_state: Arc::new(ResponsesSessionState::default()),
|
||||
endpoint_capabilities: Arc::new(ResponsesEndpointCapabilities::default()),
|
||||
hybrid_http_previous_response_id: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn definition(&self) -> &ProviderDefinition {
|
||||
&self.definition
|
||||
}
|
||||
|
||||
/// Disables opportunistic `previous_response_id` chaining for Hybrid HTTP
|
||||
/// requests made by this provider.
|
||||
///
|
||||
/// Use this when the endpoint is already known not to accept that optional
|
||||
/// Responses parameter. Hybrid requests retain their complete local replay,
|
||||
/// so disabling the optimization avoids a known-failing discovery request.
|
||||
/// Stateful requests, explicit response ids, and WebSocket transport keep
|
||||
/// their existing behavior.
|
||||
pub fn without_hybrid_http_previous_response_id(mut self) -> Self {
|
||||
self.hybrid_http_previous_response_id = false;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn session(&self) -> ResponsesSession<C> {
|
||||
ResponsesSession::new(
|
||||
self.definition.clone(),
|
||||
Arc::clone(&self.credential_source),
|
||||
self.client.clone(),
|
||||
Arc::clone(&self.session_state),
|
||||
Arc::clone(&self.endpoint_capabilities),
|
||||
self.hybrid_http_previous_response_id,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn openai_with_credential_source(credential_source: C) -> Self {
|
||||
Self::with_shared_credential_source(openai_definition(), Arc::new(credential_source))
|
||||
}
|
||||
|
||||
pub fn openrouter_with_credential_source(credential_source: C) -> Self {
|
||||
Self::with_shared_credential_source(openrouter_definition(), Arc::new(credential_source))
|
||||
}
|
||||
}
|
||||
|
||||
impl ResponsesProvider<StaticCredentialSource> {
|
||||
pub fn openai(api_key: impl Into<String>) -> Self {
|
||||
Self::openai_with_credential_source(StaticCredentialSource::new(api_key))
|
||||
}
|
||||
|
||||
pub fn openrouter(api_key: impl Into<String>) -> Self {
|
||||
Self::openrouter_with_credential_source(StaticCredentialSource::new(api_key))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn openai_definition() -> ProviderDefinition {
|
||||
build_definition(
|
||||
BuiltinProvider::OpenAI,
|
||||
"OpenAI",
|
||||
"OpenAI Responses API provider",
|
||||
DEFAULT_OPENAI_BASE_URL,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn openrouter_definition() -> ProviderDefinition {
|
||||
build_definition(
|
||||
BuiltinProvider::OpenRouter,
|
||||
"OpenRouter",
|
||||
"OpenRouter Responses API provider",
|
||||
DEFAULT_OPENROUTER_BASE_URL,
|
||||
)
|
||||
}
|
||||
|
||||
fn build_definition(
|
||||
builtin: BuiltinProvider,
|
||||
display_name: &str,
|
||||
description: &str,
|
||||
base_url: &str,
|
||||
) -> ProviderDefinition {
|
||||
let mut definition = ProviderDefinition::new(builtin);
|
||||
definition.descriptor.display_name = Some(display_name.to_string());
|
||||
definition.descriptor.description = Some(description.to_string());
|
||||
definition.wire_api = WireApi::Responses;
|
||||
definition.auth_scheme = AuthScheme::BearerToken;
|
||||
definition.capabilities = ProviderCapabilities {
|
||||
supports_model_listing: true,
|
||||
supports_streaming: true,
|
||||
supports_websockets: true,
|
||||
supports_tool_calls: true,
|
||||
supports_images: true,
|
||||
supports_history_compaction: true,
|
||||
supports_memory_summarization: true,
|
||||
supports_deferred_tools: true,
|
||||
supports_hosted_tool_search: true,
|
||||
supports_hosted_web_search: true,
|
||||
supports_image_generation: true,
|
||||
supports_reasoning_effort: true,
|
||||
reports_reasoning_tokens: true,
|
||||
reports_thoughts_tokens: false,
|
||||
supports_structured_tool_results: true,
|
||||
supports_embeddings: true,
|
||||
};
|
||||
definition.base_url = Some(base_url.to_string());
|
||||
definition.headers = Some(HashMap::new());
|
||||
definition.retry = RetryPolicy::default();
|
||||
definition
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ModelCatalog for ResponsesProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
let credentials = self.credential_source.credentials().await?;
|
||||
let request = self
|
||||
.client
|
||||
.get(
|
||||
self.definition
|
||||
.request_url_with_auth_for_path("v1/models", &credentials)?,
|
||||
)
|
||||
.headers(self.definition.build_headers(&credentials)?);
|
||||
|
||||
let response = request.send().await.map_err(ProviderError::Transport)?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::from_http_response(response).await);
|
||||
}
|
||||
|
||||
let models = response
|
||||
.json::<self::model::ResponsesModelsPage>()
|
||||
.await
|
||||
.map_err(ProviderError::Decode)?;
|
||||
|
||||
Ok(models.into_model_info(self.definition.descriptor.id.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ProviderSessionFactory for ResponsesProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn create_session(&self) -> Result<Box<dyn crate::ProviderSession>, ProviderError> {
|
||||
Ok(Box::new(self.session()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> RegisteredProvider for ResponsesProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
fn definition(&self) -> ProviderDefinition {
|
||||
self.definition.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> EmbeddingProvider for ResponsesProvider<C>
|
||||
where
|
||||
C: CredentialSource + 'static,
|
||||
{
|
||||
async fn embed_batch(
|
||||
&self,
|
||||
model: &str,
|
||||
texts: &[&str],
|
||||
) -> Result<EmbeddingResponse, ProviderError> {
|
||||
let credentials = self.credential_source.credentials().await?;
|
||||
let url = self
|
||||
.definition
|
||||
.request_url_with_auth_for_path("v1/embeddings", &credentials)?;
|
||||
let headers = self.definition.build_headers(&credentials)?;
|
||||
let body = EmbeddingRequest::batch(model, texts);
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(url)
|
||||
.headers(headers)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(ProviderError::Transport)?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(ProviderError::from_http_response(response).await);
|
||||
}
|
||||
|
||||
response
|
||||
.json::<EmbeddingResponse>()
|
||||
.await
|
||||
.map_err(ProviderError::Decode)
|
||||
}
|
||||
|
||||
fn embedding_models(&self) -> Vec<EmbeddingModelInfo> {
|
||||
// Available embedding models depend on the specific provider instance and its
|
||||
// configuration. Callers should use the /v1/models endpoint for discovery
|
||||
// rather than relying on a static list that would only be accurate for OpenAI.
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::ProviderId;
|
||||
|
||||
#[test]
|
||||
fn openai_preset_uses_responses_wire_api() {
|
||||
let provider = openai("test-key");
|
||||
let definition = provider.definition();
|
||||
|
||||
assert_eq!(
|
||||
definition.descriptor.id,
|
||||
ProviderId::from(BuiltinProvider::OpenAI)
|
||||
);
|
||||
assert_eq!(
|
||||
definition.descriptor.display_name.as_deref(),
|
||||
Some("OpenAI")
|
||||
);
|
||||
assert_eq!(definition.wire_api, WireApi::Responses);
|
||||
assert!(definition.capabilities.supports_websockets);
|
||||
assert!(definition.capabilities.supports_history_compaction);
|
||||
assert_eq!(
|
||||
definition.base_url.as_deref(),
|
||||
Some(DEFAULT_OPENAI_BASE_URL)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openrouter_preset_uses_openrouter_base_url() {
|
||||
let provider = openrouter("test-key");
|
||||
let definition = provider.definition();
|
||||
|
||||
assert_eq!(
|
||||
definition.descriptor.id,
|
||||
ProviderId::from(BuiltinProvider::OpenRouter)
|
||||
);
|
||||
assert_eq!(
|
||||
definition.descriptor.display_name.as_deref(),
|
||||
Some("OpenRouter")
|
||||
);
|
||||
assert_eq!(definition.wire_api, WireApi::Responses);
|
||||
assert!(definition.capabilities.supports_history_compaction);
|
||||
assert_eq!(
|
||||
definition.base_url.as_deref(),
|
||||
Some(DEFAULT_OPENROUTER_BASE_URL)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabling_hybrid_http_state_on_a_clone_does_not_reconfigure_the_original() {
|
||||
let original = openai("test-key");
|
||||
let disabled = original.clone().without_hybrid_http_previous_response_id();
|
||||
|
||||
assert!(original.hybrid_http_previous_response_id);
|
||||
assert!(!disabled.hybrid_http_previous_response_id);
|
||||
}
|
||||
}
|
||||
1516
vendor/mentra-provider/src/responses/model.rs
vendored
Normal file
1516
vendor/mentra-provider/src/responses/model.rs
vendored
Normal file
File diff suppressed because it is too large
Load Diff
2041
vendor/mentra-provider/src/responses/session.rs
vendored
Normal file
2041
vendor/mentra-provider/src/responses/session.rs
vendored
Normal file
File diff suppressed because it is too large
Load Diff
1542
vendor/mentra-provider/src/responses/sse.rs
vendored
Normal file
1542
vendor/mentra-provider/src/responses/sse.rs
vendored
Normal file
File diff suppressed because it is too large
Load Diff
705
vendor/mentra-provider/src/responses/websocket.rs
vendored
Normal file
705
vendor/mentra-provider/src/responses/websocket.rs
vendored
Normal file
@@ -0,0 +1,705 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::SinkExt;
|
||||
use futures_util::StreamExt;
|
||||
use http::HeaderMap;
|
||||
use http::HeaderValue;
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use serde_json::map::Map as JsonMap;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::time::Instant;
|
||||
use tokio_tungstenite::MaybeTlsStream;
|
||||
use tokio_tungstenite::WebSocketStream;
|
||||
use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::tungstenite::Error as WsError;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use url::Url;
|
||||
|
||||
use crate::ProviderError;
|
||||
use crate::ProviderEvent;
|
||||
use crate::ProviderEventStream;
|
||||
use crate::ProviderId;
|
||||
use crate::ReasoningFormat;
|
||||
use crate::ReasoningProvenance;
|
||||
use crate::ResponseHeaders;
|
||||
use crate::error::retry_after_from_header_value;
|
||||
use crate::error::retry_after_from_headers;
|
||||
|
||||
use super::SharedTurnState;
|
||||
use super::sse::StreamState;
|
||||
use super::sse::parse_json_event;
|
||||
|
||||
const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state";
|
||||
const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE: &str = "websocket_connection_limit_reached";
|
||||
const WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE: &str = "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue.";
|
||||
|
||||
pub trait ResponsesWebsocketTelemetry: Send + Sync {
|
||||
fn on_ws_request(
|
||||
&self,
|
||||
duration: Duration,
|
||||
error: Option<&ProviderError>,
|
||||
connection_reused: bool,
|
||||
);
|
||||
|
||||
fn on_ws_event(
|
||||
&self,
|
||||
result: &Result<Option<Result<Message, WsError>>, ProviderError>,
|
||||
duration: Duration,
|
||||
);
|
||||
}
|
||||
|
||||
struct WsStream {
|
||||
tx_command: mpsc::Sender<WsCommand>,
|
||||
rx_message: mpsc::UnboundedReceiver<Result<Message, WsError>>,
|
||||
pump_task: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
enum WsCommand {
|
||||
Send {
|
||||
message: Message,
|
||||
tx_result: oneshot::Sender<Result<(), WsError>>,
|
||||
},
|
||||
}
|
||||
|
||||
impl WsStream {
|
||||
fn new(inner: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
|
||||
let (tx_command, mut rx_command) = mpsc::channel::<WsCommand>(32);
|
||||
let (tx_message, rx_message) = mpsc::unbounded_channel::<Result<Message, WsError>>();
|
||||
|
||||
let pump_task = tokio::spawn(async move {
|
||||
let mut inner = inner;
|
||||
loop {
|
||||
tokio::select! {
|
||||
command = rx_command.recv() => {
|
||||
let Some(command) = command else {
|
||||
break;
|
||||
};
|
||||
match command {
|
||||
WsCommand::Send { message, tx_result } => {
|
||||
let result = inner.send(message).await;
|
||||
let should_break = result.is_err();
|
||||
let _ = tx_result.send(result);
|
||||
if should_break {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
message = inner.next() => {
|
||||
let Some(message) = message else {
|
||||
break;
|
||||
};
|
||||
match message {
|
||||
Ok(Message::Ping(payload)) => {
|
||||
if let Err(err) = inner.send(Message::Pong(payload)).await {
|
||||
let _ = tx_message.send(Err(err));
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Message::Pong(_)) => {}
|
||||
Ok(message @ (Message::Text(_)
|
||||
| Message::Binary(_)
|
||||
| Message::Close(_)
|
||||
| Message::Frame(_))) => {
|
||||
let is_close = matches!(message, Message::Close(_));
|
||||
if tx_message.send(Ok(message)).is_err() {
|
||||
break;
|
||||
}
|
||||
if is_close {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = tx_message.send(Err(err));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Self {
|
||||
tx_command,
|
||||
rx_message,
|
||||
pump_task,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(&self, message: Message) -> Result<(), WsError> {
|
||||
let (tx_result, rx_result) = oneshot::channel();
|
||||
if self
|
||||
.tx_command
|
||||
.send(WsCommand::Send { message, tx_result })
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return Err(WsError::ConnectionClosed);
|
||||
}
|
||||
rx_result.await.unwrap_or(Err(WsError::ConnectionClosed))
|
||||
}
|
||||
|
||||
async fn next(&mut self) -> Option<Result<Message, WsError>> {
|
||||
self.rx_message.recv().await
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WsStream {
|
||||
fn drop(&mut self) {
|
||||
self.pump_task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponsesWebsocketConnection {
|
||||
stream: Arc<Mutex<Option<WsStream>>>,
|
||||
idle_timeout: Duration,
|
||||
response_headers: ResponseHeaders,
|
||||
telemetry: Option<Arc<dyn ResponsesWebsocketTelemetry>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ResponsesWebsocketConnection {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ResponsesWebsocketConnection")
|
||||
.field("stream", &"<ws-stream>")
|
||||
.field("idle_timeout", &self.idle_timeout)
|
||||
.field("response_headers", &self.response_headers)
|
||||
.field("telemetry", &self.telemetry.as_ref().map(|_| "<telemetry>"))
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl ResponsesWebsocketConnection {
|
||||
pub async fn connect(
|
||||
url: Url,
|
||||
headers: HeaderMap,
|
||||
turn_state: Option<SharedTurnState>,
|
||||
idle_timeout: Duration,
|
||||
telemetry: Option<Arc<dyn ResponsesWebsocketTelemetry>>,
|
||||
) -> Result<Self, ProviderError> {
|
||||
let mut request = url
|
||||
.as_str()
|
||||
.into_client_request()
|
||||
.map_err(|error| ProviderError::InvalidRequest(error.to_string()))?;
|
||||
request.headers_mut().extend(headers);
|
||||
|
||||
let (stream, response) = connect_async(request)
|
||||
.await
|
||||
.map_err(|error| map_ws_error(error, &url))?;
|
||||
|
||||
let header_value = response
|
||||
.headers()
|
||||
.get(X_CODEX_TURN_STATE_HEADER)
|
||||
.and_then(|value| value.to_str().ok());
|
||||
if let (Some(turn_state), Some(header_value)) = (turn_state, header_value) {
|
||||
*turn_state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) =
|
||||
Some(header_value.to_string());
|
||||
}
|
||||
|
||||
let response_headers = ResponseHeaders {
|
||||
values: response
|
||||
.headers()
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
value
|
||||
.to_str()
|
||||
.ok()
|
||||
.map(|value| (name.as_str().to_string(), value.to_string()))
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
stream: Arc::new(Mutex::new(Some(WsStream::new(stream)))),
|
||||
idle_timeout,
|
||||
response_headers,
|
||||
telemetry,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn is_closed(&self) -> bool {
|
||||
self.stream.lock().await.is_none()
|
||||
}
|
||||
|
||||
pub async fn stream_request(
|
||||
&self,
|
||||
request_body: Value,
|
||||
connection_reused: bool,
|
||||
) -> Result<ProviderEventStream, ProviderError> {
|
||||
let requested_model = request_body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
self.stream_request_with_provenance(
|
||||
request_body,
|
||||
connection_reused,
|
||||
ReasoningProvenance {
|
||||
provider: ProviderId::new("openai"),
|
||||
model: requested_model,
|
||||
format: ReasoningFormat::OpenAiEncrypted,
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn stream_request_with_provenance(
|
||||
&self,
|
||||
request_body: Value,
|
||||
connection_reused: bool,
|
||||
provenance: ReasoningProvenance,
|
||||
) -> Result<ProviderEventStream, ProviderError> {
|
||||
let (tx_event, rx_event) =
|
||||
mpsc::unbounded_channel::<Result<ProviderEvent, ProviderError>>();
|
||||
let stream = Arc::clone(&self.stream);
|
||||
let idle_timeout = self.idle_timeout;
|
||||
let response_headers = self.response_headers.clone();
|
||||
let telemetry = self.telemetry.clone();
|
||||
|
||||
let request_text =
|
||||
serde_json::to_string(&request_body).map_err(ProviderError::Serialize)?;
|
||||
|
||||
tokio::spawn(async move {
|
||||
if tx_event
|
||||
.send(Ok(ProviderEvent::ResponseHeaders(response_headers)))
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let mut guard = stream.lock().await;
|
||||
let result = {
|
||||
let Some(ws_stream) = guard.as_mut() else {
|
||||
let _ = tx_event.send(Err(ProviderError::MalformedStream(
|
||||
"websocket connection is closed".to_string(),
|
||||
)));
|
||||
return;
|
||||
};
|
||||
run_websocket_response_stream(
|
||||
ws_stream,
|
||||
tx_event.clone(),
|
||||
request_text,
|
||||
idle_timeout,
|
||||
telemetry,
|
||||
connection_reused,
|
||||
StreamState::new(provenance.provider, provenance.model),
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
if let Err(err) = result {
|
||||
let failed_stream = guard.take();
|
||||
drop(guard);
|
||||
drop(failed_stream);
|
||||
let _ = tx_event.send(Err(err));
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx_event)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn merge_request_headers(
|
||||
provider_headers: &HeaderMap,
|
||||
extra_headers: HeaderMap,
|
||||
default_headers: HeaderMap,
|
||||
) -> HeaderMap {
|
||||
let mut headers = provider_headers.clone();
|
||||
headers.extend(extra_headers);
|
||||
for (name, value) in &default_headers {
|
||||
if let http::header::Entry::Vacant(entry) = headers.entry(name) {
|
||||
entry.insert(value.clone());
|
||||
}
|
||||
}
|
||||
headers
|
||||
}
|
||||
|
||||
pub(crate) fn response_create_frame(mut response: Value) -> Value {
|
||||
let Value::Object(response_object) = &mut response else {
|
||||
return json!({
|
||||
"type": "response.create",
|
||||
"response": response,
|
||||
});
|
||||
};
|
||||
|
||||
response_object.insert(
|
||||
"type".to_string(),
|
||||
Value::String("response.create".to_string()),
|
||||
);
|
||||
response_object
|
||||
.entry("instructions")
|
||||
.or_insert_with(|| Value::String(String::new()));
|
||||
response
|
||||
}
|
||||
|
||||
async fn run_websocket_response_stream(
|
||||
ws_stream: &mut WsStream,
|
||||
tx_event: mpsc::UnboundedSender<Result<ProviderEvent, ProviderError>>,
|
||||
request_text: String,
|
||||
idle_timeout: Duration,
|
||||
telemetry: Option<Arc<dyn ResponsesWebsocketTelemetry>>,
|
||||
connection_reused: bool,
|
||||
mut state: StreamState,
|
||||
) -> Result<(), ProviderError> {
|
||||
let request_start = Instant::now();
|
||||
let send_result = ws_stream.send(Message::Text(request_text.into())).await;
|
||||
let send_error = send_result
|
||||
.as_ref()
|
||||
.err()
|
||||
.map(|error| ProviderError::MalformedStream(error.to_string()));
|
||||
if let Some(t) = telemetry.as_ref() {
|
||||
t.on_ws_request(
|
||||
request_start.elapsed(),
|
||||
send_error.as_ref(),
|
||||
connection_reused,
|
||||
);
|
||||
}
|
||||
send_result.map_err(|error| ProviderError::MalformedStream(error.to_string()))?;
|
||||
|
||||
loop {
|
||||
let poll_start = Instant::now();
|
||||
let message_result = tokio::time::timeout(idle_timeout, ws_stream.next())
|
||||
.await
|
||||
.map_err(|_| {
|
||||
ProviderError::MalformedStream("idle timeout waiting for websocket".into())
|
||||
});
|
||||
if let Some(t) = telemetry.as_ref() {
|
||||
t.on_ws_event(&message_result, poll_start.elapsed());
|
||||
}
|
||||
let message = match message_result {
|
||||
Ok(Some(Ok(message))) => message,
|
||||
Ok(Some(Err(error))) => {
|
||||
return Err(ProviderError::MalformedStream(error.to_string()));
|
||||
}
|
||||
Ok(None) => {
|
||||
return Err(ProviderError::MalformedStream(
|
||||
"stream closed before response.completed".to_string(),
|
||||
));
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
|
||||
match message {
|
||||
Message::Text(text) => {
|
||||
if let Some(mapped) = parse_wrapped_websocket_error_event(&text) {
|
||||
let receiver_closed = mapped
|
||||
.headers
|
||||
.map(|headers| {
|
||||
tx_event
|
||||
.send(Ok(ProviderEvent::ResponseHeaders(headers)))
|
||||
.is_err()
|
||||
})
|
||||
.unwrap_or(false);
|
||||
if receiver_closed {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(mapped.error);
|
||||
}
|
||||
|
||||
let mut saw_message_stopped = false;
|
||||
for event in parse_json_event(&text, &mut state)? {
|
||||
if matches!(event, ProviderEvent::MessageStopped) {
|
||||
saw_message_stopped = true;
|
||||
}
|
||||
if tx_event.send(Ok(event)).is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
if saw_message_stopped {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Message::Binary(_) => {
|
||||
return Err(ProviderError::MalformedStream(
|
||||
"unexpected binary websocket event".to_string(),
|
||||
));
|
||||
}
|
||||
Message::Close(_) => {
|
||||
return Err(ProviderError::MalformedStream(
|
||||
"websocket closed by server before response.completed".to_string(),
|
||||
));
|
||||
}
|
||||
Message::Frame(_) => {}
|
||||
Message::Ping(_) | Message::Pong(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Delay hint for connection-level websocket failures. Short enough that a
|
||||
/// tunnel blip or service restart is retried promptly, long enough to avoid
|
||||
/// hammering a target that is still coming back up.
|
||||
const WS_CONNECTION_RETRY_DELAY: Duration = Duration::from_millis(750);
|
||||
|
||||
fn map_ws_error(error: WsError, url: &Url) -> ProviderError {
|
||||
match error {
|
||||
WsError::Http(response) => {
|
||||
let status = response.status();
|
||||
let retry_after = retry_after_from_headers(response.headers());
|
||||
let body = response
|
||||
.body()
|
||||
.as_ref()
|
||||
.and_then(|bytes| String::from_utf8(bytes.clone()).ok())
|
||||
.unwrap_or_else(|| format!("websocket connection failed for {url}"));
|
||||
ProviderError::Http {
|
||||
status,
|
||||
body,
|
||||
retry_after,
|
||||
}
|
||||
}
|
||||
// Connection lost mid-handshake: the peer went away before the upgrade
|
||||
// completed. That is a transient liveness failure (restart, tunnel
|
||||
// teardown), not a request the caller must fix — surface it as retryable.
|
||||
WsError::ConnectionClosed | WsError::AlreadyClosed => ProviderError::Retryable {
|
||||
message: format!("websocket connection closed while connecting to {url}"),
|
||||
delay: Some(WS_CONNECTION_RETRY_DELAY),
|
||||
},
|
||||
// Transport-level failure establishing the connection (e.g. connection
|
||||
// refused when an SSH tunnel is down, DNS blips). These clear on their
|
||||
// own once the path is back, so the whole turn is worth retrying rather
|
||||
// than falling straight through to a terminal failure.
|
||||
WsError::Io(error) => ProviderError::Retryable {
|
||||
message: format!("websocket connection failed for {url}: {error}"),
|
||||
delay: Some(WS_CONNECTION_RETRY_DELAY),
|
||||
},
|
||||
other => ProviderError::InvalidResponse(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct MappedWebsocketError {
|
||||
error: ProviderError,
|
||||
headers: Option<ResponseHeaders>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WrappedWebsocketError {
|
||||
code: Option<String>,
|
||||
message: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WrappedWebsocketErrorEvent {
|
||||
#[serde(rename = "type")]
|
||||
kind: String,
|
||||
#[serde(alias = "status_code")]
|
||||
status: Option<u16>,
|
||||
#[serde(default)]
|
||||
error: Option<WrappedWebsocketError>,
|
||||
#[serde(default)]
|
||||
headers: Option<JsonMap<String, Value>>,
|
||||
}
|
||||
|
||||
fn parse_wrapped_websocket_error_event(payload: &str) -> Option<MappedWebsocketError> {
|
||||
let event: WrappedWebsocketErrorEvent = serde_json::from_str(payload).ok()?;
|
||||
if event.kind != "error" {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(error) = event
|
||||
.error
|
||||
.as_ref()
|
||||
.filter(|error| error.code.as_deref() == Some(WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE))
|
||||
{
|
||||
return Some(MappedWebsocketError {
|
||||
error: ProviderError::Retryable {
|
||||
message: error
|
||||
.message
|
||||
.clone()
|
||||
.unwrap_or_else(|| WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE.to_string()),
|
||||
delay: None,
|
||||
},
|
||||
headers: event.headers.map(response_headers_from_json),
|
||||
});
|
||||
}
|
||||
|
||||
let status = reqwest::StatusCode::from_u16(event.status?).ok()?;
|
||||
let body = payload.to_string();
|
||||
let headers = event.headers.map(response_headers_from_json);
|
||||
// The upgrade's own headers are long gone by the time a rate limit arrives
|
||||
// as a wrapped frame, so the only `Retry-After` on this path is the one the
|
||||
// frame echoed. Read it here or it is lost.
|
||||
let retry_after = headers.as_ref().and_then(retry_after_from_response_headers);
|
||||
Some(MappedWebsocketError {
|
||||
error: ProviderError::Http {
|
||||
status,
|
||||
body,
|
||||
retry_after,
|
||||
},
|
||||
headers,
|
||||
})
|
||||
}
|
||||
|
||||
/// Reads `Retry-After` out of headers a frame carried as JSON.
|
||||
fn retry_after_from_response_headers(headers: &ResponseHeaders) -> Option<Duration> {
|
||||
headers
|
||||
.values
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("retry-after"))
|
||||
.and_then(|(_, value)| retry_after_from_header_value(value))
|
||||
}
|
||||
|
||||
fn response_headers_from_json(headers: JsonMap<String, Value>) -> ResponseHeaders {
|
||||
ResponseHeaders {
|
||||
values: headers
|
||||
.into_iter()
|
||||
.filter_map(|(name, value)| {
|
||||
json_header_value(value)
|
||||
.and_then(|value| value.to_str().ok().map(|value| (name, value.to_string())))
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn json_header_value(value: Value) -> Option<HeaderValue> {
|
||||
let value = match value {
|
||||
Value::String(value) => value,
|
||||
Value::Number(value) => value.to_string(),
|
||||
Value::Bool(value) => value.to_string(),
|
||||
_ => return None,
|
||||
};
|
||||
HeaderValue::from_str(&value).ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn merge_request_headers_matches_http_precedence() {
|
||||
let mut provider_headers = HeaderMap::new();
|
||||
provider_headers.insert(
|
||||
"originator",
|
||||
HeaderValue::from_static("provider-originator"),
|
||||
);
|
||||
provider_headers.insert("x-priority", HeaderValue::from_static("provider"));
|
||||
|
||||
let mut extra_headers = HeaderMap::new();
|
||||
extra_headers.insert("x-priority", HeaderValue::from_static("extra"));
|
||||
|
||||
let mut default_headers = HeaderMap::new();
|
||||
default_headers.insert("originator", HeaderValue::from_static("default-originator"));
|
||||
default_headers.insert("x-priority", HeaderValue::from_static("default"));
|
||||
default_headers.insert("x-default-only", HeaderValue::from_static("default-only"));
|
||||
|
||||
let merged = merge_request_headers(&provider_headers, extra_headers, default_headers);
|
||||
|
||||
assert_eq!(
|
||||
merged.get("originator"),
|
||||
Some(&HeaderValue::from_static("provider-originator"))
|
||||
);
|
||||
assert_eq!(
|
||||
merged.get("x-priority"),
|
||||
Some(&HeaderValue::from_static("extra"))
|
||||
);
|
||||
assert_eq!(
|
||||
merged.get("x-default-only"),
|
||||
Some(&HeaderValue::from_static("default-only"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_create_frame_adds_type_to_responses_request_payload() {
|
||||
let frame = response_create_frame(json!({
|
||||
"model": "gpt-5",
|
||||
"input": []
|
||||
}));
|
||||
|
||||
assert_eq!(frame["type"], "response.create");
|
||||
assert_eq!(frame["model"], "gpt-5");
|
||||
assert_eq!(frame["input"], json!([]));
|
||||
assert_eq!(frame["instructions"], "");
|
||||
assert!(frame.get("response").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrapped_websocket_error_preserves_rate_limit_headers() {
|
||||
let payload = json!({
|
||||
"type": "error",
|
||||
"status": 429,
|
||||
"error": {
|
||||
"type": "usage_limit_reached",
|
||||
"message": "The usage limit has been reached"
|
||||
},
|
||||
"headers": {
|
||||
"x-codex-primary-used-percent": "100.0"
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let mapped = parse_wrapped_websocket_error_event(&payload)
|
||||
.expect("error payload should map to a provider error");
|
||||
let ProviderError::Http { status, .. } = mapped.error else {
|
||||
panic!("expected ProviderError::Http");
|
||||
};
|
||||
assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
mapped
|
||||
.headers
|
||||
.expect("rate limit headers should be present")
|
||||
.values,
|
||||
vec![(
|
||||
"x-codex-primary-used-percent".to_string(),
|
||||
"100.0".to_string()
|
||||
)]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ws_io_error_maps_to_retryable_with_delay() {
|
||||
let url = Url::parse("wss://example.test/responses").expect("url parses");
|
||||
let io = std::io::Error::new(std::io::ErrorKind::ConnectionRefused, "Connection refused");
|
||||
let mapped = map_ws_error(WsError::Io(io), &url);
|
||||
let ProviderError::Retryable { message, delay } = mapped else {
|
||||
panic!("expected ProviderError::Retryable for a connection-level io error");
|
||||
};
|
||||
assert_eq!(delay, Some(WS_CONNECTION_RETRY_DELAY));
|
||||
// The io error string is preserved so operators can see the real cause.
|
||||
assert!(message.contains("Connection refused"), "message: {message}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ws_connection_closed_maps_to_retryable() {
|
||||
let url = Url::parse("wss://example.test/responses").expect("url parses");
|
||||
for error in [WsError::ConnectionClosed, WsError::AlreadyClosed] {
|
||||
let mapped = map_ws_error(error, &url);
|
||||
assert!(
|
||||
matches!(mapped, ProviderError::Retryable { delay: Some(_), .. }),
|
||||
"connection-closed handshake failure should be retryable",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_connection_limit_maps_to_retryable() {
|
||||
let payload = json!({
|
||||
"type": "error",
|
||||
"status": 400,
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": "websocket_connection_limit_reached",
|
||||
"message": WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let mapped = parse_wrapped_websocket_error_event(&payload)
|
||||
.expect("connection limit payload should map");
|
||||
let ProviderError::Retryable { message, delay } = mapped.error else {
|
||||
panic!("expected ProviderError::Retryable");
|
||||
};
|
||||
assert_eq!(message, WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE);
|
||||
assert_eq!(delay, None);
|
||||
}
|
||||
}
|
||||
101
vendor/mentra-provider/src/stream.rs
vendored
Normal file
101
vendor/mentra-provider/src/stream.rs
vendored
Normal file
@@ -0,0 +1,101 @@
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::{
|
||||
ReasoningProvenance, model::HostedToolSearchCall, model::HostedWebSearchCall,
|
||||
model::ImageGenerationCall, model::ImageGenerationResult, model::ImageSource, model::Role,
|
||||
model::TokenUsage, model::ToolResultContent, model::WebSearchAction,
|
||||
};
|
||||
|
||||
pub type ProviderEventStream = mpsc::UnboundedReceiver<Result<ProviderEvent, crate::ProviderError>>;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct ResponseHeaders {
|
||||
pub values: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ProviderEvent {
|
||||
ResponseHeaders(ResponseHeaders),
|
||||
ResponseCreated,
|
||||
MessageStarted {
|
||||
id: String,
|
||||
model: String,
|
||||
role: Role,
|
||||
},
|
||||
ContentBlockStarted {
|
||||
index: usize,
|
||||
kind: ContentBlockStart,
|
||||
},
|
||||
ContentBlockDelta {
|
||||
index: usize,
|
||||
delta: ContentBlockDelta,
|
||||
},
|
||||
ContentBlockStopped {
|
||||
index: usize,
|
||||
},
|
||||
MessageDelta {
|
||||
stop_reason: Option<String>,
|
||||
usage: Option<TokenUsage>,
|
||||
},
|
||||
ReasoningSummaryDelta {
|
||||
delta: String,
|
||||
summary_index: i64,
|
||||
},
|
||||
ReasoningContentDelta {
|
||||
delta: String,
|
||||
content_index: i64,
|
||||
},
|
||||
ReasoningSummaryPartAdded {
|
||||
summary_index: i64,
|
||||
},
|
||||
MessageStopped,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ContentBlockStart {
|
||||
Text,
|
||||
Thinking {
|
||||
encrypted_content: Option<String>,
|
||||
id: Option<String>,
|
||||
provenance: Option<ReasoningProvenance>,
|
||||
redacted: bool,
|
||||
},
|
||||
Image {
|
||||
source: ImageSource,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
},
|
||||
ToolResult {
|
||||
tool_use_id: String,
|
||||
is_error: bool,
|
||||
content: Option<ToolResultContent>,
|
||||
},
|
||||
HostedToolSearch {
|
||||
call: HostedToolSearchCall,
|
||||
},
|
||||
HostedWebSearch {
|
||||
call: HostedWebSearchCall,
|
||||
},
|
||||
ImageGeneration {
|
||||
call: ImageGenerationCall,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ContentBlockDelta {
|
||||
Text(String),
|
||||
ThinkingText(String),
|
||||
ThinkingSignature(String),
|
||||
ThinkingEncryptedContent(String),
|
||||
ToolUseInputJson(String),
|
||||
ToolResultContent(ToolResultContent),
|
||||
HostedToolSearchQuery(String),
|
||||
HostedToolSearchStatus(String),
|
||||
HostedWebSearchAction(WebSearchAction),
|
||||
HostedWebSearchStatus(String),
|
||||
ImageGenerationStatus(String),
|
||||
ImageGenerationRevisedPrompt(String),
|
||||
ImageGenerationResult(ImageGenerationResult),
|
||||
}
|
||||
151
vendor/mentra-provider/src/tool.rs
vendored
Normal file
151
vendor/mentra-provider/src/tool.rs
vendored
Normal file
@@ -0,0 +1,151 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
/// Declares whether a tool is loaded eagerly or deferred for provider-native tool search.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub enum ToolLoadingPolicy {
|
||||
#[default]
|
||||
Immediate,
|
||||
Deferred,
|
||||
}
|
||||
|
||||
/// Provider-visible tool kinds supported by Mentra-backed providers.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ProviderToolKind {
|
||||
#[default]
|
||||
Function,
|
||||
HostedWebSearch,
|
||||
ImageGeneration,
|
||||
}
|
||||
|
||||
/// Provider-facing description of a tool and its schemas.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ToolSpec {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub input_schema: Value,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub output_schema: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub kind: ProviderToolKind,
|
||||
#[serde(default)]
|
||||
pub loading_policy: ToolLoadingPolicy,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub strict: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub options: Option<Value>,
|
||||
}
|
||||
|
||||
impl ToolSpec {
|
||||
pub fn builder(name: impl Into<String>) -> ToolSpecBuilder {
|
||||
ToolSpecBuilder {
|
||||
name: name.into(),
|
||||
description: None,
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}),
|
||||
output_schema: None,
|
||||
kind: ProviderToolKind::Function,
|
||||
loading_policy: ToolLoadingPolicy::Immediate,
|
||||
strict: None,
|
||||
options: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolSpecBuilder {
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
input_schema: Value,
|
||||
output_schema: Option<Value>,
|
||||
kind: ProviderToolKind,
|
||||
loading_policy: ToolLoadingPolicy,
|
||||
strict: Option<bool>,
|
||||
options: Option<Value>,
|
||||
}
|
||||
|
||||
impl ToolSpecBuilder {
|
||||
pub fn description(mut self, description: impl Into<String>) -> Self {
|
||||
self.description = Some(description.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn input_schema(mut self, input_schema: Value) -> Self {
|
||||
self.input_schema = input_schema;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn output_schema(mut self, output_schema: Value) -> Self {
|
||||
self.output_schema = Some(output_schema);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn kind(mut self, kind: ProviderToolKind) -> Self {
|
||||
self.kind = kind;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn options(mut self, options: Value) -> Self {
|
||||
self.options = Some(options);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn loading_policy(mut self, loading_policy: ToolLoadingPolicy) -> Self {
|
||||
self.loading_policy = loading_policy;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn strict(mut self, strict: bool) -> Self {
|
||||
self.strict = Some(strict);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn non_strict(self) -> Self {
|
||||
self.strict(false)
|
||||
}
|
||||
|
||||
pub fn defer_loading(self, defer_loading: bool) -> Self {
|
||||
self.loading_policy(if defer_loading {
|
||||
ToolLoadingPolicy::Deferred
|
||||
} else {
|
||||
ToolLoadingPolicy::Immediate
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build(self) -> ToolSpec {
|
||||
ToolSpec {
|
||||
name: self.name,
|
||||
description: self.description,
|
||||
input_schema: self.input_schema,
|
||||
output_schema: self.output_schema,
|
||||
kind: self.kind,
|
||||
loading_policy: self.loading_policy,
|
||||
strict: self.strict,
|
||||
options: self.options,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn builder_leaves_function_tool_strictness_unset_by_default() {
|
||||
let spec = ToolSpec::builder("echo").build();
|
||||
|
||||
assert_eq!(spec.strict, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builder_sets_explicit_strictness() {
|
||||
let strict = ToolSpec::builder("strict").strict(true).build();
|
||||
let non_strict = ToolSpec::builder("loose").non_strict().build();
|
||||
|
||||
assert_eq!(strict.strict, Some(true));
|
||||
assert_eq!(non_strict.strict, Some(false));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user