ic_llm/
lib.rs

1#![doc = include_str!("../README.md")]
2use std::fmt;
3
4// Define our modules
5mod chat;
6mod tool;
7
8// Re-export public types from modules
9pub use chat::{AssistantMessage, ChatBuilder, ChatMessage, FunctionCall, Response, ToolCall};
10pub use tool::{
11    Function, ParameterBuilder, ParameterType, Parameters, Property, Tool, ToolBuilder,
12};
13
14// The principal of the LLM canister.
15const LLM_CANISTER: &str = "w36hm-eqaaa-aaaal-qr76a-cai";
16
17/// Supported LLM models.
18#[derive(Debug)]
19pub enum Model {
20    Llama3_1_8B,
21    Qwen3_32B,
22    Llama4Scout,
23}
24
25impl fmt::Display for Model {
26    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
27        let text = match self {
28            Model::Llama3_1_8B => "llama3.1:8b",
29            Model::Qwen3_32B => "qwen3:32b",
30            Model::Llama4Scout => "llama4-scout",
31        };
32        write!(f, "{}", text)
33    }
34}
35
36/// Sends a single message to a model.
37///
38/// # Example
39///
40/// ```
41/// use ic_llm::Model;
42///
43/// # async fn prompt_example() -> String {
44/// ic_llm::prompt(Model::Llama3_1_8B, "What's the speed of light?").await
45/// # }
46/// ```
47pub async fn prompt<P: ToString>(model: Model, prompt_str: P) -> String {
48    let response = ChatBuilder::new(model)
49        .with_messages(vec![ChatMessage::User {
50            content: prompt_str.to_string(),
51        }])
52        .send()
53        .await;
54
55    response.message.content.unwrap_or_default()
56}
57
58/// Creates a new ChatBuilder with the specified model.
59///
60/// This is a convenience function that returns a ChatBuilder instance initialized with the given model.
61/// You can then chain additional methods to configure the chat request before sending it.
62///
63/// # Example
64///
65/// ```
66/// use ic_llm::{Model, ChatMessage, Response};
67///
68/// # async fn chat_example() -> Response {
69/// // Basic usage
70/// ic_llm::chat(Model::Llama3_1_8B)
71///     .with_messages(vec![
72///         ChatMessage::System {
73///             content: "You are a helpful assistant".to_string(),
74///         },
75///         ChatMessage::User {
76///             content: "How big is the sun?".to_string(),
77///         },
78///     ])
79///     .send()
80///     .await
81/// # }
82/// ```
83///
84/// You can also add tools to the chat:
85///
86/// ```
87/// use ic_llm::{Model, ChatMessage, ParameterType, Response};
88///
89/// # async fn chat_with_tools_example() -> Response {
90/// ic_llm::chat(Model::Llama3_1_8B)
91///     .with_messages(vec![
92///         ChatMessage::System {
93///             content: "You are a helpful assistant".to_string(),
94///         },
95///         ChatMessage::User {
96///             content: "What's the balance of account abc123?".to_string(),
97///         },
98///     ])
99///     .with_tools(vec![
100///         ic_llm::tool("icp_account_balance")
101///             .with_description("Lookup the balance of an ICP account")
102///             .with_parameter(
103///                 ic_llm::parameter("account", ParameterType::String)
104///                     .with_description("The ICP account to look up")
105///                     .is_required()
106///             )
107///             .build()
108///     ])
109///     .send()
110///     .await
111/// # }
112/// ```
113pub fn chat(model: Model) -> ChatBuilder {
114    ChatBuilder::new(model)
115}
116
117/// Creates a new ToolBuilder with the specified name.
118///
119/// This is a convenience function that returns a ToolBuilder instance initialized with the given name.
120/// You can then chain additional methods to configure the tool before building it.
121///
122/// # Example
123///
124/// ```
125/// use ic_llm::{ParameterType, Response};
126///
127/// # fn tool_example() {
128/// // Basic usage
129/// let weather_tool = ic_llm::tool("get_weather")
130///     .with_description("Get current weather for a location")
131///     .with_parameter(
132///         ic_llm::parameter("location", ParameterType::String)
133///             .with_description("The location to get weather for")
134///             .is_required()
135///     )
136///     .build();
137/// # }
138/// ```
139pub fn tool<S: Into<String>>(name: S) -> ToolBuilder {
140    ToolBuilder::new(name)
141}
142
143/// Creates a new ParameterBuilder with the specified name and type.
144///
145/// This is a convenience function that returns a ParameterBuilder instance initialized with the given name and type.
146/// You can then chain additional methods to configure the parameter before adding it to a tool.
147///
148/// # Example
149///
150/// ```
151/// use ic_llm::ParameterType;
152///
153/// # fn parameter_example() {
154/// // Basic usage
155/// let location_param = ic_llm::parameter("location", ParameterType::String)
156///     .with_description("The location to get weather for")
157///     .is_required();
158/// # }
159/// ```
160pub fn parameter<S: Into<String>>(name: S, type_: ParameterType) -> ParameterBuilder {
161    ParameterBuilder::new(name, type_)
162}
ic_llm/lib.rs

ic_llm/
lib.rs
