use std::collections::HashMap;
use async_trait::async_trait;
use mermaid_domain::ChatRequest;
use mermaid_model::models::adapters::meta::MetaAdapter;
use mermaid_model::models::{Model, ModelCapabilities, ModelConfig, ModelError, Result};
use super::super::ctx::{FinalResponse, StreamContext, StreamEvent};
use super::{ModelProvider, RejectionCache};
pub use mermaid_model::models::adapters::meta::{DEFAULT_API_KEY_ENV, DEFAULT_BASE_URL};
pub struct MetaProvider {
adapter: MetaAdapter,
capabilities: ModelCapabilities,
rejections: RejectionCache,
}
impl MetaProvider {
pub fn new(
api_key: String,
model_name: String,
base_url: String,
extra_headers: HashMap<String, String>,
) -> Result<Self> {
let adapter = MetaAdapter::new(api_key, model_name, base_url, extra_headers)?;
let capabilities = Model::capabilities(&adapter).clone();
Ok(Self {
adapter,
capabilities,
rejections: RejectionCache::default(),
})
}
}
#[async_trait]
impl ModelProvider for MetaProvider {
fn capabilities(&self) -> &ModelCapabilities {
&self.capabilities
}
async fn chat(&self, request: ChatRequest, ctx: StreamContext) -> Result<FinalResponse> {
let config = ModelConfig::from(&request);
let chat_fut = self
.adapter
.chat(&request.messages, &config, Some(ctx.sink.clone()));
let learned_model = Model::name(&self.adapter).to_string();
let memory = self.adapter.param_memory();
self.rejections.seed("meta", &learned_model, memory).await;
let response = tokio::select! {
biased;
_ = ctx.token.cancelled() => {
return Err(ModelError::Cancelled);
},
r = chat_fut => {
self.rejections.persist("meta", &learned_model, memory).await;
r?
},
};
let usage = response.usage.clone();
let provider_continuation = response.provider_continuation.clone();
let stop_reason = response.stop_reason.clone();
let _ = ctx
.sink
.send(StreamEvent::Done {
usage: usage.clone(),
provider_continuation: provider_continuation.clone(),
stop_reason: stop_reason.clone(),
})
.await;
Ok(FinalResponse {
usage,
provider_continuation,
tool_calls: response.tool_calls.unwrap_or_default(),
stop_reason,
})
}
}