assistant: Add tool registry (#17331)

This PR adds a tool registry to hold tools that can be called by the
Assistant.

Currently we just have a `now` tool for retrieving the current datetime.

This is all behind the `assistant-tool-use` feature flag which currently
needs to be explicitly opted-in to in order for the LLM to see the
tools.

Release Notes:

- N/A

crates/assistant/src/assistant.rs

@@ -13,11 +13,13 @@ pub(crate) mod slash_command_picker;
 pub mod slash_command_settings;
 mod streaming_diff;
 mod terminal_inline_assistant;
+mod tools;
 mod workflow;
 
 pub use assistant_panel::{AssistantPanel, AssistantPanelEvent};
 use assistant_settings::AssistantSettings;
 use assistant_slash_command::SlashCommandRegistry;
+use assistant_tool::ToolRegistry;
 use client::{proto, Client};
 use command_palette_hooks::CommandPaletteFilter;
 pub use context::*;
@@ -214,6 +216,7 @@ pub fn init(
     prompt_library::init(cx);
     init_language_model_settings(cx);
     assistant_slash_command::init(cx);
+    assistant_tool::init(cx);
     assistant_panel::init(cx);
     context_servers::init(cx);
 
@@ -228,6 +231,7 @@ pub fn init(
     .map(Arc::new)
     .unwrap_or_else(|| Arc::new(prompts::PromptBuilder::new(None).unwrap()));
     register_slash_commands(Some(prompt_builder.clone()), cx);
+    register_tools(cx);
     inline_assistant::init(
         fs.clone(),
         prompt_builder.clone(),
@@ -401,6 +405,11 @@ fn update_slash_commands_from_settings(cx: &mut AppContext) {
     }
 }
 
+fn register_tools(cx: &mut AppContext) {
+    let tool_registry = ToolRegistry::global(cx);
+    tool_registry.register_tool(tools::now_tool::NowTool);
+}
+
 pub fn humanize_token_count(count: usize) -> String {
     match count {
         0..=999 => count.to_string(),

crates/assistant/src/context.rs

@@ -9,9 +9,11 @@ use anyhow::{anyhow, Context as _, Result};
 use assistant_slash_command::{
     SlashCommandOutput, SlashCommandOutputSection, SlashCommandRegistry,
 };
+use assistant_tool::ToolRegistry;
 use client::{self, proto, telemetry::Telemetry};
 use clock::ReplicaId;
 use collections::{HashMap, HashSet};
+use feature_flags::{FeatureFlag, FeatureFlagAppExt};
 use fs::{Fs, RemoveOptions};
 use futures::{
     future::{self, Shared},
@@ -27,7 +29,7 @@ use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, P
 use language_model::{
     LanguageModel, LanguageModelCacheConfiguration, LanguageModelCompletionEvent,
     LanguageModelImage, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
-    MessageContent, Role,
+    LanguageModelRequestTool, MessageContent, Role,
 };
 use open_ai::Model as OpenAiModel;
 use paths::{context_images_dir, contexts_dir};
@@ -1942,7 +1944,21 @@ impl Context {
         // Compute which messages to cache, including the last one.
         self.mark_cache_anchors(&model.cache_configuration(), false, cx);
 
-        let request = self.to_completion_request(cx);
+        let mut request = self.to_completion_request(cx);
+
+        if cx.has_flag::<ToolUseFeatureFlag>() {
+            let tool_registry = ToolRegistry::global(cx);
+            request.tools = tool_registry
+                .tools()
+                .into_iter()
+                .map(|tool| LanguageModelRequestTool {
+                    name: tool.name(),
+                    description: tool.description(),
+                    input_schema: tool.input_schema(),
+                })
+                .collect();
+        }
+
         let assistant_message = self
             .insert_message_after(last_message_id, Role::Assistant, MessageStatus::Pending, cx)
             .unwrap();
@@ -2788,6 +2804,16 @@ pub enum PendingSlashCommandStatus {
     Error(String),
 }
 
+pub(crate) struct ToolUseFeatureFlag;
+
+impl FeatureFlag for ToolUseFeatureFlag {
+    const NAME: &'static str = "assistant-tool-use";
+
+    fn enabled_for_staff() -> bool {
+        false
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct PendingToolUse {
     pub id: String,

crates/assistant/src/tools.rs

@@ -0,0 +1 @@
+pub mod now_tool;

crates/assistant/src/tools/now_tool.rs

@@ -0,0 +1,60 @@
+use std::sync::Arc;
+
+use anyhow::{anyhow, Result};
+use assistant_tool::Tool;
+use chrono::{Local, Utc};
+use gpui::{Task, WeakView, WindowContext};
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum Timezone {
+    /// Use UTC for the datetime.
+    Utc,
+    /// Use local time for the datetime.
+    Local,
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+pub struct FileToolInput {
+    /// The timezone to use for the datetime.
+    timezone: Timezone,
+}
+
+pub struct NowTool;
+
+impl Tool for NowTool {
+    fn name(&self) -> String {
+        "now".into()
+    }
+
+    fn description(&self) -> String {
+        "Returns the current datetime in RFC 3339 format.".into()
+    }
+
+    fn input_schema(&self) -> serde_json::Value {
+        let schema = schemars::schema_for!(FileToolInput);
+        serde_json::to_value(&schema).unwrap()
+    }
+
+    fn run(
+        self: Arc<Self>,
+        input: serde_json::Value,
+        _workspace: WeakView<workspace::Workspace>,
+        _cx: &mut WindowContext,
+    ) -> Task<Result<String>> {
+        let input: FileToolInput = match serde_json::from_value(input) {
+            Ok(input) => input,
+            Err(err) => return Task::ready(Err(anyhow!(err))),
+        };
+
+        let now = match input.timezone {
+            Timezone::Utc => Utc::now().to_rfc3339(),
+            Timezone::Local => Local::now().to_rfc3339(),
+        };
+        let text = format!("The current datetime is {now}.");
+
+        Task::ready(Ok(text))
+    }
+}