use async_trait::async_trait;
use futures::Stream;
use oxi_ai::{
Api, AssistantMessage, ContentBlock, Context, Model, Provider, ProviderError, ProviderEvent,
StopReason, StreamOptions, TextContent, Usage,
};
use oxi_sdk::OxiBuilder;
use std::pin::Pin;
pub struct MockProvider;
#[async_trait]
impl Provider for MockProvider {
fn name(&self) -> &str {
"mock"
}
async fn stream(
&self,
_model: &Model,
context: &Context,
_options: Option<StreamOptions>,
) -> Result<Pin<Box<dyn Stream<Item = ProviderEvent> + Send>>, ProviderError> {
let last_msg = context
.messages
.iter()
.rev()
.find_map(|m| match m {
oxi_ai::Message::User(u) => u.content.as_str().map(|s| s.to_string()),
_ => None,
})
.unwrap_or_else(|| "mock response".to_string());
let mut msg = AssistantMessage::new(Api::OpenAiCompletions, "mock", "mock/model");
msg.content
.push(ContentBlock::Text(TextContent::new(&last_msg)));
msg.stop_reason = StopReason::Stop;
msg.usage = Usage {
input: 100,
output: 50,
cache_read: 0,
cache_write: 0,
total_tokens: 150,
cost: Default::default(),
};
let partial = std::sync::Arc::new(msg.clone());
let events: Vec<ProviderEvent> = vec![
ProviderEvent::Start {
partial: partial.clone(),
},
ProviderEvent::TextDelta {
content_index: 0,
delta: last_msg,
partial,
},
ProviderEvent::Done {
reason: StopReason::Stop,
message: msg,
},
];
Ok(Box::pin(futures::stream::iter(events)))
}
}
pub fn mock_model() -> Model {
Model::new(
"model",
"Mock",
Api::OpenAiCompletions,
"mock",
"http://localhost",
)
}
pub fn mock_oxi() -> oxi_sdk::Oxi {
OxiBuilder::new()
.provider("mock", MockProvider)
.model(mock_model())
.build()
}