1mod error;
5mod structured;
6#[cfg(test)]
7mod tests;
8
9pub use error::OpenAIError;
10pub use structured::StructuredOutputMethod;
11
12use async_trait::async_trait;
13use futures_util::Stream;
14use serde::Deserialize;
15use serde_json::json;
16use std::marker::PhantomData;
17use std::pin::Pin;
18
19use super::OpenAIConfig;
20use lc_callbacks::RunType;
21use lc_core::language_models::{
22 BaseChatModel, BaseLanguageModel, LLMResult, StreamChunk, TokenUsage,
23};
24use lc_core::runnables::{run_tree_from_config, Runnable};
25use lc_core::tools::ToolDefinition;
26use lc_core::RunnableConfig;
27use lc_schema::Message;
28use schemars::JsonSchema;
29use serde::de::DeserializeOwned;
30
31#[derive(Clone)]
33pub struct OpenAIChat {
34 pub(crate) config: OpenAIConfig,
35 pub(crate) client: reqwest::Client,
36}
37
38impl OpenAIChat {
39 pub(crate) fn apply_headers(
42 mut builder: reqwest::RequestBuilder,
43 config: &OpenAIConfig,
44 ) -> reqwest::RequestBuilder {
45 builder = builder.header("Content-Type", "application/json");
46 if config.send_auth {
47 builder = builder.header("Authorization", format!("Bearer {}", config.api_key));
48 }
49 for (name, value) in &config.extra_headers {
50 builder = builder.header(name, value);
51 }
52 builder
53 }
54
55 pub fn new(config: OpenAIConfig) -> Self {
57 Self {
58 config,
59 client: crate::retry::default_client(),
62 }
63 }
64
65 pub fn from_env_result() -> Result<Self, OpenAIError> {
67 let config = OpenAIConfig::from_env_result()?;
68 Ok(Self::new(config))
69 }
70
71 fn message_to_openai_format(message: &Message) -> serde_json::Value {
73 match &message.message_type {
74 lc_schema::MessageType::System => json!({
75 "role": "system",
76 "content": message.content,
77 }),
78 lc_schema::MessageType::Human => {
79 if let Some(blocks) = crate::media::openai_user_blocks(message) {
83 json!({"role": "user", "content": blocks})
84 } else {
85 json!({"role": "user", "content": &message.content})
86 }
87 }
88 lc_schema::MessageType::AI => {
89 let mut msg = json!({
90 "role": "assistant",
91 "content": message.content,
92 });
93 if let Some(tool_calls) = &message.tool_calls {
94 msg["tool_calls"] =
95 serde_json::to_value(tool_calls).unwrap_or(serde_json::Value::Null);
96 }
97 msg
98 }
99 lc_schema::MessageType::Tool { tool_call_id } => json!({
100 "role": "tool",
101 "tool_call_id": tool_call_id,
102 "content": message.content,
103 }),
104 }
105 }
106
107 fn build_request_body(&self, messages: Vec<Message>, stream: bool) -> serde_json::Value {
109 let openai_messages: Vec<serde_json::Value> = messages
110 .iter()
111 .map(Self::message_to_openai_format)
112 .collect();
113
114 let mut body = json!({
115 "model": self.config.model,
116 "messages": openai_messages,
117 "stream": stream,
118 });
119
120 if let Some(temp) = self.config.temperature {
121 body["temperature"] = json!(temp);
122 }
123
124 if let Some(max) = self.config.max_tokens {
125 body["max_tokens"] = json!(max);
126 }
127
128 if let Some(top_p) = self.config.top_p {
129 body["top_p"] = json!(top_p);
130 }
131
132 if let Some(tools) = &self.config.tools {
133 body["tools"] = serde_json::to_value(tools).unwrap_or(serde_json::Value::Null);
134 }
135
136 if let Some(tool_choice) = &self.config.tool_choice {
137 body["tool_choice"] = json!(tool_choice);
138 }
139
140 if let Some(format) = &self.config.response_format {
145 body["response_format"] =
146 serde_json::to_value(format).unwrap_or(serde_json::Value::Null);
147 }
148
149 body
150 }
151
152 pub fn bind_tools(&self, tools: Vec<ToolDefinition>) -> Self {
154 let config = OpenAIConfig {
155 tools: Some(tools),
156 ..self.config.clone()
157 };
158 Self {
159 config,
160 client: self.client.clone(),
161 }
162 }
163
164 pub fn with_tool_choice(mut self, choice: impl Into<String>) -> Self {
166 self.config.tool_choice = Some(choice.into());
167 self
168 }
169
170 pub fn with_structured_output<T: DeserializeOwned + JsonSchema>(
172 &self,
173 ) -> StructuredOutputMethod<T> {
174 use schemars::schema_for;
175 let schema = serde_json::to_value(schema_for!(T)).unwrap_or_else(|_| {
176 serde_json::json!({"type": "object", "properties": {}})
178 });
179
180 let tool = ToolDefinition::new("structured_output", "Return structured JSON output")
181 .with_parameters(schema)
182 .with_strict(true);
183
184 let config = OpenAIConfig {
185 tools: Some(vec![tool]),
186 tool_choice: Some("auto".to_string()),
187 ..self.config.clone()
188 };
189
190 StructuredOutputMethod {
191 config,
192 client: self.client.clone(),
193 _phantom: PhantomData,
194 }
195 }
196
197 pub fn with_json_schema_output<T: DeserializeOwned + JsonSchema>(
211 &self,
212 ) -> StructuredOutputMethod<T> {
213 use schemars::schema_for;
214 let mut schema = serde_json::to_value(schema_for!(T)).unwrap_or_else(|_| {
215 serde_json::json!({"type": "object", "properties": {}})
217 });
218 crate::openai::response_format::make_strict_schema(&mut schema);
219
220 let config = OpenAIConfig {
221 response_format: Some(crate::openai::response_format::ResponseFormat::json_schema(
222 "output", schema,
223 )),
224 ..self.config.clone()
225 };
226
227 StructuredOutputMethod {
228 config,
229 client: self.client.clone(),
230 _phantom: PhantomData,
231 }
232 }
233}
234
235#[async_trait]
236impl Runnable<Vec<Message>, LLMResult> for OpenAIChat {
237 type Error = OpenAIError;
238
239 async fn invoke(
240 &self,
241 input: Vec<Message>,
242 _config: Option<RunnableConfig>,
243 ) -> Result<LLMResult, Self::Error> {
244 self.chat(input, _config).await
245 }
246
247 async fn stream(
248 &self,
249 input: Vec<Message>,
250 config: Option<RunnableConfig>,
251 ) -> Result<Pin<Box<dyn Stream<Item = Result<LLMResult, Self::Error>> + Send>>, Self::Error>
252 {
253 use futures_util::StreamExt;
254
255 let model = self.config.model.clone();
256 let (temp, max) = crate::sampling::sampling_overrides(&config);
257 let mut effective = self.clone();
258 if let Some(t) = temp {
259 effective.config.temperature = Some(t);
260 }
261 if let Some(m) = max {
262 effective.config.max_tokens = Some(m);
263 }
264 let token_stream = effective.stream_chat_internal(input).await?;
265
266 let stream = token_stream.map(move |token_result| match token_result {
269 Ok(chunk) => Ok(LLMResult {
270 content: chunk.text,
271 model: model.clone(),
272 token_usage: chunk.token_usage,
273 tool_calls: None,
274 thinking_content: None,
275 }),
276 Err(e) => Err(e),
277 });
278
279 Ok(Box::pin(stream))
280 }
281}
282
283#[async_trait]
284impl BaseLanguageModel<Vec<Message>, LLMResult> for OpenAIChat {
285 fn model_name(&self) -> &str {
286 &self.config.model
287 }
288
289 fn get_num_tokens(&self, text: &str) -> usize {
290 lc_core::token_counter::count_tokens(text).unwrap_or_else(|e| {
291 log::warn!("Token counting failed, falling back to byte-length estimation: {e}");
293 text.len()
294 })
295 }
296
297 fn temperature(&self) -> Option<f32> {
298 self.config.temperature
299 }
300
301 fn max_tokens(&self) -> Option<usize> {
302 self.config.max_tokens
303 }
304
305 fn with_temperature(mut self, temp: f32) -> Self {
306 self.config.temperature = Some(temp);
307 self
308 }
309
310 fn with_max_tokens(mut self, max: usize) -> Self {
311 self.config.max_tokens = Some(max);
312 self
313 }
314}
315
316#[async_trait]
317impl BaseChatModel for OpenAIChat {
318 async fn chat(
319 &self,
320 messages: Vec<Message>,
321 config: Option<RunnableConfig>,
322 ) -> Result<LLMResult, Self::Error> {
323 let run_name = config
324 .as_ref()
325 .and_then(|c| c.run_name.clone())
326 .unwrap_or_else(|| format!("{}:chat", self.config.model));
327
328 let mut run = run_tree_from_config(
329 run_name,
330 RunType::Llm,
331 json!({
332 "messages": messages.iter().map(|m| m.content.clone()).collect::<Vec<_>>(),
333 "model": self.config.model,
334 }),
335 config.as_ref(),
336 );
337
338 if let Some(ref cfg) = config {
339 if let Some(ref callbacks) = cfg.callbacks {
340 for handler in callbacks.handlers() {
341 handler.on_llm_start(&run, &messages).await;
342 }
343 }
344 }
345
346 let (temp, max) = crate::sampling::sampling_overrides(&config);
347 let mut effective = self.clone();
348 if let Some(t) = temp {
349 effective.config.temperature = Some(t);
350 }
351 if let Some(m) = max {
352 effective.config.max_tokens = Some(m);
353 }
354
355 let result = if effective.config.streaming {
358 let stream = effective.stream_chat_internal(messages.clone()).await?;
359 let (content, token_usage, tool_calls) = Self::aggregate_stream(stream).await?;
364 Ok(LLMResult {
365 content,
366 model: effective.config.model.clone(),
367 token_usage,
368 tool_calls,
369 thinking_content: None,
370 })
371 } else {
372 effective.chat_internal(messages.clone()).await
373 };
374
375 match result {
376 Ok(response) => {
377 run.end(json!({
378 "content": &response.content,
379 "model": &response.model,
380 "token_usage": &response.token_usage,
381 }));
382
383 if let Some(ref cfg) = config {
384 if let Some(ref callbacks) = cfg.callbacks {
385 for handler in callbacks.handlers() {
386 handler.on_llm_end(&run, &response.content).await;
387 }
388 }
389 }
390
391 Ok(response)
392 }
393 Err(e) => {
394 run.end_with_error(e.to_string());
395
396 if let Some(ref cfg) = config {
397 if let Some(ref callbacks) = cfg.callbacks {
398 for handler in callbacks.handlers() {
399 handler.on_llm_error(&run, &e.to_string()).await;
400 }
401 }
402 }
403
404 Err(e)
405 }
406 }
407 }
408
409 async fn stream_chat(
410 &self,
411 messages: Vec<Message>,
412 config: Option<RunnableConfig>,
413 ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
414 {
415 use futures_util::StreamExt;
416
417 let run_name = config
418 .as_ref()
419 .and_then(|c| c.run_name.clone())
420 .unwrap_or_else(|| format!("{}:stream", self.config.model));
421
422 let run = run_tree_from_config(
423 run_name,
424 RunType::Llm,
425 json!({
426 "messages": messages.len(),
427 "model": self.config.model,
428 }),
429 config.as_ref(),
430 );
431
432 if let Some(ref cfg) = config {
433 if let Some(ref callbacks) = cfg.callbacks {
434 for handler in callbacks.handlers() {
435 handler.on_llm_start(&run, &messages).await;
436 }
437 }
438 }
439
440 let (temp, max) = crate::sampling::sampling_overrides(&config);
441 let mut effective = self.clone();
442 if let Some(t) = temp {
443 effective.config.temperature = Some(t);
444 }
445 if let Some(m) = max {
446 effective.config.max_tokens = Some(m);
447 }
448 let stream = effective.stream_chat_internal(messages).await?;
449
450 let callbacks = config.and_then(|c| c.callbacks);
451 let stream = stream.then(move |token_result| {
452 let cbs = callbacks.clone();
453 let run = run.clone();
454 async move {
455 if let Some(ref cbs) = cbs {
456 if let Ok(ref token) = token_result {
457 for handler in cbs.handlers() {
458 handler.on_llm_new_token(&run, &token.text).await;
459 }
460 }
461 }
462 token_result
463 }
464 });
465
466 Ok(Box::pin(stream))
467 }
468
469 fn bind_tools(
470 &self,
471 tools: Vec<ToolDefinition>,
472 ) -> Option<Box<dyn BaseChatModel<Error = Self::Error> + Send + Sync>> {
473 Some(Box::new(self.bind_tools(tools)))
476 }
477}
478
479impl OpenAIChat {
480 pub(crate) async fn chat_internal(
481 &self,
482 messages: Vec<Message>,
483 ) -> Result<LLMResult, OpenAIError> {
484 let url = format!("{}/chat/completions", self.config.base_url);
485 let mut messages = messages;
486 crate::media::resolve_message_media(&mut messages, crate::media::MediaPolicy::OpenAi)
487 .await
488 .map_err(|e| OpenAIError::Api(e.to_string()))?;
489 let body = self.build_request_body(messages, false);
490
491 let response = crate::retry::send_with_retry(
498 || Self::apply_headers(self.client.post(&url), &self.config).json(&body),
499 &crate::retry::DEFAULT_RETRY,
500 )
501 .await
502 .map_err(|e| OpenAIError::Http(e.to_string()))?;
503
504 let status = response.status();
505 if !status.is_success() {
506 let error_text = response.text().await.unwrap_or_default();
507 return Err(OpenAIError::Api(format!("HTTP {}: {}", status, error_text)));
508 }
509
510 let chat_response: OpenAIChatResponse = response
511 .json()
512 .await
513 .map_err(|e| OpenAIError::Parse(e.to_string()))?;
514
515 let choice = chat_response
516 .choices
517 .first()
518 .ok_or_else(|| OpenAIError::Api("No choices in response".to_string()))?;
519 let message = &choice.message;
520
521 Ok(Self::llm_result_from_message(
522 message,
523 chat_response.model,
524 chat_response.usage,
525 ))
526 }
527
528 fn llm_result_from_message(
535 message: &OpenAIMessage,
536 model: String,
537 usage: Option<OpenAIUsage>,
538 ) -> LLMResult {
539 let content = message
540 .content
541 .clone()
542 .filter(|c| !c.is_empty())
543 .unwrap_or_default();
544
545 let thinking_content = message.reasoning_content.clone().filter(|c| !c.is_empty());
546
547 if let Some(r) = usage.as_ref().and_then(OpenAIUsage::reasoning_tokens) {
551 log::info!(target: "lc_providers::openai::reasoning", "reasoning_tokens={r}");
552 }
553 if let Some(ref r) = message.refusal {
554 log::warn!(
555 target: "lc_providers::openai::refusal",
556 "model refusal surfaced (content={}): {r}", content
557 );
558 }
559
560 LLMResult {
561 content,
562 model,
563 token_usage: usage.map(|u| TokenUsage {
564 prompt_tokens: u.prompt_tokens,
565 completion_tokens: u.completion_tokens,
566 total_tokens: u.total_tokens,
567 }),
568 tool_calls: message.tool_calls.clone(),
569 thinking_content,
570 }
571 }
572
573 async fn stream_chat_internal(
574 &self,
575 messages: Vec<Message>,
576 ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, OpenAIError>> + Send>>, OpenAIError>
577 {
578 use super::sse::{SSEParser, SseByteFramer, StreamToolCallAccumulator};
579 use std::sync::{Arc, Mutex};
580
581 let url = format!("{}/chat/completions", self.config.base_url);
582 let mut messages = messages;
583 crate::media::resolve_message_media(&mut messages, crate::media::MediaPolicy::OpenAi)
584 .await
585 .map_err(|e| OpenAIError::Api(e.to_string()))?;
586 let body = self.build_request_body(messages, true);
587
588 let response = Self::apply_headers(self.client.post(&url), &self.config)
589 .json(&body)
590 .send()
591 .await
592 .map_err(|e| OpenAIError::Http(e.to_string()))?;
593
594 let status = response.status();
595 if !status.is_success() {
596 let error_text = response.text().await.unwrap_or_default();
597 return Err(OpenAIError::Api(format!("HTTP {}: {}", status, error_text)));
598 }
599
600 let byte_stream = response.bytes_stream();
601
602 let parser = Arc::new(Mutex::new((SSEParser::new(), SseByteFramer::new())));
603
604 let parser_clone = parser.clone();
605 let (tx, rx) = tokio::sync::mpsc::channel::<Result<StreamChunk, OpenAIError>>(64);
607
608 tokio::spawn(async move {
609 use futures_util::StreamExt;
610 let mut byte_stream = byte_stream;
611 let mut tool_acc = StreamToolCallAccumulator::default();
615 let mut tool_calls_emitted = false;
616 let mut done = false;
619 let mut saw_terminal = false;
627 while let Some(chunk_result) = byte_stream.next().await {
628 let chunk_bytes = match chunk_result {
631 Ok(bytes) => bytes,
632 Err(e) => {
633 let _ = tx.send(Err(OpenAIError::Http(e.to_string()))).await;
634 return;
635 }
636 };
637
638 let events = {
639 let mut guard = parser_clone.lock().unwrap_or_else(|e| e.into_inner());
643 let mut out = Vec::new();
644 for text in guard.1.push(&chunk_bytes) {
645 out.extend(guard.0.parse(&text));
646 }
647 out
648 };
649 for event in events {
652 if event.is_done() {
653 done = true;
654 saw_terminal = true;
655 break;
656 }
657 match event.parse_openai_chunk() {
660 Ok(Some(chunk)) => {
661 if let Some(choice) = chunk.choices.first() {
662 if let Some(content) = &choice.delta.content {
663 if tx.send(Ok(StreamChunk::new(content))).await.is_err() {
664 return;
665 }
666 }
667 if let Some(deltas) = &choice.delta.tool_calls {
668 for delta in deltas {
669 tool_acc.push(delta);
670 }
671 }
672 }
673 if chunk.choices.iter().any(|c| c.finish_reason.is_some()) {
678 saw_terminal = true;
679 }
680 if let Some(usage) = chunk.usage {
687 let token_usage = TokenUsage {
688 prompt_tokens: usage.prompt_tokens,
689 completion_tokens: usage.completion_tokens,
690 total_tokens: usage.total_tokens,
691 };
692 if let Some(r) = usage.reasoning_tokens() {
695 log::info!(
696 target: "lc_providers::openai::reasoning",
697 "reasoning_tokens={r}"
698 );
699 }
700 let tool_calls = tool_acc.build();
701 if !tool_calls.is_empty() {
702 tool_calls_emitted = true;
703 }
704 let final_chunk = StreamChunk {
705 text: String::new(),
706 token_usage: Some(token_usage),
707 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
708 };
709 if tx.send(Ok(final_chunk)).await.is_err() {
710 return;
711 }
712 }
713 }
714 Ok(None) => {}
715 Err(e) => {
716 log::error!(
717 "Failed to parse streaming SSE chunk (skipping this token): {}",
718 e
719 );
720 }
721 }
722 }
723 if done {
724 break;
725 }
726 }
727 if !saw_terminal {
732 let _ = tx
733 .send(Err(OpenAIError::StreamInterrupted(
734 "connection closed before [DONE] or finish_reason".to_string(),
735 )))
736 .await;
737 return;
738 }
739 if !tool_calls_emitted {
743 let tool_calls = tool_acc.build();
744 if !tool_calls.is_empty() {
745 let _ = tx
746 .send(Ok(StreamChunk {
747 text: String::new(),
748 token_usage: None,
749 tool_calls: Some(tool_calls),
750 }))
751 .await;
752 }
753 }
754 });
755 let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
756
757 Ok(Box::pin(stream))
758 }
759
760 async fn aggregate_stream(
768 mut stream: Pin<Box<dyn Stream<Item = Result<StreamChunk, OpenAIError>> + Send>>,
769 ) -> Result<
770 (
771 String,
772 Option<TokenUsage>,
773 Option<Vec<lc_core::tools::ToolCall>>,
774 ),
775 OpenAIError,
776 > {
777 use futures_util::StreamExt;
778 let mut content = String::new();
779 let mut token_usage = None;
780 let mut tool_calls = None;
781 while let Some(item) = stream.next().await {
782 let chunk = item?;
783 content.push_str(&chunk.text);
784 if chunk.token_usage.is_some() {
787 token_usage = chunk.token_usage;
788 }
789 if chunk.tool_calls.is_some() {
790 tool_calls = chunk.tool_calls;
791 }
792 }
793 Ok((content, token_usage, tool_calls))
794 }
795}
796
797#[derive(Debug, Deserialize)]
799#[allow(dead_code)]
800struct OpenAIChatResponse {
801 id: String,
802 object: String,
803 created: i64,
804 model: String,
805 choices: Vec<OpenAIChoice>,
806 usage: Option<OpenAIUsage>,
807}
808
809#[derive(Debug, Deserialize)]
810#[allow(dead_code)]
811struct OpenAIChoice {
812 index: i32,
813 message: OpenAIMessage,
814 finish_reason: Option<String>,
815}
816
817#[derive(Debug, Deserialize)]
818#[allow(dead_code)]
819struct OpenAIMessage {
820 role: String,
821 content: Option<String>,
822 reasoning_content: Option<String>,
824 #[serde(default)]
828 refusal: Option<String>,
829 tool_calls: Option<Vec<lc_core::tools::ToolCall>>,
830}
831
832#[derive(Debug, Deserialize)]
833#[allow(dead_code)]
834struct OpenAIUsage {
835 prompt_tokens: usize,
836 completion_tokens: usize,
837 total_tokens: usize,
838 #[serde(default)]
843 reasoning_tokens: Option<usize>,
844 #[serde(default)]
845 completion_tokens_details: Option<CompletionTokensDetails>,
846}
847
848#[derive(Debug, Deserialize)]
849#[allow(dead_code)]
850struct CompletionTokensDetails {
851 #[serde(default)]
853 reasoning_tokens: Option<usize>,
854}
855
856impl OpenAIUsage {
857 fn reasoning_tokens(&self) -> Option<usize> {
859 self.reasoning_tokens.or_else(|| {
860 self.completion_tokens_details
861 .as_ref()
862 .and_then(|d| d.reasoning_tokens)
863 })
864 }
865}