1use std::sync::Arc;
9
10use af_llm::{CompletionRequest, CompletionResponse, LlmClient, LlmError};
11use async_trait::async_trait;
12use tokio::sync::mpsc::UnboundedSender;
13
14#[async_trait]
16pub trait ChatModel: Send + Sync {
17 async fn complete_streaming(
20 &self,
21 request: &CompletionRequest,
22 delta_tx: UnboundedSender<(String, bool)>,
23 ) -> Result<CompletionResponse, LlmError>;
24}
25
26#[async_trait]
28impl ChatModel for LlmClient {
29 async fn complete_streaming(
30 &self,
31 request: &CompletionRequest,
32 delta_tx: UnboundedSender<(String, bool)>,
33 ) -> Result<CompletionResponse, LlmError> {
34 self.complete_stream_single_attempt(request, |content, has_tools| {
35 let _ = delta_tx.send((content.to_string(), has_tools));
36 })
37 .await
38 }
39}
40
41#[async_trait]
45impl ChatModel for Arc<dyn ChatModel> {
46 async fn complete_streaming(
47 &self,
48 request: &CompletionRequest,
49 delta_tx: UnboundedSender<(String, bool)>,
50 ) -> Result<CompletionResponse, LlmError> {
51 (**self).complete_streaming(request, delta_tx).await
52 }
53}