ic_llm/
lib.rs

1#![doc = include_str!("../README.md")]
2use candid::Principal;
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 mainnet principal of the LLM canister.
15const MAINNET_LLM_CANISTER: &str = "w36hm-eqaaa-aaaal-qr76a-cai";
16
17/// Resolves the LLM canister principal: prefers `PUBLIC_CANISTER_ID:llm` (auto-injected
18/// by `icp deploy`) and otherwise falls back to the mainnet canister.
19pub(crate) fn default_llm_canister() -> Principal {
20    // The env-var lookup only works in a canister.
21    // Skip in unit tests.
22    #[cfg(not(test))]
23    {
24        const LLM_CANISTER_ENV: &str = "PUBLIC_CANISTER_ID:llm";
25        if ic_cdk::api::env_var_name_exists(LLM_CANISTER_ENV) {
26            let id = ic_cdk::api::env_var_value(LLM_CANISTER_ENV);
27            return Principal::from_text(&id)
28                .unwrap_or_else(|e| ic_cdk::trap(format!("invalid {LLM_CANISTER_ENV}: {e}")));
29        }
30    }
31    Principal::from_text(MAINNET_LLM_CANISTER).unwrap()
32}
33
34/// Sends a single message to a model.
35///
36/// `model` is the canister's model identifier, e.g. `"llama3.1:8b"` (free) or
37/// `"gemma3:27b"` (paid). See the README for the current list.
38///
39/// # Example
40///
41/// ```
42/// # async fn prompt_example() -> String {
43/// ic_llm::prompt("llama3.1:8b", "What's the speed of light?").await
44/// # }
45/// ```
46pub async fn prompt<P: ToString>(model: impl Into<String>, prompt_str: P) -> String {
47    let response = ChatBuilder::new(model)
48        .with_messages(vec![ChatMessage::User {
49            content: prompt_str.to_string(),
50        }])
51        .send()
52        .await;
53
54    response.message.content.unwrap_or_default()
55}
56
57/// Creates a new ChatBuilder with the specified model.
58///
59/// This is a convenience function that returns a ChatBuilder instance initialized with the given model.
60/// You can then chain additional methods to configure the chat request before sending it.
61///
62/// # Example
63///
64/// ```
65/// use ic_llm::{ChatMessage, Response};
66///
67/// # async fn chat_example() -> Response {
68/// // Basic usage
69/// ic_llm::chat("llama3.1:8b")
70///     .with_messages(vec![
71///         ChatMessage::System {
72///             content: "You are a helpful assistant".to_string(),
73///         },
74///         ChatMessage::User {
75///             content: "How big is the sun?".to_string(),
76///         },
77///     ])
78///     .send()
79///     .await
80/// # }
81/// ```
82///
83/// You can also add tools to the chat:
84///
85/// ```
86/// use ic_llm::{ChatMessage, ParameterType, Response};
87///
88/// # async fn chat_with_tools_example() -> Response {
89/// ic_llm::chat("llama3.1:8b")
90///     .with_messages(vec![
91///         ChatMessage::System {
92///             content: "You are a helpful assistant".to_string(),
93///         },
94///         ChatMessage::User {
95///             content: "What's the balance of account abc123?".to_string(),
96///         },
97///     ])
98///     .with_tools(vec![
99///         ic_llm::tool("icp_account_balance")
100///             .with_description("Lookup the balance of an ICP account")
101///             .with_parameter(
102///                 ic_llm::parameter("account", ParameterType::String)
103///                     .with_description("The ICP account to look up")
104///                     .is_required()
105///             )
106///             .build()
107///     ])
108///     .send()
109///     .await
110/// # }
111/// ```
112pub fn chat(model: impl Into<String>) -> ChatBuilder {
113    ChatBuilder::new(model)
114}
115
116/// Creates a new ToolBuilder with the specified name.
117///
118/// This is a convenience function that returns a ToolBuilder instance initialized with the given name.
119/// You can then chain additional methods to configure the tool before building it.
120///
121/// # Example
122///
123/// ```
124/// use ic_llm::{ParameterType, Response};
125///
126/// # fn tool_example() {
127/// // Basic usage
128/// let weather_tool = ic_llm::tool("get_weather")
129///     .with_description("Get current weather for a location")
130///     .with_parameter(
131///         ic_llm::parameter("location", ParameterType::String)
132///             .with_description("The location to get weather for")
133///             .is_required()
134///     )
135///     .build();
136/// # }
137/// ```
138pub fn tool<S: Into<String>>(name: S) -> ToolBuilder {
139    ToolBuilder::new(name)
140}
141
142/// Creates a new ParameterBuilder with the specified name and type.
143///
144/// This is a convenience function that returns a ParameterBuilder instance initialized with the given name and type.
145/// You can then chain additional methods to configure the parameter before adding it to a tool.
146///
147/// # Example
148///
149/// ```
150/// use ic_llm::ParameterType;
151///
152/// # fn parameter_example() {
153/// // Basic usage
154/// let location_param = ic_llm::parameter("location", ParameterType::String)
155///     .with_description("The location to get weather for")
156///     .is_required();
157/// # }
158/// ```
159pub fn parameter<S: Into<String>>(name: S, type_: ParameterType) -> ParameterBuilder {
160    ParameterBuilder::new(name, type_)
161}
ic_llm/lib.rs

ic_llm/
lib.rs
