Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Reuse conversation cache when streaming edits (#30245)

Browse source 
Release Notes:

- Improved latency when the agent applies edits.
This commit is contained in:
Antonio Scandurra 2025-05-08 14:36:34 +02:00 • committed by GitHub
parent 032022e37b
commit 9f6809a28d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
50 changed files with 847 additions and 21557 deletions
Download patch file Download diff file Expand all files Collapse all files

17
crates/language_models/src/provider/anthropic.rs
View file

@ -15,7 +15,8 @@ use language_model::{
AuthenticateError, LanguageModel, LanguageModelCacheConfiguration,
LanguageModelCompletionError, LanguageModelId, LanguageModelKnownError, LanguageModelName,
LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
LanguageModelProviderState, LanguageModelRequest, MessageContent, RateLimiter, Role,
LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, MessageContent,
RateLimiter, Role,
};
use language_model::{LanguageModelCompletionEvent, LanguageModelToolUse, StopReason};
use schemars::JsonSchema;
@ -420,6 +421,14 @@ impl LanguageModel for AnthropicModel {
true
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto
| LanguageModelToolChoice::Any
| LanguageModelToolChoice::None => true,
}
}
fn telemetry_id(&self) -> String {
format!("anthropic/{}", self.model.id())
}
@ -620,7 +629,11 @@ pub fn into_anthropic(
input_schema: tool.input_schema,
})
.collect(),
tool_choice: None,
tool_choice: request.tool_choice.map(|choice| match choice {
LanguageModelToolChoice::Auto => anthropic::ToolChoice::Auto,
LanguageModelToolChoice::Any => anthropic::ToolChoice::Any,
LanguageModelToolChoice::None => anthropic::ToolChoice::None,
}),
metadata: None,
stop_sequences: Vec::new(),
temperature: request.temperature.or(Some(default_temperature)),

38
crates/language_models/src/provider/bedrock.rs
View file

