From 73bb5fe1d4e5bdfc2370d603038fc9162eee5d77 Mon Sep 17 00:00:00 2001 From: dstoc <539597+dstoc@users.noreply.github.com> Date: Fri, 12 Sep 2025 07:06:31 +0000 Subject: [PATCH 1/2] feat(llm): add a synthetic _id param to tool call arguments --- crates/llm/AGENTS.md | 1 + crates/llm/src/gemini_rust.rs | 2 +- crates/llm/src/harmony.rs | 6 ++++-- crates/llm/src/lib.rs | 11 +++++++++++ crates/llm/src/ollama.rs | 4 +++- crates/llm/src/openai_chat.rs | 4 +++- 6 files changed, 23 insertions(+), 5 deletions(-) diff --git a/crates/llm/AGENTS.md b/crates/llm/AGENTS.md index bfc5afb..b4b9808 100644 --- a/crates/llm/AGENTS.md +++ b/crates/llm/AGENTS.md @@ -56,6 +56,7 @@ Trait-based LLM client implementations for multiple providers. - Core message and tool types defined locally instead of re-exporting from `ollama-rs` - tool calls hold name and arguments via `JsonResult`, preserving unparseable argument strings in the `error` variant - tool info stores name, description, and parameters without wrapper enums + - when converting assistant tool calls into provider requests, the function `arguments` include the original tool-call `id` under the `_id` field - chat messages are an enum of `UserMessage`, `AssistantMessage`, `SystemMessage`, and `ToolMessage`, each with only relevant fields - `AssistantMessage` holds a `Vec` for text, tool calls, and thinking segments - tool calls include an `id` string, assigned locally when missing diff --git a/crates/llm/src/gemini_rust.rs b/crates/llm/src/gemini_rust.rs index fc7edda..14b3f7b 100644 --- a/crates/llm/src/gemini_rust.rs +++ b/crates/llm/src/gemini_rust.rs @@ -74,7 +74,7 @@ impl LlmClient for GeminiRustClient { } AssistantPart::ToolCall(tc) => { let args = match tc.arguments { - JsonResult::Content { content } => content, + JsonResult::Content { .. } => tc.arguments_content_with_id(), JsonResult::Error { .. } => Value::Null, }; parts_vec.push(Part::FunctionCall { diff --git a/crates/llm/src/harmony.rs b/crates/llm/src/harmony.rs index e418b5b..33c310e 100644 --- a/crates/llm/src/harmony.rs +++ b/crates/llm/src/harmony.rs @@ -134,7 +134,9 @@ fn build_prompt( } AssistantPart::ToolCall(tc) => { let args = match &tc.arguments { - JsonResult::Content { content } => content.to_string(), + JsonResult::Content { .. } => { + tc.arguments_content_with_id().to_string() + } JsonResult::Error { error } => error.clone(), }; convo_msgs.push( @@ -474,7 +476,7 @@ mod tests { ), ]); assert!(prefill_tokens.is_none()); - let args = json!({"a": 2, "b": 2}).to_string(); + let args = json!({"a": 2, "b": 2, "_id": "1"}).to_string(); let result = json!({"sum": 4}).to_string(); let expected_tail = format!( concat!( diff --git a/crates/llm/src/lib.rs b/crates/llm/src/lib.rs index 7872267..379ab50 100644 --- a/crates/llm/src/lib.rs +++ b/crates/llm/src/lib.rs @@ -92,6 +92,17 @@ pub struct ToolCall { pub arguments: JsonResult, } +impl ToolCall { + pub fn arguments_content_with_id(&self) -> Value { + let mut result = (*self.arguments.as_result().unwrap()).clone(); + result + .as_object_mut() + .unwrap() + .insert("_id".into(), Value::String(self.id.clone())); + result + } +} + #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub enum AssistantPart { diff --git a/crates/llm/src/ollama.rs b/crates/llm/src/ollama.rs index 9014b27..aed7006 100644 --- a/crates/llm/src/ollama.rs +++ b/crates/llm/src/ollama.rs @@ -62,7 +62,9 @@ impl LlmClient for OllamaClient { } AssistantPart::ToolCall(tc) => { let args = match tc.arguments { - JsonResult::Content { content } => content, + JsonResult::Content { .. } => { + tc.arguments_content_with_id() + } JsonResult::Error { .. } => Value::Null, }; msg.tool_calls.push(OllamaToolCall { diff --git a/crates/llm/src/openai_chat.rs b/crates/llm/src/openai_chat.rs index cf34239..cfb8494 100644 --- a/crates/llm/src/openai_chat.rs +++ b/crates/llm/src/openai_chat.rs @@ -92,7 +92,9 @@ impl LlmClient for OpenAiChatClient { .into_iter() .map(|tc| { let args = match &tc.arguments { - JsonResult::Content { content } => content.to_string(), + JsonResult::Content { .. } => { + tc.arguments_content_with_id().to_string() + } JsonResult::Error { error } => error.clone(), }; ChatCompletionMessageToolCall { From bd813c3bea14884becf8215e1011f9071641d1de Mon Sep 17 00:00:00 2001 From: dstoc <539597+dstoc@users.noreply.github.com> Date: Fri, 12 Sep 2025 07:09:30 +0000 Subject: [PATCH 2/2] feat(llment): add a new builtin function, discard_function_response --- crates/llment/AGENTS.md | 6 +++ .../prompts/snippets/instructions/context.md | 12 +++++ crates/llment/prompts/system/default.md | 1 + crates/llment/src/builtins.rs | 52 ++++++++++++++++--- 4 files changed, 64 insertions(+), 7 deletions(-) create mode 100644 crates/llment/prompts/snippets/instructions/context.md diff --git a/crates/llment/AGENTS.md b/crates/llment/AGENTS.md index 44ac830..1485be5 100644 --- a/crates/llment/AGENTS.md +++ b/crates/llment/AGENTS.md @@ -163,6 +163,12 @@ Basic terminal chat interface scaffold using a bespoke component framework built - `ChatMessageRequest` includes MCP `tool_infos` before enabling thinking - Built-in tools registered via `setup_builtin_tools` - `setup_builtin_tools` returns a running `McpService` inserted into the shared `McpContext` + - Tools: + - `chat.get_message_count`: returns the number of chat messages + - `chat.discard_function_response`: + - parameters: `{ id: string }` + - finds a `Tool` message by matching `id` and clears its result text + - returns `"ok"` on success or a not-found message if no matching entry exists - `McpContext` retains running service handles - `get_message_count` returns the number of chat messages diff --git a/crates/llment/prompts/snippets/instructions/context.md b/crates/llment/prompts/snippets/instructions/context.md new file mode 100644 index 0000000..07f633a --- /dev/null +++ b/crates/llment/prompts/snippets/instructions/context.md @@ -0,0 +1,12 @@ +{% if tool_enabled("chat.discard_function_response") %} +## Context + +In this environment there is limited history and context size available. + +It's important that you remove function responses (FR) that are no longer necessary by calling the chat. +discard_function_response. + +Summarize the necessary parts of the FR with chain-of-thought, then proactively discard the FR as soon as possible -- before proceeding with other function calls. If the contents of the FR needs to be part of a message to the user, wait for a subsequent round before discarding. + +Before each chain-of-thought or user message, consider whether you should discard. +{% endif %} diff --git a/crates/llment/prompts/system/default.md b/crates/llment/prompts/system/default.md index d9426f2..ff01b96 100644 --- a/crates/llment/prompts/system/default.md +++ b/crates/llment/prompts/system/default.md @@ -2,6 +2,7 @@ {% include "snippets/instructions/focus.md" %} {% include "snippets/instructions/task.md" %} {% include "snippets/instructions/autonomy.md" %} +{% include "snippets/instructions/context.md" %} {% for file in glob("snippets/env/*.md") %} {% include file %} {% endfor %} diff --git a/crates/llment/src/builtins.rs b/crates/llment/src/builtins.rs index c66fa33..4165f62 100644 --- a/crates/llment/src/builtins.rs +++ b/crates/llment/src/builtins.rs @@ -1,21 +1,28 @@ use std::sync::{Arc, Mutex}; use arc_swap::ArcSwap; -use llm::{ChatMessage, ToolInfo, mcp::McpService}; +use llm::{ChatMessage, JsonResult, ToolInfo, mcp::McpService}; use rmcp::{ ServerHandler, - handler::server::router::tool::ToolRouter, + handler::server::{router::tool::ToolRouter, tool::Parameters}, model::{ServerCapabilities, ServerInfo}, service::{RoleClient, RunningService, ServiceExt}, tool, tool_handler, tool_router, }; use schemars::{JsonSchema, schema_for}; use serde::{Deserialize, Serialize}; +use serde_json::Value; use tokio::io::duplex; #[derive(Serialize, Deserialize, JsonSchema)] pub struct GetMessageCountParams {} +#[derive(Serialize, Deserialize, JsonSchema)] +pub struct DiscardFunctionResponseParams { + /// The id of the ToolCall/Tool response to discard + pub id: String, +} + #[derive(Clone)] struct BuiltinTools { chat_history: Arc>>, @@ -38,6 +45,30 @@ impl BuiltinTools { fn get_message_count(&self) -> String { self.chat_history.lock().unwrap().len().to_string() } + + #[tool( + name = "discard_function_response", + description = "Removes the content from a tool response in history by id" + )] + fn discard_function_response( + &self, + Parameters(params): Parameters, + ) -> String { + let mut history = self.chat_history.lock().unwrap(); + if let Some((idx, _)) = history.iter().enumerate().rev().find(|(_, m)| match m { + ChatMessage::Tool(t) => t.id == params.id, + _ => false, + }) { + if let ChatMessage::Tool(t) = &mut history[idx] { + t.content = JsonResult::Content { + content: Value::String("".into()), + }; + } + "ok".into() + } else { + format!("Tool response with id '{}' not found", params.id) + } + } } #[tool_handler(router = self.tool_router)] @@ -59,11 +90,18 @@ pub async fn setup_builtin_tools( builtins.clone().serve(server_transport), McpService { prefix: "chat".into(), - tools: ArcSwap::new(Arc::new(vec![ToolInfo { - name: "get_message_count".into(), - description: "Returns the number of chat messages".into(), - parameters: schema_for!(GetMessageCountParams), - }])), + tools: ArcSwap::new(Arc::new(vec![ + ToolInfo { + name: "get_message_count".into(), + description: "Returns the number of chat messages".into(), + parameters: schema_for!(GetMessageCountParams), + }, + ToolInfo { + name: "discard_function_response".into(), + description: "Removes the content from a tool response in history by id".into(), + parameters: schema_for!(DiscardFunctionResponseParams), + }, + ])), } .serve(client_transport) );