llmservice_flows/
lib.rs

1//! LLM Service integration for [Flows.network](https://flows.network)
2//!
3//! # Quick Start
4//!
5//! To get started, let's write a tiny flow function.
6//!
7//! ```rust
8//! use llmservice_flows::{
9//!     chat::ChatOptions,
10//!     LLMServiceFlows,
11//! };
12//! use lambda_flows::{request_received, send_response};
13//! use serde_json::Value;
14//! use std::collections::HashMap;
15//!
16//! #[no_mangle]
17//! #[tokio::main(flavor = "current_thread")]
18//! pub async fn run() {
19//!     request_received(handler).await;
20//! }
21//!
22//! async fn handler(_qry: HashMap<String, Value>, body: Vec<u8>) {
23//!     let co = ChatOptions {
24//!       model: Some("gpt-4"),
25//!       token_limit: 8192,
26//!       ..Default::default()
27//!     };
28//!     let mut lf = LLMServiceFlows::new("https://api.openai.com/v1");
29//!     lf.set_api_key("your api key");
30//!
31//!     let r = match lf.chat_completion(
32//!         "any_conversation_id",
33//!         String::from_utf8_lossy(&body).into_owned().as_str(),
34//!         &co,
35//!     )
36//!     .await
37//!     {
38//!         Ok(c) => c.choice,
39//!         Err(e) => e,
40//!     };
41//!
42//!     send_response(
43//!         200,
44//!         vec![(
45//!             String::from("content-type"),
46//!             String::from("text/plain; charset=UTF-8"),
47//!         )],
48//!         r.as_bytes().to_vec(),
49//!     );
50//! }
51//! ```
52//!
53//! When the Lambda request is received, chat
54//! using [LLMServiceFlows::chat_completion] then send the response.
55//!
56
57use lazy_static::lazy_static;
58use std::time::Duration;
59use tokio::time::sleep;
60
61pub mod audio;
62pub mod chat;
63pub mod embeddings;
64
65lazy_static! {
66    static ref LLM_API_PREFIX: String =
67        String::from(std::option_env!("LLM_API_PREFIX").unwrap_or("https://llm.flows.network/api"));
68}
69
70extern "C" {
71    fn get_flows_user(p: *mut u8) -> i32;
72    fn get_flow_id(p: *mut u8) -> i32;
73}
74
75unsafe fn _get_flows_user() -> String {
76    let mut flows_user = Vec::<u8>::with_capacity(100);
77    let c = get_flows_user(flows_user.as_mut_ptr());
78    flows_user.set_len(c as usize);
79    String::from_utf8(flows_user).unwrap()
80}
81
82unsafe fn _get_flow_id() -> String {
83    let mut flow_id = Vec::<u8>::with_capacity(100);
84    let c = get_flow_id(flow_id.as_mut_ptr());
85    if c == 0 {
86        panic!("Failed to get flow id");
87    }
88    flow_id.set_len(c as usize);
89    String::from_utf8(flow_id).unwrap()
90}
91
92const MAX_RETRY_TIMES: u8 = 10;
93const RETRY_INTERVAL: u64 = 10; // Wait 10 seconds before retry
94
95/// The main struct for setting the basic configuration
96/// for LLM Service interface.
97pub struct LLMServiceFlows<'a> {
98    /// Exposed url of the LLM Service
99    service_endpoint: &'a str,
100
101    /// API Key
102    api_key: Option<&'a str>,
103
104    /// Use retry_times to set the number of retries when requesting
105    /// LLM Service's api encounters a problem. Default is 0 and max number is 10.
106    retry_times: u8,
107}
108
109pub(crate) trait LLMApi {
110    type Output;
111    async fn api(&self, endpoint: &str, api_key: &str) -> Retry<Self::Output>;
112}
113
114impl<'a> LLMServiceFlows<'a> {
115    pub fn new(service_endpoint: &'a str) -> LLMServiceFlows<'a> {
116        LLMServiceFlows {
117            service_endpoint,
118            api_key: None,
119            retry_times: 0,
120        }
121    }
122
123    pub fn set_retry_times(&mut self, retry_times: u8) {
124        self.retry_times = retry_times;
125    }
126
127    pub fn set_api_key(&mut self, api_key: &'a str) {
128        self.api_key = Some(api_key);
129    }
130
131    async fn keep_trying<F: LLMApi>(&self, llmapi: F) -> Result<<F as LLMApi>::Output, String> {
132        let mut retry_times = match self.retry_times {
133            r if r > MAX_RETRY_TIMES => MAX_RETRY_TIMES,
134            r => r,
135        };
136
137        loop {
138            match llmapi
139                .api(self.service_endpoint, self.api_key.unwrap_or_default())
140                .await
141            {
142                Retry::Yes(s) => match retry_times > 0 {
143                    true => {
144                        sleep(Duration::from_secs(crate::RETRY_INTERVAL)).await;
145                        retry_times = retry_times - 1;
146                        continue;
147                    }
148                    false => return Err(s),
149                },
150                Retry::No(r) => return r,
151            }
152        }
153    }
154}
155
156enum Retry<T> {
157    Yes(String),
158    No(Result<T, String>),
159}
llmservice_flows/lib.rs

llmservice_flows/
lib.rs
