1use async_openai::{
2 Client,
3 config::OpenAIConfig,
4 error::OpenAIError,
5 types::{chat::CreateChatCompletionStreamResponse, stream::StreamResponse},
6};
7use serde::de::DeserializeOwned;
8use serde_json::Value as JsonValue;
9use std::future::Future;
10
11use crate::event_bus::token_bus::TokenEvent;
12
13#[cfg(feature = "deepseek")]
14pub mod deepseek;
15#[cfg(feature = "openai")]
16pub mod openai;
17
18pub trait StreamChunkExt: DeserializeOwned + Send + 'static {
23 fn extract_events(&self) -> Vec<TokenEvent>;
25}
26
27pub trait ChatProvider: Send + Sync + 'static {
34 type Chunk: StreamChunkExt;
36
37 fn build_request_json(
42 model: &str,
43 messages: &[JsonValue],
44 skill_content: &str,
45 tools_json: &JsonValue,
46 ) -> JsonValue;
47
48 fn create_stream(
53 client: &Client<OpenAIConfig>,
54 request_json: JsonValue,
55 ) -> impl Future<Output = Result<StreamResponse<Self::Chunk>, OpenAIError>> + Send;
56}
57
58impl StreamChunkExt for CreateChatCompletionStreamResponse {
61 fn extract_events(&self) -> Vec<TokenEvent> {
62 let mut events = Vec::new();
63 for choice in &self.choices {
64 if let Some(finish_reason) = choice.finish_reason {
65 events.push(TokenEvent::Finish(finish_reason));
66 }
67 match (
68 choice.delta.content.as_deref(),
69 choice.delta.tool_calls.as_ref(),
70 ) {
71 (Some(text), Some(tool_calls)) => {
72 if !text.is_empty() {
73 events.push(TokenEvent::Text(text.to_string()));
74 }
75 for tc in tool_calls {
76 events.push(TokenEvent::ToolDelta {
77 index: tc.index as usize,
78 call_id: tc.id.clone().unwrap_or_default(),
79 name: tc.function.clone().and_then(|f| f.name),
80 args_chunk: tc.function.clone().and_then(|f| f.arguments),
81 });
82 }
83 }
84 (Some(text), None) => {
85 if !text.is_empty() {
86 events.push(TokenEvent::Text(text.to_string()));
87 }
88 }
89 (None, Some(tool_calls)) => {
90 for tc in tool_calls {
91 events.push(TokenEvent::ToolDelta {
92 index: tc.index as usize,
93 call_id: tc.id.clone().unwrap_or_default(),
94 name: tc.function.clone().and_then(|f| f.name),
95 args_chunk: tc.function.clone().and_then(|f| f.arguments),
96 });
97 }
98 }
99 (None, None) => {}
100 }
101 }
102 events
103 }
104}
105
106pub fn build_standard_request_json(
108 model: &str,
109 messages: &[JsonValue],
110 skill_content: &str,
111 tools_json: &JsonValue,
112) -> JsonValue {
113 let mut msgs: Vec<JsonValue> = messages.to_vec();
114 if !skill_content.is_empty() {
115 msgs.push(JsonValue::Object(
116 [
117 ("role".into(), "system".into()),
118 ("content".into(), skill_content.into()),
119 ]
120 .into_iter()
121 .collect(),
122 ));
123 }
124 let mut req = JsonValue::Object(
125 [
126 ("model".into(), model.into()),
127 ("messages".into(), JsonValue::Array(msgs)),
128 ("stream".into(), true.into()),
129 ]
130 .into_iter()
131 .collect(),
132 );
133 if let Some(arr) = tools_json.as_array()
134 && !arr.is_empty()
135 {
136 req.as_object_mut()
137 .unwrap()
138 .insert("tools".into(), tools_json.clone());
139 }
140 req
141}