a_agent/provider/
chat_completion.rs

1use std::collections::BTreeMap;
2
3use anyhow::Result;
4use async_openai::Client;
5use async_openai::config::OpenAIConfig;
6use async_openai::types::stream::StreamResponse;
7use async_trait::async_trait;
8use futures_util::StreamExt;
9use serde_json::Value;
10use tokio_util::sync::CancellationToken;
11
12use crate::config::ProviderConfig;
13use crate::model::{ContentBlock, ModelRequest, ModelTurn, Role, StreamEvent, ToolCall, Usage};
14
15use super::{EventSink, Provider, merge_request_fields, tool_definitions};
16
17pub struct ChatCompletionProvider {
18    client: Client<OpenAIConfig>,
19    config: ProviderConfig,
20}
21
22impl ChatCompletionProvider {
23    pub fn new(config: ProviderConfig, api_key: String) -> Result<Self> {
24        let base_url = config
25            .base_url
26            .clone()
27            .unwrap_or_else(|| "https://api.openai.com/v1".into());
28        let mut sdk_config = OpenAIConfig::new()
29            .with_api_key(api_key)
30            .with_api_base(base_url.trim_end_matches('/'));
31        for (key, value) in &config.headers {
32            sdk_config = sdk_config.with_header(
33                reqwest::header::HeaderName::from_bytes(key.as_bytes())?,
34                value.as_str(),
35            )?;
36        }
37        Ok(Self {
38            client: Client::with_config(sdk_config),
39            config,
40        })
41    }
42
43    fn request_body(&self, request: ModelRequest) -> Value {
44        let mut messages =
45            vec![serde_json::json!({"role":"system","content":request.system_prompt})];
46        for message in request.messages {
47            match message.role {
48                Role::User => messages.push(
49                    serde_json::json!({"role":"user","content":text_blocks(&message.blocks)}),
50                ),
51                Role::Assistant => {
52                    let calls = message.blocks.iter().filter_map(|block| match block {
53                        ContentBlock::ToolCall(call) => Some(serde_json::json!({
54                            "id":call.id,"type":"function","function":{"name":call.name,"arguments":call.arguments}
55                        })), _ => None
56                    }).collect::<Vec<_>>();
57                    let text = text_blocks(&message.blocks);
58                    let mut item = serde_json::json!({"role":"assistant","content":if text.is_empty() { Value::Null } else { Value::String(text) }});
59                    if !calls.is_empty() {
60                        item["tool_calls"] = Value::Array(calls);
61                    }
62                    messages.push(item);
63                }
64                Role::Tool => {
65                    for block in message.blocks {
66                        if let ContentBlock::ToolResult(result) = block {
67                            messages.push(serde_json::json!({"role":"tool","tool_call_id":result.call_id,"content":result.output}));
68                        }
69                    }
70                }
71                Role::System => {}
72            }
73        }
74        let tools = if request.include_tools {
75            tool_definitions()
76                .into_iter()
77                .map(|tool| serde_json::json!({"type":"function","function":tool}))
78                .collect()
79        } else {
80            Vec::new()
81        };
82        let mut body = serde_json::Map::new();
83        merge_request_fields(&mut body, &self.config);
84        body.entry("stream_options")
85            .or_insert_with(|| serde_json::json!({"include_usage":true}));
86        body.insert("model".into(), Value::String(self.config.model.clone()));
87        body.insert("max_tokens".into(), Value::from(self.config.max_tokens));
88        body.insert("messages".into(), Value::Array(messages));
89        body.insert("tools".into(), Value::Array(tools));
90        body.insert("stream".into(), Value::Bool(true));
91        Value::Object(body)
92    }
93}
94
95#[async_trait]
96impl Provider for ChatCompletionProvider {
97    async fn stream_turn(
98        &self,
99        request: ModelRequest,
100        events: EventSink,
101        cancel: CancellationToken,
102    ) -> Result<ModelTurn> {
103        let chat = self.client.chat();
104        let create = chat.create_stream_byot(self.request_body(request));
105        tokio::pin!(create);
106        let mut stream: StreamResponse<Value> = tokio::select! {
107            _ = cancel.cancelled() => anyhow::bail!("Chat Completions request cancelled"),
108            result = &mut create => result?,
109        };
110        let mut values = Vec::new();
111        let mut live = ChatLive::default();
112        loop {
113            tokio::select! {
114                _ = cancel.cancelled() => anyhow::bail!("Chat Completions request cancelled"),
115                item = stream.next() => match item {
116                    Some(Ok(value)) => { live.emit(&value, &events); values.push(value); }
117                    Some(Err(error)) => return Err(error.into()),
118                    None => break,
119                }
120            }
121        }
122        normalize_events(values).map(|(turn, _)| turn)
123    }
124}
125
126#[derive(Default)]
127struct ChatLive {
128    calls: BTreeMap<usize, (String, String, bool)>,
129}
130
131impl ChatLive {
132    fn emit(&mut self, value: &Value, sink: &EventSink) {
133        if let Some(raw) = value.get("usage") {
134            sink.emit(StreamEvent::Usage(normalize_usage(raw)));
135        }
136        let Some(delta) = value.pointer("/choices/0/delta") else {
137            return;
138        };
139        if let Some(part) = delta.get("content").and_then(Value::as_str) {
140            sink.emit(StreamEvent::TextDelta { delta: part.into() });
141        }
142        if let Some(part) = delta
143            .get("reasoning_content")
144            .or_else(|| delta.get("reasoning"))
145            .and_then(Value::as_str)
146        {
147            sink.emit(StreamEvent::ReasoningDelta { delta: part.into() });
148        }
149        for raw in delta
150            .get("tool_calls")
151            .and_then(Value::as_array)
152            .into_iter()
153            .flatten()
154        {
155            let index = raw.get("index").and_then(Value::as_u64).unwrap_or_default() as usize;
156            let state = self.calls.entry(index).or_default();
157            if let Some(id) = raw.get("id").and_then(Value::as_str) {
158                state.0.push_str(id);
159            }
160            if let Some(name) = raw.pointer("/function/name").and_then(Value::as_str) {
161                state.1.push_str(name);
162            }
163            if !state.2 && !state.0.is_empty() && !state.1.is_empty() {
164                state.2 = true;
165                sink.emit(StreamEvent::ToolCallStart {
166                    id: state.0.clone(),
167                    name: state.1.clone(),
168                });
169            }
170            if let Some(part) = raw.pointer("/function/arguments").and_then(Value::as_str) {
171                sink.emit(StreamEvent::ToolCallArgsDelta {
172                    id: state.0.clone(),
173                    delta: part.into(),
174                });
175            }
176        }
177        if value
178            .pointer("/choices/0/finish_reason")
179            .is_some_and(|value| !value.is_null())
180        {
181            for (id, _, started) in self.calls.values() {
182                if *started {
183                    sink.emit(StreamEvent::ToolCallEnd { id: id.clone() });
184                }
185            }
186            sink.emit(StreamEvent::Done);
187        }
188    }
189}
190
191fn text_blocks(blocks: &[ContentBlock]) -> String {
192    blocks
193        .iter()
194        .filter_map(|block| match block {
195            ContentBlock::Text(text) => Some(text.as_str()),
196            _ => None,
197        })
198        .collect::<Vec<_>>()
199        .join("\n")
200}
201
202pub fn normalize_events(values: Vec<Value>) -> Result<(ModelTurn, Vec<StreamEvent>)> {
203    let mut text = String::new();
204    let mut reasoning = String::new();
205    let mut calls: BTreeMap<usize, ToolCall> = BTreeMap::new();
206    let mut started = BTreeMap::new();
207    let mut events = Vec::new();
208    let mut usage = None;
209
210    for value in values {
211        if let Some(error) = value.get("error") {
212            anyhow::bail!("provider error: {error}");
213        }
214        if let Some(raw) = value.get("usage") {
215            usage = Some(normalize_usage(raw));
216        }
217        let Some(delta) = value.pointer("/choices/0/delta") else {
218            continue;
219        };
220        if let Some(part) = delta.get("content").and_then(Value::as_str) {
221            text.push_str(part);
222            events.push(StreamEvent::TextDelta { delta: part.into() });
223        }
224        if let Some(part) = delta
225            .get("reasoning_content")
226            .or_else(|| delta.get("reasoning"))
227            .and_then(Value::as_str)
228        {
229            reasoning.push_str(part);
230            events.push(StreamEvent::ReasoningDelta { delta: part.into() });
231        }
232        for raw in delta
233            .get("tool_calls")
234            .and_then(Value::as_array)
235            .into_iter()
236            .flatten()
237        {
238            let index = raw.get("index").and_then(Value::as_u64).unwrap_or_default() as usize;
239            let call = calls
240                .entry(index)
241                .or_insert_with(|| ToolCall::new("", "", ""));
242            if let Some(id) = raw.get("id").and_then(Value::as_str) {
243                call.id.push_str(id);
244            }
245            if let Some(name) = raw.pointer("/function/name").and_then(Value::as_str) {
246                call.name.push_str(name);
247            }
248            if !started.get(&index).copied().unwrap_or(false)
249                && !call.id.is_empty()
250                && !call.name.is_empty()
251            {
252                events.push(StreamEvent::ToolCallStart {
253                    id: call.id.clone(),
254                    name: call.name.clone(),
255                });
256                started.insert(index, true);
257            }
258            if let Some(part) = raw.pointer("/function/arguments").and_then(Value::as_str) {
259                call.arguments.push_str(part);
260                events.push(StreamEvent::ToolCallArgsDelta {
261                    id: call.id.clone(),
262                    delta: part.into(),
263                });
264            }
265        }
266    }
267    let tool_calls = calls.into_values().collect::<Vec<_>>();
268    for call in &tool_calls {
269        events.push(StreamEvent::ToolCallEnd {
270            id: call.id.clone(),
271        });
272    }
273    let mut blocks = Vec::new();
274    if !reasoning.is_empty() {
275        blocks.push(ContentBlock::Reasoning(reasoning));
276    }
277    if !text.is_empty() {
278        blocks.push(ContentBlock::Text(text));
279    }
280    blocks.extend(tool_calls.iter().cloned().map(ContentBlock::ToolCall));
281    if let Some(usage) = usage {
282        events.push(StreamEvent::Usage(usage));
283    }
284    events.push(StreamEvent::Done);
285    Ok((
286        ModelTurn {
287            blocks,
288            tool_calls,
289            usage,
290            provider_state: None,
291        },
292        events,
293    ))
294}
295
296fn normalize_usage(raw: &Value) -> Usage {
297    let cached_tokens = raw
298        .pointer("/prompt_tokens_details/cached_tokens")
299        .and_then(Value::as_u64)
300        .or_else(|| raw.get("prompt_cache_hit_tokens").and_then(Value::as_u64))
301        .or_else(|| raw.get("cached_tokens").and_then(Value::as_u64));
302    let cache_write_tokens = raw
303        .pointer("/prompt_tokens_details/cache_write_tokens")
304        .and_then(Value::as_u64);
305    Usage {
306        input_tokens: raw
307            .get("prompt_tokens")
308            .and_then(Value::as_u64)
309            .map(|input| {
310                input.saturating_sub(cached_tokens.unwrap_or(0) + cache_write_tokens.unwrap_or(0))
311            }),
312        output_tokens: raw.get("completion_tokens").and_then(Value::as_u64),
313        cached_tokens,
314        cache_write_tokens,
315        total_tokens: raw.get("total_tokens").and_then(Value::as_u64),
316    }
317}
a_agent/provider/chat_completion.rs

a_agent/provider/
chat_completion.rs
