93 lines
2.6 KiB
Rust
93 lines
2.6 KiB
Rust
use std::time::Duration;
|
|
|
|
use openrouter_rs::{
|
|
Message,
|
|
api::chat::ChatCompletionRequest,
|
|
types::{Tool, ToolCall},
|
|
};
|
|
use tracing::instrument;
|
|
|
|
pub struct ChatResult {
|
|
pub message: String,
|
|
pub cost: Option<f64>,
|
|
}
|
|
|
|
pub struct ToolChatResult {
|
|
pub message: Option<String>,
|
|
pub tool_calls: Vec<ToolCall>,
|
|
pub cost: Option<f64>,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct OpenRouterClient {
|
|
client: openrouter_rs::OpenRouterClient,
|
|
model: String,
|
|
}
|
|
|
|
impl OpenRouterClient {
|
|
pub fn new(token: &str, model: &str, timeout: u64) -> anyhow::Result<Self> {
|
|
Ok(Self {
|
|
client: openrouter_rs::OpenRouterClient::builder()
|
|
.api_key(token)
|
|
.http_client(
|
|
reqwest::Client::builder()
|
|
.timeout(Duration::from_secs(timeout))
|
|
.build()?,
|
|
)
|
|
.build()?,
|
|
model: String::from(model),
|
|
})
|
|
}
|
|
|
|
#[instrument(skip(self), err)]
|
|
pub async fn chat(&self, msg: &str) -> anyhow::Result<ChatResult> {
|
|
let request = ChatCompletionRequest::builder()
|
|
.model(&self.model)
|
|
.enable_reasoning()
|
|
.messages(vec![Message::new(openrouter_rs::types::Role::User, msg)])
|
|
.build()?;
|
|
|
|
let response = self.client.chat().create(&request).await?;
|
|
|
|
Ok(ChatResult {
|
|
message: response.choices[0]
|
|
.content()
|
|
.map(String::from)
|
|
.ok_or(anyhow::anyhow!("No content"))?,
|
|
cost: response.usage.and_then(|u| u.cost),
|
|
})
|
|
}
|
|
|
|
/// Sends a conversation with tool definitions and returns either a final
|
|
/// message or the tool calls requested by the model.
|
|
#[instrument(skip(self, messages, tools), err)]
|
|
pub async fn chat_with_tools(
|
|
&self,
|
|
messages: Vec<Message>,
|
|
tools: Vec<Tool>,
|
|
) -> anyhow::Result<ToolChatResult> {
|
|
let request = ChatCompletionRequest::builder()
|
|
.model(&self.model)
|
|
.enable_reasoning()
|
|
.messages(messages)
|
|
.tools(tools)
|
|
.tool_choice_auto()
|
|
.build()?;
|
|
|
|
let response = self.client.chat().create(&request).await?;
|
|
let choice = response
|
|
.choices
|
|
.first()
|
|
.ok_or_else(|| anyhow::anyhow!("No choices in response"))?;
|
|
|
|
Ok(ToolChatResult {
|
|
message: choice.content().map(String::from),
|
|
tool_calls: choice
|
|
.tool_calls()
|
|
.map(<[ToolCall]>::to_vec)
|
|
.unwrap_or_default(),
|
|
cost: response.usage.and_then(|u| u.cost),
|
|
})
|
|
}
|
|
}
|