@ -15,11 +15,11 @@ use bedrock::bedrock_client::types::{
StopReason,
};
use bedrock::{
BedrockAutoToolChoice, BedrockBlob, BedrockError, BedrockInnerContent, BedrockMessage,
BedrockModelMode, BedrockStreamingResponse, BedrockThinkingBlock, BedrockThinkingTextBlock,
BedrockTool, BedrockToolChoice, BedrockToolConfig, BedrockToolInputSchema,
BedrockToolResultBlock, BedrockToolResultContentBlock, BedrockToolResultStatus,
BedrockToolSpec, BedrockToolUseBlock, Model, value_to_aws_document,
BedrockAnyToolChoice, BedrockAutoToolChoice, BedrockBlob, BedrockError, BedrockInnerContent,
BedrockMessage, BedrockModelMode, BedrockStreamingResponse, BedrockThinkingBlock,
BedrockThinkingTextBlock, BedrockTool, BedrockToolChoice, BedrockToolConfig,
BedrockToolInputSchema, BedrockToolResultBlock, BedrockToolResultContentBlock,
BedrockToolResultStatus, BedrockToolSpec, BedrockToolUseBlock, Model, value_to_aws_document,
};
use collections::{BTreeMap, HashMap};
use credentials_provider::CredentialsProvider;
@ -35,8 +35,8 @@ use language_model::{
AuthenticateError, LanguageModel, LanguageModelCacheConfiguration,
LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
LanguageModelProviderState, LanguageModelRequest, LanguageModelToolUse, MessageContent,
RateLimiter, Role, TokenUsage,
LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
LanguageModelToolUse, MessageContent, RateLimiter, Role, TokenUsage,
};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@ -520,6 +520,15 @@ impl LanguageModel for BedrockModel {
self.model.supports_tool_use()
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => {
self.model.supports_tool_use()
}
LanguageModelToolChoice::None => false,
}
}
fn telemetry_id(&self) -> String {
format!("bedrock/{}", self.model.id())
}
@ -719,11 +728,20 @@ pub fn into_bedrock(
})
.collect();
let tool_choice = match request.tool_choice {
Some(LanguageModelToolChoice::Auto) | None => {
BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
}
Some(LanguageModelToolChoice::Any) => {
BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build())
}
Some(LanguageModelToolChoice::None) => {
return Err(anyhow!("LanguageModelToolChoice::None is not supported"));
}
};
let tool_config: BedrockToolConfig = BedrockToolConfig::builder()
.set_tools(Some(tool_spec))
.tool_choice(BedrockToolChoice::Auto(
BedrockAutoToolChoice::builder().build(),
))
.tool_choice(tool_choice)
.build()?;
Ok(bedrock::Request {

13
crates/language_models/src/provider/cloud.rs
View file

@ -14,8 +14,9 @@ use language_model::{
AuthenticateError, CloudModel, LanguageModel, LanguageModelCacheConfiguration,
LanguageModelCompletionError, LanguageModelId, LanguageModelKnownError, LanguageModelName,
LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
LanguageModelProviderTosView, LanguageModelRequest, LanguageModelToolSchemaFormat,
ModelRequestLimitReachedError, RateLimiter, RequestUsage, ZED_CLOUD_PROVIDER_ID,
LanguageModelProviderTosView, LanguageModelRequest, LanguageModelToolChoice,
LanguageModelToolSchemaFormat, ModelRequestLimitReachedError, RateLimiter, RequestUsage,
ZED_CLOUD_PROVIDER_ID,
};
use language_model::{
LanguageModelAvailability, LanguageModelCompletionEvent, LanguageModelProvider, LlmApiToken,
@ -686,6 +687,14 @@ impl LanguageModel for CloudLanguageModel {
}
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto
| LanguageModelToolChoice::Any
| LanguageModelToolChoice::None => true,
}
}
fn telemetry_id(&self) -> String {
format!("zed.dev/{}", self.model.id())
}

18
crates/language_models/src/provider/copilot_chat.rs
View file

@ -20,8 +20,8 @@ use language_model::{
AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,
LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,
LanguageModelRequestMessage, LanguageModelToolUse, MessageContent, RateLimiter, Role,
StopReason,
LanguageModelRequestMessage, LanguageModelToolChoice, LanguageModelToolUse, MessageContent,
RateLimiter, Role, StopReason,
};
use settings::SettingsStore;
use std::time::Duration;
@ -197,6 +197,14 @@ impl LanguageModel for CopilotChatLanguageModel {
}
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto
| LanguageModelToolChoice::Any
| LanguageModelToolChoice::None => self.supports_tools(),
}
}
fn telemetry_id(&self) -> String {
format!("copilot_chat/{}", self.model.id())
}
@ -541,7 +549,11 @@ impl CopilotChatLanguageModel {
model,
messages,
tools,
tool_choice: None,
tool_choice: request.tool_choice.map(|choice| match choice {
LanguageModelToolChoice::Auto => copilot::copilot_chat::ToolChoice::Auto,
LanguageModelToolChoice::Any => copilot::copilot_chat::ToolChoice::Any,
LanguageModelToolChoice::None => copilot::copilot_chat::ToolChoice::None,
}),
})
}
}

7
crates/language_models/src/provider/deepseek.rs
View file

@ -11,7 +11,8 @@ use http_client::HttpClient;
use language_model::{
AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,
LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, RateLimiter, Role,
LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,
LanguageModelToolChoice, RateLimiter, Role,
};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@ -282,6 +283,10 @@ impl LanguageModel for DeepSeekLanguageModel {
false
}
fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool {
false
}
fn telemetry_id(&self) -> String {
format!("deepseek/{}", self.model.id())
}

23
crates/language_models/src/provider/google.rs
View file

