oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add tool calling support for Gemini models (#27772)

Browse Source 
Release Notes:

- N/A
This commit is contained in:
Bennet Bo Fenner
2025-03-31 17:46:42 +02:00
committed by GitHub
parent f6d58f76e4
commit c8a9a74e6a
32 changed files with 735 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/assistant2/src/active_thread.rs
+4 -1
View File
@@ -1787,12 +1787,13 @@ impl ActiveThread {
fn handle_deny_tool(
&mut self,
tool_use_id: LanguageModelToolUseId,
tool_name: Arc<str>,
_: &ClickEvent,
_window: &mut Window,
cx: &mut Context<Self>,
) {
self.thread.update(cx, |thread, cx| {
thread.deny_tool_use(tool_use_id, cx);
thread.deny_tool_use(tool_use_id, tool_name, cx);
});
}
@@ -1865,10 +1866,12 @@ impl ActiveThread {
})
.child({
let tool_id = tool.id.clone();
let tool_name = tool.name.clone();
Button::new("deny-tool", "Deny").on_click(cx.listener(
move |this, event, window, cx| {
this.handle_deny_tool(
tool_id.clone(),
tool_name.clone(),
event,
window,
cx,
crates/assistant2/src/thread.rs
+29 -14
View File
@@ -777,7 +777,7 @@ impl Thread {
LanguageModelRequestTool {
name: tool.name(),
description: tool.description(),
input_schema: tool.input_schema(),
input_schema: tool.input_schema(model.tool_input_format()),
}
}));
@@ -1030,17 +1030,23 @@ impl Thread {
}
}
LanguageModelCompletionEvent::ToolUse(tool_use) => {
if let Some(last_assistant_message) = thread
let last_assistant_message_id = thread
.messages
.iter()
.rfind(|message| message.role == Role::Assistant)
{
thread.tool_use.request_tool_use(
last_assistant_message.id,
tool_use,
cx,
);
}
.map(|message| message.id)
.unwrap_or_else(|| {
thread.insert_message(
Role::Assistant,
vec![MessageSegment::Text("Using tool...".to_string())],
cx,
)
});
thread.tool_use.request_tool_use(
last_assistant_message_id,
tool_use,
cx,
);
}
}
@@ -1257,6 +1263,7 @@ impl Thread {
tool: Arc<dyn Tool>,
cx: &mut Context<Thread>,
) -> Task<()> {
let tool_name: Arc<str> = tool.name().into();
let run_tool = tool.run(
input,
messages,
@@ -1271,9 +1278,11 @@ impl Thread {
thread
.update(cx, |thread, cx| {
let pending_tool_use = thread
.tool_use
.insert_tool_output(tool_use_id.clone(), output);
let pending_tool_use = thread.tool_use.insert_tool_output(
tool_use_id.clone(),
tool_name,
output,
);
cx.emit(ThreadEvent::ToolFinished {
tool_use_id,
@@ -1561,12 +1570,18 @@ impl Thread {
self.cumulative_token_usage.clone()
}
pub fn deny_tool_use(&mut self, tool_use_id: LanguageModelToolUseId, cx: &mut Context<Self>) {
pub fn deny_tool_use(
&mut self,
tool_use_id: LanguageModelToolUseId,
tool_name: Arc<str>,
cx: &mut Context<Self>,
) {
let err = Err(anyhow::anyhow!(
"Permission to run tool action denied by user"
));
self.tool_use.insert_tool_output(tool_use_id.clone(), err);
self.tool_use
.insert_tool_output(tool_use_id.clone(), tool_name, err);
cx.emit(ThreadEvent::ToolFinished {
tool_use_id,
crates/assistant2/src/tool_use.rs
+6
View File
@@ -113,6 +113,7 @@ impl ToolUseState {
tool_use_id.clone(),
LanguageModelToolResult {
tool_use_id,
tool_name: tool_use.clone(),
is_error: tool_result.is_error,
content: tool_result.content.clone(),
},
@@ -134,6 +135,7 @@ impl ToolUseState {
tool_use_id.clone(),
LanguageModelToolResult {
tool_use_id,
tool_name: tool_use.name.clone(),
content: "Tool canceled by user".into(),
is_error: true,
},
@@ -313,6 +315,7 @@ impl ToolUseState {
pub fn insert_tool_output(
&mut self,
tool_use_id: LanguageModelToolUseId,
tool_name: Arc<str>,
output: Result<String>,
) -> Option<PendingToolUse> {
match output {
@@ -321,6 +324,7 @@ impl ToolUseState {
tool_use_id.clone(),
LanguageModelToolResult {
tool_use_id: tool_use_id.clone(),
tool_name,
content: tool_result.into(),
is_error: false,
},
@@ -332,6 +336,7 @@ impl ToolUseState {
tool_use_id.clone(),
LanguageModelToolResult {
tool_use_id: tool_use_id.clone(),
tool_name,
content: err.to_string().into(),
is_error: true,
},
@@ -379,6 +384,7 @@ impl ToolUseState {
request_message.content.push(MessageContent::ToolResult(
LanguageModelToolResult {
tool_use_id: tool_use_id.clone(),
tool_name: tool_result.tool_name.clone(),
is_error: tool_result.is_error,
content: if tool_result.content.is_empty() {
// Surprisingly, the API fails if we return an empty string here.
crates/assistant_tool/src/assistant_tool.rs
+2 -1
View File
@@ -11,6 +11,7 @@ use anyhow::Result;
use gpui::{App, Entity, SharedString, Task};
use icons::IconName;
use language_model::LanguageModelRequestMessage;
use language_model::LanguageModelToolSchemaFormat;
use project::Project;
pub use crate::action_log::*;
@@ -50,7 +51,7 @@ pub trait Tool: 'static + Send + Sync {
fn needs_confirmation(&self) -> bool;
/// Returns the JSON schema that describes the tool's input.
fn input_schema(&self) -> serde_json::Value {
fn input_schema(&self, _: LanguageModelToolSchemaFormat) -> serde_json::Value {
serde_json::Value::Object(serde_json::Map::default())
}
crates/assistant_tools/src/assistant_tools.rs
+1
View File
@@ -18,6 +18,7 @@ mod path_search_tool;
mod read_file_tool;
mod regex_search_tool;
mod replace;
mod schema;
mod symbol_info_tool;
mod thinking_tool;
crates/assistant_tools/src/bash_tool.rs
+4 -4
View File
@@ -1,7 +1,8 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Context as _, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -38,9 +39,8 @@ impl Tool for BashTool {
IconName::Terminal
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(BashToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<BashToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/batch_tool.rs
+4 -4
View File
@@ -1,8 +1,9 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool, ToolWorkingSet};
use futures::future::join_all;
use gpui::{App, AppContext, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -162,9 +163,8 @@ impl Tool for BatchTool {
IconName::Cog
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(BatchToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<BatchToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/code_symbols_tool.rs
+4 -4
View File
@@ -7,7 +7,7 @@ use assistant_tool::{ActionLog, Tool};
use collections::IndexMap;
use gpui::{App, AsyncApp, Entity, Task};
use language::{CodeLabel, Language, LanguageRegistry};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use lsp::SymbolKind;
use project::{DocumentSymbol, Project, Symbol};
use regex::{Regex, RegexBuilder};
@@ -17,6 +17,7 @@ use ui::IconName;
use util::markdown::MarkdownString;
use crate::code_symbol_iter::{CodeSymbolIterator, Entry};
use crate::schema::json_schema_for;
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
pub struct CodeSymbolsInput {
@@ -93,9 +94,8 @@ impl Tool for CodeSymbolsTool {
IconName::Eye
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(CodeSymbolsInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<CodeSymbolsInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/copy_path_tool.rs
+4 -3
View File
@@ -1,7 +1,9 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, AppContext, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::LanguageModelToolSchemaFormat;
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -53,9 +55,8 @@ impl Tool for CopyPathTool {
IconName::Clipboard
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(CopyPathToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<CopyPathToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/create_directory_tool.rs
+4 -3
View File
@@ -1,7 +1,9 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::LanguageModelToolSchemaFormat;
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -43,9 +45,8 @@ impl Tool for CreateDirectoryTool {
IconName::Folder
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(CreateDirectoryToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<CreateDirectoryToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/create_file_tool.rs
+4 -3
View File
@@ -1,7 +1,9 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::LanguageModelToolSchemaFormat;
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -50,9 +52,8 @@ impl Tool for CreateFileTool {
IconName::FileCreate
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(CreateFileToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<CreateFileToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/delete_path_tool.rs
+4 -4
View File
@@ -1,8 +1,9 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use futures::{channel::mpsc, SinkExt, StreamExt};
use gpui::{App, AppContext, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::{Project, ProjectPath};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -44,9 +45,8 @@ impl Tool for DeletePathTool {
IconName::FileDelete
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(DeletePathToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<DeletePathToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/diagnostics_tool.rs
+4 -4
View File
@@ -1,8 +1,9 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, Entity, Task};
use language::{DiagnosticSeverity, OffsetRangeExt};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -57,9 +58,8 @@ impl Tool for DiagnosticsTool {
IconName::Warning
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(DiagnosticsToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<DiagnosticsToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/edit_files_tool.rs
+4 -3
View File
@@ -2,12 +2,14 @@ mod edit_action;
pub mod log;
use crate::replace::{replace_exact, replace_with_flexible_indent};
use crate::schema::json_schema_for;
use anyhow::{anyhow, Context, Result};
use assistant_tool::{ActionLog, Tool};
use collections::HashSet;
use edit_action::{EditAction, EditActionParser};
use futures::{channel::mpsc, SinkExt, StreamExt};
use gpui::{App, AppContext, AsyncApp, Entity, Task};
use language_model::LanguageModelToolSchemaFormat;
use language_model::{
LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
};
@@ -91,9 +93,8 @@ impl Tool for EditFilesTool {
IconName::Pencil
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(EditFilesToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<EditFilesToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/fetch_tool.rs
+4 -4
View File
@@ -2,13 +2,14 @@ use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;
use crate::schema::json_schema_for;
use anyhow::{anyhow, bail, Context as _, Result};
use assistant_tool::{ActionLog, Tool};
use futures::AsyncReadExt as _;
use gpui::{App, AppContext as _, Entity, Task};
use html_to_markdown::{convert_html_to_markdown, markdown, TagHandler};
use http_client::{AsyncBody, HttpClientWithUrl};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -127,9 +128,8 @@ impl Tool for FetchTool {
IconName::Globe
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(FetchToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<FetchToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/find_replace_file_tool.rs
+4 -4
View File
@@ -1,7 +1,8 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Context as _, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, AppContext, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -140,9 +141,8 @@ impl Tool for FindReplaceFileTool {
IconName::Pencil
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(FindReplaceFileToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<FindReplaceFileToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/list_directory_tool.rs
+4 -4
View File
@@ -1,7 +1,8 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -55,9 +56,8 @@ impl Tool for ListDirectoryTool {
IconName::Folder
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(ListDirectoryToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<ListDirectoryToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/move_path_tool.rs
+4 -4
View File
@@ -1,7 +1,8 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, AppContext, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -53,9 +54,8 @@ impl Tool for MovePathTool {
IconName::ArrowRightLeft
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(MovePathToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<MovePathToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/now_tool.rs
+4 -4
View File
@@ -1,10 +1,11 @@
use std::sync::Arc;
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use chrono::{Local, Utc};
use gpui::{App, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -44,9 +45,8 @@ impl Tool for NowTool {
IconName::Info
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(NowToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<NowToolInput>(format)
}
fn ui_text(&self, _input: &serde_json::Value) -> String {
crates/assistant_tools/src/open_tool.rs
+4 -4
View File
@@ -1,7 +1,8 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Context as _, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, AppContext, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -34,9 +35,8 @@ impl Tool for OpenTool {
IconName::ExternalLink
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(OpenToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<OpenToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/path_search_tool.rs
+4 -4
View File
@@ -1,7 +1,8 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, AppContext, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -52,9 +53,8 @@ impl Tool for PathSearchTool {
IconName::SearchCode
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(PathSearchToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<PathSearchToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/read_file_tool.rs
+4 -4
View File
@@ -1,11 +1,12 @@
use std::path::Path;
use std::sync::Arc;
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, Entity, Task};
use itertools::Itertools;
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -58,9 +59,8 @@ impl Tool for ReadFileTool {
IconName::Eye
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(ReadFileToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<ReadFileToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/regex_search_tool.rs
+4 -4
View File
@@ -1,9 +1,10 @@
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use futures::StreamExt;
use gpui::{App, Entity, Task};
use language::OffsetRangeExt;
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::{
search::{SearchQuery, SearchResult},
Project,
@@ -55,9 +56,8 @@ impl Tool for RegexSearchTool {
IconName::Regex
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(RegexSearchToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<RegexSearchToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/schema.rs
+120
View File
@@ -0,0 +1,120 @@
use anyhow::Result;
use language_model::LanguageModelToolSchemaFormat;
use schemars::{
schema::{RootSchema, Schema, SchemaObject},
JsonSchema,
};
pub fn json_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> serde_json::Value {
let schema = root_schema_for::<T>(format);
schema_to_json(&schema, format).expect("Failed to convert tool calling schema to JSON")
}
pub fn schema_to_json(
schema: &RootSchema,
format: LanguageModelToolSchemaFormat,
) -> Result<serde_json::Value> {
let mut value = serde_json::to_value(schema)?;
match format {
LanguageModelToolSchemaFormat::JsonSchema => Ok(value),
LanguageModelToolSchemaFormat::JsonSchemaSubset => {
transform_fields_to_json_schema_subset(&mut value);
Ok(value)
}
}
}
fn root_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> RootSchema {
let mut generator = match format {
LanguageModelToolSchemaFormat::JsonSchema => schemars::SchemaGenerator::default(),
LanguageModelToolSchemaFormat::JsonSchemaSubset => {
schemars::r#gen::SchemaSettings::default()
.with(|settings| {
settings.meta_schema = None;
settings.inline_subschemas = true;
settings
.visitors
.push(Box::new(TransformToJsonSchemaSubsetVisitor));
})
.into_generator()
}
};
generator.root_schema_for::<T>()
}
#[derive(Debug, Clone)]
struct TransformToJsonSchemaSubsetVisitor;
impl schemars::visit::Visitor for TransformToJsonSchemaSubsetVisitor {
fn visit_root_schema(&mut self, root: &mut RootSchema) {
schemars::visit::visit_root_schema(self, root)
}
fn visit_schema(&mut self, schema: &mut Schema) {
schemars::visit::visit_schema(self, schema)
}
fn visit_schema_object(&mut self, schema: &mut SchemaObject) {
// Ensure that the type field is not an array, this happens when we use
// Option<T>, the type will be [T, "null"].
if let Some(instance_type) = schema.instance_type.take() {
schema.instance_type = match instance_type {
schemars::schema::SingleOrVec::Single(t) => {
Some(schemars::schema::SingleOrVec::Single(t))
}
schemars::schema::SingleOrVec::Vec(items) => items
.into_iter()
.next()
.map(schemars::schema::SingleOrVec::from),
};
}
// One of is not supported, use anyOf instead.
if let Some(subschema) = schema.subschemas.as_mut() {
if let Some(one_of) = subschema.one_of.take() {
subschema.any_of = Some(one_of);
}
}
schemars::visit::visit_schema_object(self, schema)
}
}
fn transform_fields_to_json_schema_subset(json: &mut serde_json::Value) {
if let serde_json::Value::Object(obj) = json {
if let Some(default) = obj.get("default") {
let is_null = default.is_null();
//Default is not supported, so we need to remove it.
obj.remove("default");
if is_null {
obj.insert("nullable".to_string(), serde_json::Value::Bool(true));
}
}
// If a type is not specified for an input parameter we need to add it.
if obj.contains_key("description")
&& !obj.contains_key("type")
&& !(obj.contains_key("anyOf")
|| obj.contains_key("oneOf")
|| obj.contains_key("allOf"))
{
obj.insert(
"type".to_string(),
serde_json::Value::String("string".to_string()),
);
}
//Format field is only partially supported (e.g. not uint compatibility)
obj.remove("format");
for (_, value) in obj.iter_mut() {
if let serde_json::Value::Object(_) | serde_json::Value::Array(_) = value {
transform_fields_to_json_schema_subset(value);
}
}
} else if let serde_json::Value::Array(arr) = json {
for item in arr.iter_mut() {
transform_fields_to_json_schema_subset(item);
}
}
}
crates/assistant_tools/src/symbol_info_tool.rs
+5 -4
View File
@@ -2,7 +2,7 @@ use anyhow::{anyhow, Context as _, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, AsyncApp, Entity, Task};
use language::{self, Anchor, Buffer, BufferSnapshot, Location, Point, ToPoint, ToPointUtf16};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -10,6 +10,8 @@ use std::{fmt::Write, ops::Range, sync::Arc};
use ui::IconName;
use util::markdown::MarkdownString;
use crate::schema::json_schema_for;
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
pub struct SymbolInfoToolInput {
/// The relative path to the file containing the symbol.
@@ -82,9 +84,8 @@ impl Tool for SymbolInfoTool {
IconName::Eye
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(SymbolInfoToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<SymbolInfoToolInput>(format)
}
fn ui_text(&self, input: &serde_json::Value) -> String {
crates/assistant_tools/src/thinking_tool.rs
+4 -4
View File
@@ -1,9 +1,10 @@
use std::sync::Arc;
use crate::schema::json_schema_for;
use anyhow::{anyhow, Result};
use assistant_tool::{ActionLog, Tool};
use gpui::{App, Entity, Task};
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -35,9 +36,8 @@ impl Tool for ThinkingTool {
IconName::Brain
}
fn input_schema(&self) -> serde_json::Value {
let schema = schemars::schema_for!(ThinkingToolInput);
serde_json::to_value(&schema).unwrap()
fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
json_schema_for::<ThinkingToolInput>(format)
}
fn ui_text(&self, _input: &serde_json::Value) -> String {
crates/context_server/src/context_server_tool.rs
+2 -2
View File
@@ -4,7 +4,7 @@ use anyhow::{anyhow, bail, Result};
use assistant_tool::{ActionLog, Tool, ToolSource};
use gpui::{App, Entity, Task};
use icons::IconName;
use language_model::LanguageModelRequestMessage;
use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
use project::Project;
use crate::manager::ContextServerManager;
@@ -53,7 +53,7 @@ impl Tool for ContextServerTool {
true
}
fn input_schema(&self) -> serde_json::Value {
fn input_schema(&self, _: LanguageModelToolSchemaFormat) -> serde_json::Value {
match &self.tool.input_schema {
serde_json::Value::Null => {
serde_json::json!({ "type": "object", "properties": [] })
crates/google_ai/src/google_ai.rs
+77
View File
@@ -127,6 +127,10 @@ pub struct GenerateContentRequest {
pub contents: Vec<Content>,
pub generation_config: Option<GenerationConfig>,
pub safety_settings: Option<Vec<SafetySetting>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Tool>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_config: Option<ToolConfig>,
}
#[derive(Debug, Serialize, Deserialize)]
@@ -134,6 +138,7 @@ pub struct GenerateContentRequest {
pub struct GenerateContentResponse {
pub candidates: Option<Vec<GenerateContentCandidate>>,
pub prompt_feedback: Option<PromptFeedback>,
pub usage_metadata: Option<UsageMetadata>,
}
#[derive(Debug, Serialize, Deserialize)]
@@ -166,6 +171,8 @@ pub enum Role {
pub enum Part {
TextPart(TextPart),
InlineDataPart(InlineDataPart),
FunctionCallPart(FunctionCallPart),
FunctionResponsePart(FunctionResponsePart),
}
#[derive(Debug, Serialize, Deserialize)]
@@ -187,6 +194,18 @@ pub struct GenerativeContentBlob {
pub data: String,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FunctionCallPart {
pub function_call: FunctionCall,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FunctionResponsePart {
pub function_response: FunctionResponse,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CitationSource {
@@ -210,6 +229,17 @@ pub struct PromptFeedback {
pub block_reason_message: Option<String>,
}
#[derive(Debug, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct UsageMetadata {
pub prompt_token_count: Option<usize>,
pub cached_content_token_count: Option<usize>,
pub candidates_token_count: Option<usize>,
pub tool_use_prompt_token_count: Option<usize>,
pub thoughts_token_count: Option<usize>,
pub total_token_count: Option<usize>,
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GenerationConfig {
@@ -298,6 +328,53 @@ pub struct CountTokensResponse {
pub total_tokens: usize,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct FunctionCall {
pub name: String,
pub args: serde_json::Value,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct FunctionResponse {
pub name: String,
pub response: serde_json::Value,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Tool {
pub function_declarations: Vec<FunctionDeclaration>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolConfig {
pub function_calling_config: FunctionCallingConfig,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FunctionCallingConfig {
pub mode: FunctionCallingMode,
#[serde(skip_serializing_if = "Option::is_none")]
pub allowed_function_names: Option<Vec<String>>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum FunctionCallingMode {
Auto,
Any,
None,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct FunctionDeclaration {
pub name: String,
pub description: String,
pub parameters: serde_json::Value,
}
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[derive(Clone, Default, Debug, Deserialize, Serialize, PartialEq, Eq, strum::EnumIter)]
pub enum Model {
crates/language_model/src/language_model.rs
+13
View File
@@ -66,6 +66,15 @@ pub enum LanguageModelCompletionEvent {
UsageUpdate(TokenUsage),
}
/// Indicates the format used to define the input schema for a language model tool.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum LanguageModelToolSchemaFormat {
/// A JSON schema, see https://json-schema.org
JsonSchema,
/// A subset of an OpenAPI 3.0 schema object supported by Google AI, see https://ai.google.dev/api/caching#Schema
JsonSchemaSubset,
}
#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
@@ -176,6 +185,10 @@ pub trait LanguageModel: Send + Sync {
LanguageModelAvailability::Public
}
fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
LanguageModelToolSchemaFormat::JsonSchema
}
fn max_token_count(&self) -> usize;
fn max_output_tokens(&self) -> Option<u32> {
None
crates/language_model/src/request.rs
+1
View File
@@ -167,6 +167,7 @@ impl LanguageModelImage {
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
pub struct LanguageModelToolResult {
pub tool_use_id: LanguageModelToolUseId,
pub tool_name: Arc<str>,
pub is_error: bool,
pub content: Arc<str>,
}
crates/language_models/src/provider/google.rs
+277 -39
View File
@@ -2,13 +2,17 @@ use anyhow::{anyhow, Context as _, Result};
use collections::BTreeMap;
use credentials_provider::CredentialsProvider;
use editor::{Editor, EditorElement, EditorStyle};
use futures::Stream;
use futures::{future::BoxFuture, FutureExt, StreamExt};
use google_ai::stream_generate_content;
use google_ai::{FunctionDeclaration, GenerateContentResponse, Part, UsageMetadata};
use gpui::{
AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace,
};
use http_client::HttpClient;
use language_model::{AuthenticateError, LanguageModelCompletionEvent};
use language_model::{
AuthenticateError, LanguageModelCompletionEvent, LanguageModelToolSchemaFormat,
LanguageModelToolUse, LanguageModelToolUseId, StopReason,
};
use language_model::{
LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
@@ -17,7 +21,8 @@ use language_model::{
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use settings::{Settings, SettingsStore};
use std::{future, sync::Arc};
use std::pin::Pin;
use std::sync::Arc;
use strum::IntoEnumIterator;
use theme::ThemeSettings;
use ui::{prelude::*, Icon, IconName, List, Tooltip};
@@ -174,7 +179,7 @@ impl LanguageModelProvider for GoogleLanguageModelProvider {
model,
state: self.state.clone(),
http_client: self.http_client.clone(),
rate_limiter: RateLimiter::new(4),
request_limiter: RateLimiter::new(4),
}))
}
@@ -211,7 +216,7 @@ impl LanguageModelProvider for GoogleLanguageModelProvider {
model,
state: self.state.clone(),
http_client: self.http_client.clone(),
rate_limiter: RateLimiter::new(4),
request_limiter: RateLimiter::new(4),
}) as Arc<dyn LanguageModel>
})
.collect()
@@ -240,7 +245,39 @@ pub struct GoogleLanguageModel {
model: google_ai::Model,
state: gpui::Entity<State>,
http_client: Arc<dyn HttpClient>,
rate_limiter: RateLimiter,
request_limiter: RateLimiter,
}
impl GoogleLanguageModel {
fn stream_completion(
&self,
request: google_ai::GenerateContentRequest,
cx: &AsyncApp,
) -> BoxFuture<
'static,
Result<futures::stream::BoxStream<'static, Result<GenerateContentResponse>>>,
> {
let http_client = self.http_client.clone();
let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| {
let settings = &AllLanguageModelSettings::get_global(cx).google;
(state.api_key.clone(), settings.api_url.clone())
}) else {
return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
};
async move {
let api_key = api_key.ok_or_else(|| anyhow!("Missing Google API key"))?;
let request = google_ai::stream_generate_content(
http_client.as_ref(),
&api_url,
&api_key,
request,
);
request.await.context("failed to stream completion")
}
.boxed()
}
}
impl LanguageModel for GoogleLanguageModel {
@@ -260,6 +297,10 @@ impl LanguageModel for GoogleLanguageModel {
LanguageModelProviderName(PROVIDER_NAME.into())
}
fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
LanguageModelToolSchemaFormat::JsonSchemaSubset
}
fn telemetry_id(&self) -> String {
format!("google/{}", self.model.id())
}
@@ -305,40 +346,67 @@ impl LanguageModel for GoogleLanguageModel {
Result<futures::stream::BoxStream<'static, Result<LanguageModelCompletionEvent>>>,
> {
let request = into_google(request, self.model.id().to_string());
let http_client = self.http_client.clone();
let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| {
let settings = &AllLanguageModelSettings::get_global(cx).google;
(state.api_key.clone(), settings.api_url.clone())
}) else {
return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
};
let future = self.rate_limiter.stream(async move {
let api_key = api_key.ok_or_else(|| anyhow!("Missing Google API Key"))?;
let response =
stream_generate_content(http_client.as_ref(), &api_url, &api_key, request);
let events = response.await?;
Ok(google_ai::extract_text_from_events(events).boxed())
let request = self.stream_completion(request, cx);
let future = self.request_limiter.stream(async move {
let response = request.await.map_err(|err| anyhow!(err))?;
Ok(map_to_language_model_completion_events(response))
});
async move {
Ok(future
.await?
.map(|result| result.map(LanguageModelCompletionEvent::Text))
.boxed())
}
.boxed()
async move { Ok(future.await?.boxed()) }.boxed()
}
fn use_any_tool(
&self,
_request: LanguageModelRequest,
_name: String,
_description: String,
_schema: serde_json::Value,
_cx: &AsyncApp,
request: LanguageModelRequest,
name: String,
description: String,
schema: serde_json::Value,
cx: &AsyncApp,
) -> BoxFuture<'static, Result<futures::stream::BoxStream<'static, Result<String>>>> {
future::ready(Err(anyhow!("not implemented"))).boxed()
let mut request = into_google(request, self.model.id().to_string());
request.tools = Some(vec![google_ai::Tool {
function_declarations: vec![google_ai::FunctionDeclaration {
name: name.clone(),
description,
parameters: schema,
}],
}]);
request.tool_config = Some(google_ai::ToolConfig {
function_calling_config: google_ai::FunctionCallingConfig {
mode: google_ai::FunctionCallingMode::Any,
allowed_function_names: Some(vec![name]),
},
});
let response = self.stream_completion(request, cx);
self.request_limiter
.run(async move {
let response = response.await?;
Ok(response
.filter_map(|event| async move {
match event {
Ok(response) => {
if let Some(candidates) = &response.candidates {
for candidate in candidates {
for part in &candidate.content.parts {
if let google_ai::Part::FunctionCallPart(
function_call_part,
) = part
{
return Some(Ok(serde_json::to_string(
&function_call_part.function_call.args,
)
.unwrap_or_default()));
}
}
}
}
None
}
Err(e) => Some(Err(e)),
}
})
.boxed())
})
.boxed()
}
}
@@ -351,11 +419,41 @@ pub fn into_google(
contents: request
.messages
.into_iter()
.map(|msg| google_ai::Content {
parts: vec![google_ai::Part::TextPart(google_ai::TextPart {
text: msg.string_contents(),
})],
role: match msg.role {
.map(|message| google_ai::Content {
parts: message
.content
.into_iter()
.filter_map(|content| match content {
language_model::MessageContent::Text(text) => {
if !text.is_empty() {
Some(Part::TextPart(google_ai::TextPart { text }))
} else {
None
}
}
language_model::MessageContent::Image(_) => None,
language_model::MessageContent::ToolUse(tool_use) => {
Some(Part::FunctionCallPart(google_ai::FunctionCallPart {
function_call: google_ai::FunctionCall {
name: tool_use.name.to_string(),
args: tool_use.input,
},
}))
}
language_model::MessageContent::ToolResult(tool_result) => Some(
Part::FunctionResponsePart(google_ai::FunctionResponsePart {
function_response: google_ai::FunctionResponse {
name: tool_result.tool_name.to_string(),
// The API expects a valid JSON object
response: serde_json::json!({
"output": tool_result.content
}),
},
}),
),
})
.collect(),
role: match message.role {
Role::User => google_ai::Role::User,
Role::Assistant => google_ai::Role::Model,
Role::System => google_ai::Role::User, // Google AI doesn't have a system role
@@ -371,9 +469,119 @@ pub fn into_google(
top_k: None,
}),
safety_settings: None,
tools: Some(
request
.tools
.into_iter()
.map(|tool| google_ai::Tool {
function_declarations: vec![FunctionDeclaration {
name: tool.name,
description: tool.description,
parameters: tool.input_schema,
}],
})
.collect(),
),
tool_config: None,
}
}
pub fn map_to_language_model_completion_events(
events: Pin<Box<dyn Send + Stream<Item = Result<GenerateContentResponse>>>>,
) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> {
use std::sync::atomic::{AtomicU64, Ordering};
static TOOL_CALL_COUNTER: AtomicU64 = AtomicU64::new(0);
struct State {
events: Pin<Box<dyn Send + Stream<Item = Result<GenerateContentResponse>>>>,
usage: UsageMetadata,
stop_reason: StopReason,
}
futures::stream::unfold(
State {
events,
usage: UsageMetadata::default(),
stop_reason: StopReason::EndTurn,
},
|mut state| async move {
if let Some(event) = state.events.next().await {
match event {
Ok(event) => {
let mut events: Vec<Result<LanguageModelCompletionEvent>> = Vec::new();
let mut wants_to_use_tool = false;
if let Some(usage_metadata) = event.usage_metadata {
update_usage(&mut state.usage, &usage_metadata);
events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(
convert_usage(&state.usage),
)))
}
if let Some(candidates) = event.candidates {
for candidate in candidates {
if let Some(finish_reason) = candidate.finish_reason.as_deref() {
state.stop_reason = match finish_reason {
"STOP" => StopReason::EndTurn,
"MAX_TOKENS" => StopReason::MaxTokens,
_ => {
log::error!(
"Unexpected google finish_reason: {finish_reason}"
);
StopReason::EndTurn
}
};
}
candidate
.content
.parts
.into_iter()
.for_each(|part| match part {
Part::TextPart(text_part) => events.push(Ok(
LanguageModelCompletionEvent::Text(text_part.text),
)),
Part::InlineDataPart(_) => {}
Part::FunctionCallPart(function_call_part) => {
wants_to_use_tool = true;
let name: Arc<str> =
function_call_part.function_call.name.into();
let next_tool_id =
TOOL_CALL_COUNTER.fetch_add(1, Ordering::SeqCst);
let id: LanguageModelToolUseId =
format!("{}-{}", name, next_tool_id).into();
events.push(Ok(LanguageModelCompletionEvent::ToolUse(
LanguageModelToolUse {
id,
name,
input: function_call_part.function_call.args,
},
)));
}
Part::FunctionResponsePart(_) => {}
});
}
}
// Even when Gemini wants to use a Tool, the API
// responds with `finish_reason: STOP`
if wants_to_use_tool {
state.stop_reason = StopReason::ToolUse;
}
events.push(Ok(LanguageModelCompletionEvent::Stop(state.stop_reason)));
return Some((events, state));
}
Err(err) => {
return Some((vec![Err(anyhow!(err))], state));
}
}
}
None
},
)
.flat_map(futures::stream::iter)
}
pub fn count_google_tokens(
request: LanguageModelRequest,
cx: &App,
@@ -403,6 +611,36 @@ pub fn count_google_tokens(
.boxed()
}
fn update_usage(usage: &mut UsageMetadata, new: &UsageMetadata) {
if let Some(prompt_token_count) = new.prompt_token_count {
usage.prompt_token_count = Some(prompt_token_count);
}
if let Some(cached_content_token_count) = new.cached_content_token_count {
usage.cached_content_token_count = Some(cached_content_token_count);
}
if let Some(candidates_token_count) = new.candidates_token_count {
usage.candidates_token_count = Some(candidates_token_count);
}
if let Some(tool_use_prompt_token_count) = new.tool_use_prompt_token_count {
usage.tool_use_prompt_token_count = Some(tool_use_prompt_token_count);
}
if let Some(thoughts_token_count) = new.thoughts_token_count {
usage.thoughts_token_count = Some(thoughts_token_count);
}
if let Some(total_token_count) = new.total_token_count {
usage.total_token_count = Some(total_token_count);
}
}
fn convert_usage(usage: &UsageMetadata) -> language_model::TokenUsage {
language_model::TokenUsage {
input_tokens: usage.prompt_token_count.unwrap_or(0) as u32,
output_tokens: usage.candidates_token_count.unwrap_or(0) as u32,
cache_read_input_tokens: usage.cached_content_token_count.unwrap_or(0) as u32,
cache_creation_input_tokens: 0,
}
}
struct ConfigurationView {
api_key_editor: Entity<Editor>,
state: gpui::Entity<State>,