nihility-model 0.2.2

nihility project ai model module
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

use crate::config::ModelType;
use crate::error::Result;
use crate::{ModelClient, NihilityModel, NihilityModelError};
use std::pin::Pin;
use async_openai::types::completions::CreateCompletionRequestArgs;
use tokio::spawn;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokio_stream::{Stream, StreamExt};
use tracing::error;

pub type TextResponseStream =
    Pin<Box<dyn Stream<Item = core::result::Result<String, NihilityModelError>> + Send>>;

impl NihilityModel {
    pub async fn completion(&self, model_type: ModelType, prompt: &str) -> Result<String> {
        let (client, model) = self
            .config
            .model_client(self.http_client.clone(), model_type)?;
        let request = CreateCompletionRequestArgs::default()
            .model(model)
            .prompt(prompt)
            .stream(false)
            .build()?;
        Ok(match client {
            ModelClient::OpenAI(openai_client) => openai_client
                .completions()
                .create(request)
                .await?
                .choices
                .first()
                .ok_or(NihilityModelError::Response)?
                .text
                .clone(),
            ModelClient::Azure(azure_client) => azure_client
                .completions()
                .create(request)
                .await?
                .choices
                .first()
                .ok_or(NihilityModelError::Response)?
                .text
                .clone(),
        })
    }

    pub async fn completion_stream(
        &self,
        model_type: ModelType,
        prompt: &str,
    ) -> Result<TextResponseStream> {
        let (client, model) = self
            .config
            .model_client(self.http_client.clone(), model_type)?;
        let request = CreateCompletionRequestArgs::default()
            .model(model)
            .prompt(prompt)
            .stream(true)
            .build()?;
        let (tx, rx) = mpsc::channel(10);
        let mut resp_stream = match client {
            ModelClient::OpenAI(openai_client) => {
                openai_client.completions().create_stream(request).await?
            }
            ModelClient::Azure(azure_client) => {
                azure_client.completions().create_stream(request).await?
            }
        };
        spawn(async move {
            while let Some(resp) = resp_stream.next().await {
                match resp {
                    Ok(ccr) => {
                        let text = match ccr.choices.first() {
                            None => {
                                error!("Received empty choices");
                                if let Err(e) = tx.send(Err(NihilityModelError::Response)).await {
                                    error!("Send result to response error: {}", e);
                                    break;
                                }
                                continue;
                            }
                            Some(choice) => choice.text.clone(),
                        };
                        if let Err(e) = tx.send(Ok(text)).await {
                            error!("Send result to response error: {}", e);
                            break;
                        }
                    }
                    Err(e) => {
                        error!("Response Stream resp Error: {}", e);
                        if let Err(e) = tx.send(Err(NihilityModelError::from(e))).await {
                            error!("Send result to response error: {}", e);
                            break;
                        }
                        break;
                    }
                }
            }
        });
        Ok(Box::pin(ReceiverStream::new(rx)) as TextResponseStream)
    }
}