headless_lms_chatbot/chatbot_tools/
mod.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4use sqlx::PgConnection;
5
6use crate::{
7    azure_chatbot::ChatbotUserContext,
8    chatbot_error::chatbot_err,
9    chatbot_tools::{
10        custom_tools::{
11            course_progress::CourseProgressTool, course_structure::CourseStructureTool,
12            document_lookup::DocumentLookupTool,
13        },
14        provider_tools::azure_ai_search::AzureAISearchToolDefinition,
15    },
16    prelude::{BackendError, ChatbotError, ChatbotErrorType, ChatbotResult},
17};
18
19pub mod custom_tools;
20pub mod provider_tools;
21
22pub trait ChatbotTool {
23    type State;
24    type Arguments: Serialize;
25
26    /// Parse the LLM-generated function arguments and clean them
27    fn parse_arguments(args_string: String) -> ChatbotResult<Self::Arguments>;
28
29    /// Create a new instance after parsing arguments
30    fn from_db_and_arguments(
31        conn: &mut PgConnection,
32        arguments: Self::Arguments,
33        user_context: &ChatbotUserContext,
34    ) -> impl std::future::Future<Output = ChatbotResult<Self>> + Send
35    where
36        Self: Sized;
37
38    /// Output the result of the tool call in LLM-readable form
39    fn output(&self) -> String;
40
41    /// Additional instructions for the LLM on how to describe and
42    /// communicate the tool output. Just-in-time prompt.
43    fn output_description_instructions(&self) -> Option<String>;
44
45    /// Get and format tool output and instructions for LLM
46    fn get_tool_output(&self) -> String {
47        let output = self.output();
48        let instructions = self.output_description_instructions();
49
50        if let Some(i) = instructions {
51            format!(
52                "Result: [output]{output}[/output]\n\nInstructions for describing the output: [instructions]{i}[/instructions]"
53            )
54        } else {
55            output
56        }
57    }
58
59    /// Get parsed arguments
60    fn get_arguments(&self) -> &Self::Arguments;
61
62    /// Get a AzureLLMToolDefinition struct that represents this tool.
63    /// The definition is sent to the LLM as part of a chat request.
64    fn get_tool_definition() -> AzureLLMFunctionToolDefinition;
65
66    /// Create a new instance from connection, args and context
67    fn new(
68        conn: &mut PgConnection,
69        args_string: String,
70        user_context: &ChatbotUserContext,
71    ) -> impl std::future::Future<Output = ChatbotResult<Self>> + Send
72    where
73        Self: Sized,
74    {
75        async {
76            let parsed = Self::parse_arguments(args_string)?;
77            Self::from_db_and_arguments(conn, parsed, user_context).await
78        }
79    }
80}
81
82pub struct ToolProperties<S, A: Serialize> {
83    state: S,
84    arguments: A,
85}
86
87#[derive(Clone, Debug, Deserialize, Serialize)]
88#[serde(untagged)]
89pub enum AzureLLMToolDefinition {
90    Function(AzureLLMFunctionToolDefinition),
91    Search(AzureAISearchToolDefinition),
92}
93
94/// A tool definition that is formatted for Azure.
95/// Defines a tool (function) that the LLM can call.
96#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
97pub struct AzureLLMFunctionToolDefinition {
98    #[serde(rename = "type")]
99    pub tool_type: LLMToolType,
100    pub name: String,
101    pub description: String,
102    pub parameters: LLMToolParams,
103    /// Ensures that the LLM calls the tool with the correct params. Should be `true`
104    pub strict: bool,
105}
106
107/// Parameters that a chatbot tool accepts in an AzureLLMToolDefinition
108#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
109#[serde(rename_all = "camelCase")]
110pub struct LLMToolParams {
111    #[serde(rename = "type")]
112    pub tool_type: LLMToolParamType,
113    pub properties: HashMap<String, LLMToolParamProperties>,
114    pub required: Vec<String>,
115    /// required to be false
116    pub additional_properties: bool,
117}
118
119#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
120pub struct LLMToolParamProperties {
121    #[serde(rename = "type")]
122    pub param_type: String,
123    pub description: String,
124}
125
126#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
127#[serde(rename_all = "snake_case")]
128pub enum LLMToolParamType {
129    Object,
130}
131
132#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
133#[serde(rename_all = "snake_case")]
134pub enum LLMToolType {
135    Function,
136}
137
138/// Get a vec of AzureLLMToolDefinitions for all available chatbot tools
139pub fn get_chatbot_tool_definitions() -> Vec<AzureLLMToolDefinition> {
140    vec![
141        AzureLLMToolDefinition::Function(CourseProgressTool::get_tool_definition()),
142        AzureLLMToolDefinition::Function(DocumentLookupTool::get_tool_definition()),
143        AzureLLMToolDefinition::Function(CourseStructureTool::get_tool_definition()),
144    ]
145}
146
147pub struct ChatbotToolCallResult {
148    pub arguments: String,
149    pub output: String,
150}
151
152/// Call a chatbot tool with LLM-provided arguments by matching the tool call
153/// made by the LLM. User context and db connection are needed for some tools.
154pub async fn call_chatbot_tool(
155    conn: &mut PgConnection,
156    fn_name: &str,
157    fn_args: String,
158    user_context: &ChatbotUserContext,
159) -> ChatbotResult<ChatbotToolCallResult> {
160    let (arguments, output) = match fn_name {
161        "course_progress" => {
162            let tool = CourseProgressTool::new(conn, "".to_string(), user_context).await?;
163            let args = tool.get_arguments();
164            (serde_json::to_string(args)?, tool.output())
165        }
166        "document_lookup" => {
167            let tool = DocumentLookupTool::new(conn, fn_args, user_context).await?;
168            let args = tool.get_arguments();
169            (serde_json::to_string(args)?, tool.output())
170        }
171        "course_structure" => {
172            let tool = CourseStructureTool::new(conn, fn_args, user_context).await?;
173            let args = tool.get_arguments();
174            (serde_json::to_string(args)?, tool.output())
175        }
176        _ => {
177            return Err(chatbot_err!(
178                InvalidToolName,
179                "Incorrect or unknown function name".to_string()
180            ));
181        }
182    };
183    Ok(ChatbotToolCallResult { arguments, output })
184}
headless_lms_chatbot/chatbot_tools/mod.rs

headless_lms_chatbot/chatbot_tools/
mod.rs
