aboutsummaryrefslogtreecommitdiffstats
path: root/src/openrouter
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--src/openrouter/mod.rs134
1 files changed, 134 insertions, 0 deletions
diff --git a/src/openrouter/mod.rs b/src/openrouter/mod.rs
new file mode 100644
index 0000000..31cee85
--- /dev/null
+++ b/src/openrouter/mod.rs
@@ -0,0 +1,134 @@
1use reqwest::Client;
2use serde::{Deserialize, Serialize};
3
4use crate::error::Result;
5use crate::lang;
6
7pub struct OpenRouterClient {
8 api_key: String,
9 model: String,
10 client: Client,
11}
12
13#[derive(Serialize)]
14struct ChatRequest {
15 model: String,
16 messages: Vec<Message>,
17 max_tokens: u32,
18}
19
20#[derive(Serialize)]
21struct Message {
22 role: String,
23 content: String,
24}
25
26#[derive(Deserialize)]
27struct ChatResponse {
28 choices: Vec<Choice>,
29}
30
31#[derive(Deserialize)]
32struct Choice {
33 message: MessageContent,
34}
35
36#[derive(Deserialize)]
37struct MessageContent {
38 content: String,
39}
40
41impl OpenRouterClient {
42 pub fn new(api_key: String, model: String) -> Self {
43 Self {
44 api_key,
45 model,
46 client: Client::new(),
47 }
48 }
49
50 pub async fn chat(&self, system_prompt: &str, user_message: &str) -> Result<String> {
51 let body = ChatRequest {
52 model: self.model.clone(),
53 messages: vec![
54 Message {
55 role: "system".into(),
56 content: system_prompt.to_string(),
57 },
58 Message {
59 role: "user".into(),
60 content: user_message.to_string(),
61 },
62 ],
63 max_tokens: 500,
64 };
65
66 let resp = self
67 .client
68 .post("https://openrouter.ai/api/v1/chat/completions")
69 .header("Authorization", format!("Bearer {}", self.api_key))
70 .json(&body)
71 .send()
72 .await?;
73
74 if !resp.status().is_success() {
75 let status = resp.status();
76 let text = resp.text().await.unwrap_or_default();
77 return Err(crate::error::Error::Custom(lang::fmt_openrouter_err(
78 &status, &text,
79 )));
80 }
81
82 let chat_resp: ChatResponse = resp.json().await?;
83 Ok(chat_resp
84 .choices
85 .into_iter()
86 .next()
87 .map(|c| c.message.content)
88 .unwrap_or_default())
89 }
90
91 pub async fn chat_with_history(
92 &self,
93 system_prompt: &str,
94 messages: Vec<(String, String)>, // (role, content) pairs
95 ) -> Result<String> {
96 let mut msg_vec = vec![Message {
97 role: "system".into(),
98 content: system_prompt.to_string(),
99 }];
100 for (role, content) in messages {
101 msg_vec.push(Message { role, content });
102 }
103
104 let body = ChatRequest {
105 model: self.model.clone(),
106 messages: msg_vec,
107 max_tokens: 500,
108 };
109
110 let resp = self
111 .client
112 .post("https://openrouter.ai/api/v1/chat/completions")
113 .header("Authorization", format!("Bearer {}", self.api_key))
114 .json(&body)
115 .send()
116 .await?;
117
118 if !resp.status().is_success() {
119 let status = resp.status();
120 let text = resp.text().await.unwrap_or_default();
121 return Err(crate::error::Error::Custom(lang::fmt_openrouter_err(
122 &status, &text,
123 )));
124 }
125
126 let chat_resp: ChatResponse = resp.json().await?;
127 Ok(chat_resp
128 .choices
129 .into_iter()
130 .next()
131 .map(|c| c.message.content)
132 .unwrap_or_default())
133 }
134}