use crate::client::{
self, BearerAuth, Capabilities, Capable, DebugExt, Nothing, Provider, ProviderBuilder,
ProviderClient,
};
use crate::http_client::{self, HttpClientExt};
use crate::message::{Document, DocumentSourceKind};
use crate::providers::openai;
use crate::providers::openai::send_compatible_streaming_request;
use crate::streaming::StreamingCompletionResponse;
use crate::{
OneOrMany,
completion::{self, CompletionError, CompletionRequest},
message::{self, AssistantContent, Message, UserContent},
};
use serde::{Deserialize, Serialize};
use std::string::FromUtf8Error;
use thiserror::Error;
use tracing::{self, Instrument, info_span};
#[derive(Debug, Default, Clone, Copy)]
pub struct MiraExt;
#[derive(Debug, Default, Clone, Copy)]
pub struct MiraBuilder;
type MiraApiKey = BearerAuth;
impl Provider for MiraExt {
type Builder = MiraBuilder;
const VERIFY_PATH: &'static str = "/user-credits";
}
impl<H> Capabilities<H> for MiraExt {
type Completion = Capable<CompletionModel<H>>;
type Embeddings = Nothing;
type Transcription = Nothing;
type ModelListing = Nothing;
#[cfg(feature = "image")]
type ImageGeneration = Nothing;
#[cfg(feature = "audio")]
type AudioGeneration = Nothing;
}
impl DebugExt for MiraExt {}
impl ProviderBuilder for MiraBuilder {
type Extension<H>
= MiraExt
where
H: HttpClientExt;
type ApiKey = MiraApiKey;
const BASE_URL: &'static str = MIRA_API_BASE_URL;
fn build<H>(
_builder: &crate::client::ClientBuilder<Self, Self::ApiKey, H>,
) -> http_client::Result<Self::Extension<H>>
where
H: HttpClientExt,
{
Ok(MiraExt)
}
}
pub type Client<H = reqwest::Client> = client::Client<MiraExt, H>;
pub type ClientBuilder<H = reqwest::Client> = client::ClientBuilder<MiraBuilder, MiraApiKey, H>;
#[derive(Debug, Error)]
pub enum MiraError {
#[error("Invalid API key")]
InvalidApiKey,
#[error("API error: {0}")]
ApiError(u16),
#[error("Request error: {0}")]
RequestError(#[from] http_client::Error),
#[error("UTF-8 error: {0}")]
Utf8Error(#[from] FromUtf8Error),
#[error("JSON error: {0}")]
JsonError(#[from] serde_json::Error),
}
#[derive(Debug, Deserialize)]
struct ApiErrorResponse {
message: String,
}
#[derive(Debug, Deserialize, Clone, Serialize)]
pub struct RawMessage {
pub role: String,
pub content: String,
}
const MIRA_API_BASE_URL: &str = "https://api.mira.network";
impl TryFrom<RawMessage> for message::Message {
type Error = CompletionError;
fn try_from(raw: RawMessage) -> Result<Self, Self::Error> {
match raw.role.as_str() {
"system" => Ok(message::Message::System {
content: raw.content,
}),
"user" => Ok(message::Message::User {
content: OneOrMany::one(UserContent::Text(message::Text { text: raw.content })),
}),
"assistant" => Ok(message::Message::Assistant {
id: None,
content: OneOrMany::one(AssistantContent::Text(message::Text {
text: raw.content,
})),
}),
_ => Err(CompletionError::ResponseError(format!(
"Unsupported message role: {}",
raw.role
))),
}
}
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(untagged)]
pub enum CompletionResponse {
Structured {
id: String,
object: String,
created: u64,
model: String,
choices: Vec<ChatChoice>,
#[serde(skip_serializing_if = "Option::is_none")]
usage: Option<Usage>,
},
Simple(String),
}
#[derive(Debug, Deserialize, Serialize)]
pub struct ChatChoice {
pub message: RawMessage,
#[serde(default)]
pub finish_reason: Option<String>,
#[serde(default)]
pub index: Option<usize>,
}
#[derive(Debug, Deserialize, Serialize)]
struct ModelsResponse {
data: Vec<ModelInfo>,
}
#[derive(Debug, Deserialize, Serialize)]
struct ModelInfo {
id: String,
}
impl<T> Client<T>
where
T: HttpClientExt + 'static,
{
pub async fn list_models(&self) -> Result<Vec<String>, MiraError> {
let req = self.get("/v1/models").and_then(|req| {
req.body(http_client::NoBody)
.map_err(http_client::Error::Protocol)
})?;
let response = self.send(req).await?;
let status = response.status();
if !status.is_success() {
let error_text = http_client::text(response).await.unwrap_or_default();
tracing::error!("Error response: {}", error_text);
return Err(MiraError::ApiError(status.as_u16()));
}
let response_text = http_client::text(response).await?;
let models: ModelsResponse = serde_json::from_str(&response_text).map_err(|e| {
tracing::error!("Failed to parse response: {}", e);
MiraError::JsonError(e)
})?;
Ok(models.data.into_iter().map(|model| model.id).collect())
}
}
impl ProviderClient for Client {
type Input = String;
fn from_env() -> Self {
let api_key = std::env::var("MIRA_API_KEY").expect("MIRA_API_KEY not set");
Self::new(&api_key).unwrap()
}
fn from_val(input: Self::Input) -> Self {
Self::new(&input).unwrap()
}
}
#[derive(Debug, Serialize, Deserialize)]
pub(super) struct MiraCompletionRequest {
model: String,
pub messages: Vec<RawMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u64>,
pub stream: bool,
}
impl TryFrom<(&str, CompletionRequest)> for MiraCompletionRequest {
type Error = CompletionError;
fn try_from((model, req): (&str, CompletionRequest)) -> Result<Self, Self::Error> {
if req.output_schema.is_some() {
tracing::warn!("Structured outputs currently not supported for Mira");
}
let model = req.model.clone().unwrap_or_else(|| model.to_string());
let mut messages = Vec::new();
if let Some(content) = &req.preamble {
messages.push(RawMessage {
role: "user".to_string(),
content: content.to_string(),
});
}
if let Some(Message::User { content }) = req.normalized_documents() {
let text = content
.into_iter()
.filter_map(|doc| match doc {
UserContent::Document(Document {
data: DocumentSourceKind::Base64(data) | DocumentSourceKind::String(data),
..
}) => Some(data),
UserContent::Text(text) => Some(text.text),
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
messages.push(RawMessage {
role: "user".to_string(),
content: text,
});
}
for msg in req.chat_history {
let (role, content) = match msg {
Message::System { content } => ("system", content),
Message::User { content } => {
let text = content
.iter()
.map(|c| match c {
UserContent::Text(text) => &text.text,
_ => "",
})
.collect::<Vec<_>>()
.join("\n");
("user", text)
}
Message::Assistant { content, .. } => {
let text = content
.iter()
.map(|c| match c {
AssistantContent::Text(text) => &text.text,
_ => "",
})
.collect::<Vec<_>>()
.join("\n");
("assistant", text)
}
};
messages.push(RawMessage {
role: role.to_string(),
content,
});
}
Ok(Self {
model: model.to_string(),
messages,
temperature: req.temperature,
max_tokens: req.max_tokens,
stream: false,
})
}
}
#[derive(Clone)]
pub struct CompletionModel<T = reqwest::Client> {
client: Client<T>,
pub model: String,
}
impl<T> CompletionModel<T> {
pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
Self {
client,
model: model.into(),
}
}
}
impl<T> completion::CompletionModel for CompletionModel<T>
where
T: HttpClientExt + Clone + Default + std::fmt::Debug + Send + 'static,
{
type Response = CompletionResponse;
type StreamingResponse = openai::StreamingCompletionResponse;
type Client = Client<T>;
fn make(client: &Self::Client, model: impl Into<String>) -> Self {
Self::new(client.clone(), model)
}
async fn completion(
&self,
completion_request: CompletionRequest,
) -> Result<completion::CompletionResponse<CompletionResponse>, CompletionError> {
let span = if tracing::Span::current().is_disabled() {
info_span!(
target: "rig::completions",
"chat",
gen_ai.operation.name = "chat",
gen_ai.provider.name = "mira",
gen_ai.request.model = self.model,
gen_ai.system_instructions = tracing::field::Empty,
gen_ai.response.id = tracing::field::Empty,
gen_ai.response.model = tracing::field::Empty,
gen_ai.usage.output_tokens = tracing::field::Empty,
gen_ai.usage.input_tokens = tracing::field::Empty,
gen_ai.usage.cached_tokens = tracing::field::Empty,
)
} else {
tracing::Span::current()
};
span.record("gen_ai.system_instructions", &completion_request.preamble);
if !completion_request.tools.is_empty() {
tracing::warn!(target: "rig::completions",
"Tool calls are not supported by Mira AI. {len} tools will be ignored.",
len = completion_request.tools.len()
);
}
if completion_request.tool_choice.is_some() {
tracing::warn!("WARNING: `tool_choice` not supported on Mira AI");
}
if completion_request.additional_params.is_some() {
tracing::warn!("WARNING: Additional parameters not supported on Mira AI");
}
let request = MiraCompletionRequest::try_from((self.model.as_ref(), completion_request))?;
if tracing::enabled!(tracing::Level::TRACE) {
tracing::trace!(target: "rig::completions",
"Mira completion request: {}",
serde_json::to_string_pretty(&request)?
);
}
let body = serde_json::to_vec(&request)?;
let req = self
.client
.post("/v1/chat/completions")?
.body(body)
.map_err(http_client::Error::from)?;
let async_block = async move {
let response = self
.client
.send::<_, bytes::Bytes>(req)
.await
.map_err(|e| CompletionError::ProviderError(e.to_string()))?;
let status = response.status();
let response_body = response.into_body().into_future().await?.to_vec();
if !status.is_success() {
let status = status.as_u16();
let error_text = String::from_utf8_lossy(&response_body).to_string();
return Err(CompletionError::ProviderError(format!(
"API error: {status} - {error_text}"
)));
}
let response: CompletionResponse = serde_json::from_slice(&response_body)?;
if tracing::enabled!(tracing::Level::TRACE) {
tracing::trace!(target: "rig::completions",
"Mira completion response: {}",
serde_json::to_string_pretty(&response)?
);
}
if let CompletionResponse::Structured {
id, model, usage, ..
} = &response
{
let span = tracing::Span::current();
span.record("gen_ai.response.model_name", model);
span.record("gen_ai.response.id", id);
if let Some(usage) = usage {
span.record("gen_ai.usage.input_tokens", usage.prompt_tokens);
span.record(
"gen_ai.usage.output_tokens",
usage.total_tokens - usage.prompt_tokens,
);
}
}
response.try_into()
};
async_block.instrument(span).await
}
async fn stream(
&self,
completion_request: CompletionRequest,
) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
let span = if tracing::Span::current().is_disabled() {
info_span!(
target: "rig::completions",
"chat_streaming",
gen_ai.operation.name = "chat_streaming",
gen_ai.provider.name = "mira",
gen_ai.request.model = self.model,
gen_ai.system_instructions = tracing::field::Empty,
gen_ai.response.id = tracing::field::Empty,
gen_ai.response.model = tracing::field::Empty,
gen_ai.usage.output_tokens = tracing::field::Empty,
gen_ai.usage.input_tokens = tracing::field::Empty,
gen_ai.usage.cached_tokens = tracing::field::Empty,
)
} else {
tracing::Span::current()
};
span.record("gen_ai.system_instructions", &completion_request.preamble);
if !completion_request.tools.is_empty() {
tracing::warn!(target: "rig::completions",
"Tool calls are not supported by Mira AI. {len} tools will be ignored.",
len = completion_request.tools.len()
);
}
if completion_request.tool_choice.is_some() {
tracing::warn!("WARNING: `tool_choice` not supported on Mira AI");
}
if completion_request.additional_params.is_some() {
tracing::warn!("WARNING: Additional parameters not supported on Mira AI");
}
let mut request =
MiraCompletionRequest::try_from((self.model.as_ref(), completion_request))?;
request.stream = true;
if tracing::enabled!(tracing::Level::TRACE) {
tracing::trace!(target: "rig::completions",
"Mira completion request: {}",
serde_json::to_string_pretty(&request)?
);
}
let body = serde_json::to_vec(&request)?;
let req = self
.client
.post("/v1/chat/completions")?
.body(body)
.map_err(http_client::Error::from)?;
send_compatible_streaming_request(self.client.clone(), req)
.instrument(span)
.await
}
}
impl From<ApiErrorResponse> for CompletionError {
fn from(err: ApiErrorResponse) -> Self {
CompletionError::ProviderError(err.message)
}
}
impl TryFrom<CompletionResponse> for completion::CompletionResponse<CompletionResponse> {
type Error = CompletionError;
fn try_from(response: CompletionResponse) -> Result<Self, Self::Error> {
let (content, usage) = match &response {
CompletionResponse::Structured { choices, usage, .. } => {
let choice = choices.first().ok_or_else(|| {
CompletionError::ResponseError("Response contained no choices".to_owned())
})?;
let usage = usage
.as_ref()
.map(|usage| completion::Usage {
input_tokens: usage.prompt_tokens as u64,
output_tokens: (usage.total_tokens - usage.prompt_tokens) as u64,
total_tokens: usage.total_tokens as u64,
cached_input_tokens: 0,
cache_creation_input_tokens: 0,
})
.unwrap_or_default();
let message = message::Message::try_from(choice.message.clone())?;
let content = match message {
Message::Assistant { content, .. } => {
if content.is_empty() {
return Err(CompletionError::ResponseError(
"Response contained empty content".to_owned(),
));
}
for c in content.iter() {
if !matches!(c, AssistantContent::Text(_)) {
tracing::warn!(target: "rig",
"Unsupported content type encountered: {:?}. The Mira provider currently only supports text content", c
);
}
}
content.iter().map(|c| {
match c {
AssistantContent::Text(text) => Ok(completion::AssistantContent::text(&text.text)),
other => Err(CompletionError::ResponseError(
format!("Unsupported content type: {other:?}. The Mira provider currently only supports text content")
))
}
}).collect::<Result<Vec<_>, _>>()?
}
Message::User { .. } => {
tracing::warn!(target: "rig", "Received user message in response where assistant message was expected");
return Err(CompletionError::ResponseError(
"Received user message in response where assistant message was expected".to_owned()
));
}
Message::System { .. } => {
tracing::warn!(target: "rig", "Received system message in response where assistant message was expected");
return Err(CompletionError::ResponseError(
"Received system message in response where assistant message was expected".to_owned(),
));
}
};
(content, usage)
}
CompletionResponse::Simple(text) => (
vec![completion::AssistantContent::text(text)],
completion::Usage::new(),
),
};
let choice = OneOrMany::many(content).map_err(|_| {
CompletionError::ResponseError(
"Response contained no message or tool call (empty)".to_owned(),
)
})?;
Ok(completion::CompletionResponse {
choice,
usage,
raw_response: response,
message_id: None,
})
}
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Usage {
pub prompt_tokens: usize,
pub total_tokens: usize,
}
impl std::fmt::Display for Usage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"Prompt tokens: {} Total tokens: {}",
self.prompt_tokens, self.total_tokens
)
}
}
impl From<Message> for serde_json::Value {
fn from(msg: Message) -> Self {
match msg {
Message::System { content } => serde_json::json!({
"role": "system",
"content": content
}),
Message::User { content } => {
let text = content
.iter()
.map(|c| match c {
UserContent::Text(text) => &text.text,
_ => "",
})
.collect::<Vec<_>>()
.join("\n");
serde_json::json!({
"role": "user",
"content": text
})
}
Message::Assistant { content, .. } => {
let text = content
.iter()
.map(|c| match c {
AssistantContent::Text(text) => &text.text,
_ => "",
})
.collect::<Vec<_>>()
.join("\n");
serde_json::json!({
"role": "assistant",
"content": text
})
}
}
}
}
impl TryFrom<serde_json::Value> for Message {
type Error = CompletionError;
fn try_from(value: serde_json::Value) -> Result<Self, Self::Error> {
let role = value["role"].as_str().ok_or_else(|| {
CompletionError::ResponseError("Message missing role field".to_owned())
})?;
let content = match value.get("content") {
Some(content) => match content {
serde_json::Value::String(s) => s.clone(),
serde_json::Value::Array(arr) => arr
.iter()
.filter_map(|c| {
c.get("text")
.and_then(|t| t.as_str())
.map(|text| text.to_string())
})
.collect::<Vec<_>>()
.join("\n"),
_ => {
return Err(CompletionError::ResponseError(
"Message content must be string or array".to_owned(),
));
}
},
None => {
return Err(CompletionError::ResponseError(
"Message missing content field".to_owned(),
));
}
};
match role {
"system" => Ok(Message::System { content }),
"user" => Ok(Message::User {
content: OneOrMany::one(UserContent::Text(message::Text { text: content })),
}),
"assistant" => Ok(Message::Assistant {
id: None,
content: OneOrMany::one(AssistantContent::Text(message::Text { text: content })),
}),
_ => Err(CompletionError::ResponseError(format!(
"Unsupported message role: {role}"
))),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::message::UserContent;
use serde_json::json;
#[test]
fn test_deserialize_message() {
let assistant_message_json = json!({
"role": "assistant",
"content": "Hello there, how may I assist you today?"
});
let user_message_json = json!({
"role": "user",
"content": "What can you help me with?"
});
let assistant_message_array_json = json!({
"role": "assistant",
"content": [{
"type": "text",
"text": "Hello there, how may I assist you today?"
}]
});
let assistant_message = Message::try_from(assistant_message_json).unwrap();
let user_message = Message::try_from(user_message_json).unwrap();
let assistant_message_array = Message::try_from(assistant_message_array_json).unwrap();
match assistant_message {
Message::Assistant { content, .. } => {
assert_eq!(
content.first(),
AssistantContent::Text(message::Text {
text: "Hello there, how may I assist you today?".to_string()
})
);
}
_ => panic!("Expected assistant message"),
}
match user_message {
Message::User { content } => {
assert_eq!(
content.first(),
UserContent::Text(message::Text {
text: "What can you help me with?".to_string()
})
);
}
_ => panic!("Expected user message"),
}
match assistant_message_array {
Message::Assistant { content, .. } => {
assert_eq!(
content.first(),
AssistantContent::Text(message::Text {
text: "Hello there, how may I assist you today?".to_string()
})
);
}
_ => panic!("Expected assistant message"),
}
}
#[test]
fn test_message_conversion() {
let original_message = message::Message::User {
content: OneOrMany::one(message::UserContent::text("Hello")),
};
let mira_value: serde_json::Value = original_message.clone().into();
let converted_message: Message = mira_value.try_into().unwrap();
assert_eq!(original_message, converted_message);
}
#[test]
fn test_completion_response_conversion() {
let mira_response = CompletionResponse::Structured {
id: "resp_123".to_string(),
object: "chat.completion".to_string(),
created: 1234567890,
model: "deepseek-r1".to_string(),
choices: vec![ChatChoice {
message: RawMessage {
role: "assistant".to_string(),
content: "Test response".to_string(),
},
finish_reason: Some("stop".to_string()),
index: Some(0),
}],
usage: Some(Usage {
prompt_tokens: 10,
total_tokens: 20,
}),
};
let completion_response: completion::CompletionResponse<CompletionResponse> =
mira_response.try_into().unwrap();
assert_eq!(
completion_response.choice.first(),
completion::AssistantContent::text("Test response")
);
}
#[test]
fn test_client_initialization() {
let _client =
crate::providers::mira::Client::new("dummy-key").expect("Client::new() failed");
let _client_from_builder = crate::providers::mira::Client::builder()
.api_key("dummy-key")
.build()
.expect("Client::builder() failed");
}
}