@ -12,8 +12,8 @@ use gpui::{
use http_client::HttpClient;
use language_model::{
AuthenticateError, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelToolSchemaFormat, LanguageModelToolUse, LanguageModelToolUseId, MessageContent,
StopReason,
LanguageModelToolChoice, LanguageModelToolSchemaFormat, LanguageModelToolUse,
LanguageModelToolUseId, MessageContent, StopReason,
};
use language_model::{
LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
@ -313,6 +313,14 @@ impl LanguageModel for GoogleLanguageModel {
true
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto
| LanguageModelToolChoice::Any
| LanguageModelToolChoice::None => true,
}
}
fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
LanguageModelToolSchemaFormat::JsonSchemaSubset
}
@ -484,7 +492,16 @@ pub fn into_google(
.collect(),
}]
}),
tool_config: None,
tool_config: request.tool_choice.map(|choice| google_ai::ToolConfig {
function_calling_config: google_ai::FunctionCallingConfig {
mode: match choice {
LanguageModelToolChoice::Auto => google_ai::FunctionCallingMode::Auto,
LanguageModelToolChoice::Any => google_ai::FunctionCallingMode::Any,
LanguageModelToolChoice::None => google_ai::FunctionCallingMode::None,
},
allowed_function_names: None,
},
}),
}
}

5
crates/language_models/src/provider/lmstudio.rs
View file

@ -4,6 +4,7 @@ use gpui::{AnyView, App, AsyncApp, Context, Subscription, Task};
use http_client::HttpClient;
use language_model::{
AuthenticateError, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelToolChoice,
};
use language_model::{
LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
@ -284,6 +285,10 @@ impl LanguageModel for LmStudioLanguageModel {
false
}
fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool {
false
}
fn telemetry_id(&self) -> String {
format!("lmstudio/{}", self.model.id())
}

7
crates/language_models/src/provider/mistral.rs
View file

@ -10,7 +10,8 @@ use http_client::HttpClient;
use language_model::{
AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,
LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, RateLimiter, Role,
LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,
LanguageModelToolChoice, RateLimiter, Role,
};
use futures::stream::BoxStream;
@ -302,6 +303,10 @@ impl LanguageModel for MistralLanguageModel {
false
}
fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool {
false
}
fn telemetry_id(&self) -> String {
format!("mistral/{}", self.model.id())
}

11
crates/language_models/src/provider/ollama.rs
View file

@ -5,7 +5,8 @@ use gpui::{AnyView, App, AsyncApp, Context, Subscription, Task};
use http_client::HttpClient;
use language_model::{
AuthenticateError, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelRequestTool, LanguageModelToolUse, LanguageModelToolUseId, StopReason,
LanguageModelRequestTool, LanguageModelToolChoice, LanguageModelToolUse,
LanguageModelToolUseId, StopReason,
};
use language_model::{
LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
@ -324,6 +325,14 @@ impl LanguageModel for OllamaLanguageModel {
self.model.supports_tools.unwrap_or(false)
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto => false,
LanguageModelToolChoice::Any => false,
LanguageModelToolChoice::None => false,
}
}
fn telemetry_id(&self) -> String {
format!("ollama/{}", self.model.id())
}

16
crates/language_models/src/provider/open_ai.rs
View file

@ -12,7 +12,7 @@ use language_model::{
AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,
LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,
LanguageModelToolUse, MessageContent, RateLimiter, Role, StopReason,
LanguageModelToolChoice, LanguageModelToolUse, MessageContent, RateLimiter, Role, StopReason,
};
use open_ai::{Model, ResponseStreamEvent, stream_completion};
use schemars::JsonSchema;
@ -295,6 +295,14 @@ impl LanguageModel for OpenAiLanguageModel {
true
}
fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
match choice {
LanguageModelToolChoice::Auto => true,
LanguageModelToolChoice::Any => true,
LanguageModelToolChoice::None => true,
}
}
fn telemetry_id(&self) -> String {
format!("openai/{}", self.model.id())
}
@ -417,7 +425,11 @@ pub fn into_open_ai(
},
})
.collect(),
tool_choice: None,
tool_choice: request.tool_choice.map(|choice| match choice {
LanguageModelToolChoice::Auto => open_ai::ToolChoice::Auto,
LanguageModelToolChoice::Any => open_ai::ToolChoice::Required,
LanguageModelToolChoice::None => open_ai::ToolChoice::None,
}),
}
}