This commit is contained in:
qpismont committed
1 parent
9f03384bc9
commit
f5f0d3999e
30 files changed
+96
-93
No files matched your search
@@ -0,0 +1,473 @@
|
||||
//! Minimal OpenRouter chat-completions client.
|
||||
//!
|
||||
//! Steep only needs a non-streaming `POST /chat/completions` with optional
|
||||
//! tool calling, so the wire types are implemented in-tree instead of pulling a
|
||||
//! third-party SDK (and its own `reqwest` version) into the workspace.
|
||||
//!
|
||||
//! Only the response fields Steep consumes are modelled: `content`,
|
||||
//! `tool_calls` and `usage.cost`. Unknown fields are ignored.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Context;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tracing::instrument;
|
||||
|
||||
/// OpenRouter API root, version prefix included.
|
||||
const BASE_URL: &str = "https://openrouter.ai/api/v1";
|
||||
|
||||
/// What the model is allowed to do on a given turn.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ToolChoice {
|
||||
/// The model decides whether to call a tool.
|
||||
Auto,
|
||||
/// The model must answer, without calling any tool. The tools stay declared,
|
||||
/// so the conversation remains the same shape as on the other turns.
|
||||
None,
|
||||
}
|
||||
|
||||
impl ToolChoice {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Auto => "auto",
|
||||
Self::None => "none",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a completion that may contain tool calls.
|
||||
pub struct ToolChatResult {
|
||||
pub message: Option<String>,
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
pub cost: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum Role {
|
||||
System,
|
||||
User,
|
||||
Assistant,
|
||||
Tool,
|
||||
}
|
||||
|
||||
/// A single turn of the conversation.
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct Message {
|
||||
pub role: Role,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_call_id: Option<String>,
|
||||
}
|
||||
|
||||
impl Message {
|
||||
pub fn new(role: Role, content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
role,
|
||||
content: Some(content.into()),
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Assistant turn requesting tool calls. Models often answer with tool
|
||||
/// calls but no text, in which case the content is sent as `null`.
|
||||
pub fn assistant_with_tool_calls(
|
||||
content: impl Into<String>,
|
||||
tool_calls: Vec<ToolCall>,
|
||||
) -> Self {
|
||||
let content = content.into();
|
||||
|
||||
Self {
|
||||
role: Role::Assistant,
|
||||
content: (!content.is_empty()).then_some(content),
|
||||
tool_calls: Some(tool_calls),
|
||||
tool_call_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a tool call, linked to the request by `tool_call_id`.
|
||||
pub fn tool_response(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
|
||||
Self {
|
||||
role: Role::Tool,
|
||||
content: Some(content.into()),
|
||||
tool_calls: None,
|
||||
tool_call_id: Some(tool_call_id.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A tool the model may call.
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct Tool {
|
||||
#[serde(rename = "type", default = "function_type")]
|
||||
pub kind: String,
|
||||
pub function: FunctionDefinition,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct FunctionDefinition {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
/// JSON schema describing the accepted arguments.
|
||||
pub parameters: Value,
|
||||
}
|
||||
|
||||
impl Tool {
|
||||
pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: Value) -> Self {
|
||||
Self {
|
||||
kind: function_type(),
|
||||
function: FunctionDefinition {
|
||||
name: name.into(),
|
||||
description: description.into(),
|
||||
parameters,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A tool call requested by the model.
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct ToolCall {
|
||||
pub id: String,
|
||||
#[serde(rename = "type", default = "function_type")]
|
||||
pub kind: String,
|
||||
pub function: FunctionCall,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct FunctionCall {
|
||||
pub name: String,
|
||||
/// Arguments as a JSON-encoded string. Kept verbatim so that echoing the
|
||||
/// call back into the conversation does not re-encode or corrupt it.
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
impl ToolCall {
|
||||
pub fn id(&self) -> &str {
|
||||
&self.id
|
||||
}
|
||||
|
||||
pub fn name(&self) -> &str {
|
||||
&self.function.name
|
||||
}
|
||||
|
||||
pub fn arguments_json(&self) -> &str {
|
||||
&self.function.arguments
|
||||
}
|
||||
}
|
||||
|
||||
fn function_type() -> String {
|
||||
String::from("function")
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChatRequest<'a> {
|
||||
model: &'a str,
|
||||
messages: &'a [Message],
|
||||
reasoning: Reasoning,
|
||||
tools: &'a [Tool],
|
||||
tool_choice: &'a str,
|
||||
}
|
||||
|
||||
/// OpenRouter reasoning config; Steep always asks for reasoning.
|
||||
#[derive(Serialize)]
|
||||
struct Reasoning {
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatResponse {
|
||||
#[serde(default)]
|
||||
choices: Vec<Choice>,
|
||||
#[serde(default)]
|
||||
usage: Option<Usage>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Choice {
|
||||
message: ResponseMessage,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ResponseMessage {
|
||||
#[serde(default)]
|
||||
content: Option<Content>,
|
||||
#[serde(default)]
|
||||
tool_calls: Option<Vec<ToolCall>>,
|
||||
}
|
||||
|
||||
/// Message content is a plain string for most models and a list of typed parts
|
||||
/// for others; both are collapsed to text.
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum Content {
|
||||
Text(String),
|
||||
Parts(Vec<ContentPart>),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ContentPart {
|
||||
#[serde(default)]
|
||||
text: Option<String>,
|
||||
}
|
||||
|
||||
impl Content {
|
||||
fn into_text(self) -> String {
|
||||
match self {
|
||||
Self::Text(text) => text,
|
||||
Self::Parts(parts) => parts
|
||||
.into_iter()
|
||||
.filter_map(|part| part.text)
|
||||
.collect::<Vec<_>>()
|
||||
.join(""),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Usage {
|
||||
#[serde(default)]
|
||||
cost: Option<f64>,
|
||||
}
|
||||
|
||||
/// Error payload returned by OpenRouter for a failed request.
|
||||
#[derive(Deserialize)]
|
||||
struct ErrorResponse {
|
||||
error: ErrorDetail,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ErrorDetail {
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OpenRouterClient {
|
||||
client: reqwest::Client,
|
||||
api_key: String,
|
||||
model: String,
|
||||
}
|
||||
|
||||
impl OpenRouterClient {
|
||||
pub fn new(token: &str, model: &str, timeout: u64) -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
client: reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(timeout))
|
||||
.build()?,
|
||||
api_key: String::from(token),
|
||||
model: String::from(model),
|
||||
})
|
||||
}
|
||||
|
||||
/// Sends a conversation with tool definitions and returns either a final
|
||||
/// message or the tool calls requested by the model.
|
||||
///
|
||||
/// The conversation is borrowed rather than taken by value: the caller keeps
|
||||
/// appending to it between iterations, and a copy per iteration would be pure
|
||||
/// waste.
|
||||
#[instrument(skip(self, messages, tools), err)]
|
||||
pub async fn chat_with_tools(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
tool_choice: ToolChoice,
|
||||
) -> anyhow::Result<ToolChatResult> {
|
||||
let response = self.complete(messages, tools, tool_choice).await?;
|
||||
|
||||
let cost = response.usage.and_then(|usage| usage.cost);
|
||||
let message = response
|
||||
.choices
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| anyhow::anyhow!("No choices in response"))?
|
||||
.message;
|
||||
|
||||
Ok(ToolChatResult {
|
||||
message: message.content.map(Content::into_text),
|
||||
tool_calls: message.tool_calls.unwrap_or_default(),
|
||||
cost,
|
||||
})
|
||||
}
|
||||
|
||||
async fn complete(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
tool_choice: ToolChoice,
|
||||
) -> anyhow::Result<ChatResponse> {
|
||||
let request = ChatRequest {
|
||||
model: &self.model,
|
||||
messages,
|
||||
reasoning: Reasoning { enabled: true },
|
||||
tools,
|
||||
tool_choice: tool_choice.as_str(),
|
||||
};
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(format!("{BASE_URL}/chat/completions"))
|
||||
.bearer_auth(&self.api_key)
|
||||
.json(&request)
|
||||
.send()
|
||||
.await
|
||||
.context("failed to reach OpenRouter")?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
anyhow::bail!("OpenRouter returned {status}: {}", error_message(&body));
|
||||
}
|
||||
|
||||
response
|
||||
.json::<ChatResponse>()
|
||||
.await
|
||||
.context("invalid OpenRouter response")
|
||||
}
|
||||
}
|
||||
|
||||
/// Extracts the message from an OpenRouter error body, falling back to the raw
|
||||
/// body when it is not the expected JSON shape.
|
||||
fn error_message(body: &str) -> String {
|
||||
serde_json::from_str::<ErrorResponse>(body)
|
||||
.map(|response| response.error.message)
|
||||
.unwrap_or_else(|_| body.trim().to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
/// Builds a tool call the way the API returns one, so the fixture also
|
||||
/// covers deserialization.
|
||||
fn tool_call(id: &str, name: &str, arguments: &str) -> ToolCall {
|
||||
serde_json::from_value(json!({
|
||||
"id": id,
|
||||
"type": "function",
|
||||
"function": { "name": name, "arguments": arguments }
|
||||
}))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_call_arguments_are_kept_verbatim() {
|
||||
let call = tool_call("call_1", "grep", r#"{"pattern":"fn main"}"#);
|
||||
let message = Message::assistant_with_tool_calls("", vec![call]);
|
||||
|
||||
let serialized = serde_json::to_value(&message).unwrap();
|
||||
|
||||
assert_eq!(serialized["content"], Value::Null);
|
||||
assert_eq!(serialized["tool_calls"][0]["type"], "function");
|
||||
assert_eq!(serialized["tool_calls"][0]["function"]["name"], "grep");
|
||||
assert_eq!(
|
||||
serialized["tool_calls"][0]["function"]["arguments"],
|
||||
r#"{"pattern":"fn main"}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_response_carries_role_and_tool_call_id() {
|
||||
let message = Message::tool_response("call_1", "src/main.rs");
|
||||
let serialized = serde_json::to_value(&message).unwrap();
|
||||
|
||||
assert_eq!(serialized["role"], "tool");
|
||||
assert_eq!(serialized["tool_call_id"], "call_1");
|
||||
assert_eq!(serialized["content"], "src/main.rs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_serializes_tool_choice_and_reasoning() {
|
||||
let messages = [Message::new(Role::User, "hi")];
|
||||
let tools = [Tool::new("ls", "List files", json!({"type": "object"}))];
|
||||
|
||||
let request = serde_json::to_value(ChatRequest {
|
||||
model: "some/model",
|
||||
messages: &messages,
|
||||
reasoning: Reasoning { enabled: true },
|
||||
tools: &tools,
|
||||
tool_choice: ToolChoice::Auto.as_str(),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(request["model"], "some/model");
|
||||
assert_eq!(request["reasoning"]["enabled"], true);
|
||||
assert_eq!(request["tool_choice"], "auto");
|
||||
assert_eq!(request["tools"][0]["function"]["name"], "ls");
|
||||
assert_eq!(request["messages"][0]["role"], "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forced_answer_turn_disables_the_tools() {
|
||||
assert_eq!(ToolChoice::None.as_str(), "none");
|
||||
assert_eq!(ToolChoice::Auto.as_str(), "auto");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_parses_tool_calls_and_cost() {
|
||||
let response: ChatResponse = serde_json::from_value(json!({
|
||||
"choices": [{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": null,
|
||||
"tool_calls": [{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": { "name": "ls", "arguments": "{\"path\":\"src\"}" }
|
||||
}]
|
||||
}
|
||||
}],
|
||||
"usage": { "cost": 0.0021 }
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.usage.and_then(|usage| usage.cost), Some(0.0021));
|
||||
|
||||
let message = response.choices.into_iter().next().unwrap().message;
|
||||
assert!(message.content.is_none());
|
||||
|
||||
let calls = message.tool_calls.unwrap();
|
||||
assert_eq!(calls[0].name(), "ls");
|
||||
assert_eq!(calls[0].id(), "call_1");
|
||||
assert_eq!(calls[0].arguments_json(), r#"{"path":"src"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_parses_parts_content() {
|
||||
let response: ChatResponse = serde_json::from_value(json!({
|
||||
"choices": [{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{ "type": "text", "text": "hello " },
|
||||
{ "type": "text", "text": "world" }
|
||||
]
|
||||
}
|
||||
}]
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert!(response.usage.is_none());
|
||||
|
||||
let message = response.choices.into_iter().next().unwrap().message;
|
||||
assert_eq!(message.content.unwrap().into_text(), "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_message_prefers_api_message() {
|
||||
let body = r#"{"error":{"message":"No auth credentials found","code":401}}"#;
|
||||
assert_eq!(error_message(body), "No auth credentials found");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_message_falls_back_to_raw_body() {
|
||||
assert_eq!(
|
||||
error_message(" <html>bad gateway</html> "),
|
||||
"<html>bad gateway</html>"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user