Tool calling

[zackangelo]

Summary by CodeRabbit

• New Features

  • Pluggable tool support for model sequences, enabling custom tool invocation during generation.
  • New example programs demonstrating basic usage and tool-augmented workflows.
  • Package bumped to v0.2.0 and default feature now enables the client.

• Breaking Changes

  • Python API: removed tools_enabled parameter from open().
  • Stream behavior: generation streams now surface errors as results for improved error handling.

[coderabbitai]

Walkthrough

Adds a pluggable Toolbox for tool registration/invocation, integrates tools into sequence open and event handling, refactors GenStream to carry Result items, adds example programs and utilities demonstrating tool usage, updates Cargo.toml with dev-dependencies, and removes the Python tools_enabled parameter.

Changes

Cohort / File(s)	Summary
Manifest & Dependencies Cargo.toml	Package version bumped to 0.2.0; default features set to ["client"]; tokio dependency simplified to features = ["full"], optional = true; new [dev-dependencies] added (anyhow, clap, colorize, futures, meval, tokio (dev), tracing-subscriber).
Tool Support Module src/client/tools.rs	New public toolbox implementation: Toolbox, Tool trait, ToolDefinition/parameters/property/schema types with serde, call_tools(), and tool_def_prompt(); includes tests for serialization.
Client Module Integration src/client/mod.rs	Exposes pub mod tools; replaces OpenOpts.tools_enabled: bool with toolbox: Option<Toolbox>; added SeqClosed error variant; ModelSocketError::Other uses #[from] anyhow::Error; opening logic derives tools availability from toolbox and may append tool definition prompt.
Sequence & Stream Changes src/client/seq.rs	Seq gains toolbox: Arc<Option<Mutex<Toolbox>>> and constructor updated; event handling extended with SeqToolCall → on_tool_call; GenStream now yields Result<GenChunk, ModelSocketError>; close/fork logic updated to drain/send SeqClosed to pending senders; text helpers unwrap new Result type.
Python Bindings src/python.rs	Removed tools_enabled parameter from PyBlockingModelSocketClient.open signature and PyO3 metadata; PyBlockingGenStream.__next__ changed to use stream.next().await with error mapping and imports adjusted for futures::StreamExt.
Examples & Utilities examples/simple.rs, examples/tool_call.rs, examples/utils/mod.rs	New examples/simple.rs (basic ModelSocket usage). New examples/tool_call.rs defines WeatherTool and CalculatorTool, builds a Toolbox, and demonstrates tool-augmented chat. New examples/utils/mod.rs adds pub async fn print_stream(GenStream, hidden, colors) to print stream events with optional filtering and ANSI colors.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant App as Application
    participant Seq as Sequence
    participant Toolbox
    participant Model as ModelSocket

    User->>App: register tools (add_tool)
    App->>Toolbox: add_tool(Weather, Calculator)
    App->>Seq: open(model, toolbox)
    Seq->>Model: SeqOpenReq (tools_enabled=true)
    Model-->>Seq: SeqOpenRes
    Seq->>Seq: append(tool_def_prompt)

    User->>Seq: append_message(user_query)
    User->>Seq: generate()
    Seq->>Model: generation request
    Model-->>Seq: SeqToolCall event
    Seq->>Toolbox: call_tools(SeqToolCall[])
    Toolbox-->>Seq: ToolResult[]
    Seq->>Model: SeqCommand::ToolReturn (with results)
    Model-->>Seq: GenChunk(s)
    Seq-->>App: GenStream emits Result<GenChunk, ModelSocketError>
    App->>App: print_stream(stream)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Areas needing extra attention:

• GenStream type change to Result and all unwrap/propagation sites.
• Seq::on_tool_call and toolbox invocation/error handling (closing seq on failures).
• Correct sharing of toolbox via Arc<Option<Mutex>> across forks and seq lifecycle.
• Integration point that appends tool definition prompt after open (timing/format).
• Python binding changes ensuring iteration error mapping works correctly.

Poem

🐰
I hopped a prompt and found a kit,
Tools gathered neatly, each a wit,
Weather whispers, numbers play,
Streams now carry wins and fray,
Sequences hum — tools lead the skit! ✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.35% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title "Tool calling" directly reflects the main objective of adding tool calling functionality throughout the codebase, including new tools module, toolbox integration, and example implementations.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch tools

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/python.rs (1)

88-107: Python open text_signature advertises removed tools_enabled parameter

The text_signature still includes tools_enabled=False, but the actual function signature removed this parameter:

#[pyo3(
    text_signature = "($self, model, /, *, tools_enabled=False, tool_prompt=None, skip_prelude=False)",
    signature = (model, tool_prompt=None, skip_prelude=None)
)]
pub fn open(
    &self,
    model: &str,
    // tools_enabled: Option<bool>,
    tool_prompt: Option<&str>,
    skip_prelude: Option<bool>,
) -> PyResult<PyBlockingSeq> {

Update text_signature to match and remove commented-out tools_enabled code:

-    text_signature = "($self, model, /, *, tools_enabled=False, tool_prompt=None, skip_prelude=False)",
+    text_signature = "($self, model, /, *, tool_prompt=None, skip_prelude=False)",
     signature = (model, tool_prompt=None, skip_prelude=None)
 )]
 pub fn open(
     &self,
     model: &str,
-    // tools_enabled: Option<bool>,
     tool_prompt: Option<&str>,
     skip_prelude: Option<bool>,
 ) -> PyResult<PyBlockingSeq> {
     let client = self.inner.clone();
 
     let seq = Python::attach(|py| {
         block_on(py, async move {
             let mut opts = OpenOpts::default();
-            // opts.tools_enabled = tools_enabled.unwrap_or(false);
             opts.tool_prompt = tool_prompt.map(|s| s.to_string());
             opts.skip_prelude = skip_prelude.unwrap_or(false);

🧹 Nitpick comments (11)

examples/utils/mod.rs (1)

6-35: Streaming + error handling look correct

print_stream correctly handles the Result<GenChunk, ModelSocketError> stream, respects the hidden/colors flags, and propagates errors via anyhow::Result. Per-event flushing is appropriate for a demo; if this ever gets reused in high‑volume contexts, consider buffering or less frequent flushes.

Cargo.toml (1)

12-26: Confirm making client a default feature is intended

Setting:

[features]
default = ["client"]

means downstream crates will now pull in the async client stack (tokio, tungstenite, etc.) unless they opt out of default features. If the goal is to keep modelsocket lightweight as a library dependency, you may want to leave default = [] and ask users to enable client explicitly. Otherwise, this change is fine but is a notable public behavior shift.

Also applies to: 37-37, 50-57

src/client/mod.rs (1)

2-8: Toolbox integration and error surface look solid, with a small edge-case to consider

• Exporting pub mod tools; and threading Toolbox through OpenOpts and open() into Seq::new is clean and keeps tool support opt‑in at the call site.
• ModelSocketError::SeqClosed and Other(#[from] anyhow::Error) give you a clear way to signal local shutdown and wrap toolbox/meval/etc. errors.

One edge case: if appending the tool_def_prompt as a system message fails:

if let Some(tool_def_prompt) = tool_def_prompt {
    seq.append(tool_def_prompt, AppendOpts::system()).await?;
}

open() will return an error after:

• the sequence has already been created, and
• it has been registered in self.seqs.

In practice the server will likely close the sequence on the same underlying failure, and on_seq_closed will clean it up, but you might want to either:

• treat the tool-definition append as best-effort (log but don’t fail open()), or
• explicitly remove the seq from self.seqs before returning Err here.

The new GenOpts/AppendOpts helpers (assistant(), user(), system()) are straightforward and make call sites much nicer.

Also applies to: 19-50, 177-231, 321-335

src/client/seq.rs (2)

62-87: Tool call handling is correct; consider remote-close behavior

on_tool_call correctly:

• no-ops (with a debug log) if no toolbox is configured,
• delegates to toolbox.call_tools(tool_calls).await,
• sends a SeqCommand::ToolReturn with the collected results and default gen_opts,
• and, on error, logs and tries to close the sequence.

That’s a solid first pass for tool integration.

Separately, on_close (used when the server sends MSEvent::SeqClosed) still just clears cmds/gen_streams:

let mut cmds = self.cmds.lock().await;
cmds.clear();

let mut gen_streams = self.gen_streams.lock().await;
gen_streams.clear();

so callers waiting on an in‑flight append/generate don’t receive a SeqClosed error when the server initiates the close. You’ve already introduced ModelSocketError::SeqClosed and are using it in close(), so you might want to mirror that here by sending Err(ModelSocketError::SeqClosed) to any remaining senders before clearing the maps.

Also applies to: 159-195

---

274-343: Good use of SeqClosed in close(); no obvious race issues

The new close() implementation:

• drains all outstanding cmds and gen_streams, sending Err(ModelSocketError::SeqClosed) to each,
• then re‑inserts a fresh cid/sender pair for the close command,
• sends SeqCommand::Close and waits for the confirmation via that channel.

This avoids leaving waiting callers hanging and is structured to avoid holding the cmds lock while doing network I/O. The TODO about marking the sequence as closed to prevent future commands is a reasonable follow‑up, but the current behavior is sound.

src/client/tools.rs (6)

23-34: Consider handling duplicate tool names in add_tool

add_tool will silently overwrite any existing tool with the same definition.name. That may be surprising if multiple modules register tools and accidentally collide on names. Consider either:

• returning a Result / Option<Box<dyn Tool>> to signal replacement, or
• logging a warning when an entry is replaced.

Also, the + 'static bound on T is redundant given the Tool trait already has 'static in its bounds.

---

36-48: Double‑check call_tools error and execution semantics

call_tools executes tools sequentially and short‑circuits on the first error (? on each call). That’s perfectly valid, but it means:

• later tools are not called if an earlier one fails,
• callers get only the first error, not a per‑tool result set with mixed successes/failures.

If the intended contract is “fire a batch and always get back one result per requested call,” you might instead want to collect per‑tool Results or include error info in ToolResult. If short‑circuiting is desired, documenting that behavior in the public API would help.

---

50-59: Improve diagnostics when tools are missing

The "tool {} not found" error and debug!(tool = name, "tool call") log are useful, but you might further improve debuggability by:

• including the list of available tool names in the error or log (especially useful when a name is slightly off),
• or at least distinguishing “lookup failure” logs from successful call logs.

Not required, but worth considering if tool registration is dynamic or user‑driven.

---

61-67: Avoid panic in tool_def_prompt and consider deterministic ordering

Two minor points here:

1. serde_json::to_string_pretty(&tool.parameters).unwrap() will panic on serialization failure. Given the types are under your control this should “never” fail, but panics from library code in a client API can still be surprising. Consider:

   • returning Result<String> from tool_def_prompt and propagating the error, or
   • using .expect("failed to serialize tool parameters") for a clearer panic message if you truly consider it impossible.

2. tool_def_prompt(&self) iterates self.tools.values() from a HashMap, so the ordering of tool definitions in the final prompt is non‑deterministic. If reproducible prompts matter (e.g., for tests or predictable UX), switching to BTreeMap or sorting by name before join would give stable ordering.

Also applies to: 70-77

---

79-82: Tool trait bounds are strong; may want a more ergonomic Future type

Requiring Tool: Debug + Send + Sync + 'static is reasonable given the toolbox is shared across threads. The call signature, however, effectively requires implementors to return a 'static Future (because of the trait‑object dyn Future), which prevents capturing &self directly in an async block.

If you want to make implementors’ lives easier while still avoiding async_trait, consider a boxed future alias with an explicit lifetime, e.g.:

pub type ToolFuture<'a> = Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>>;

pub trait Tool: fmt::Debug + Send + Sync + 'static {
    fn definition(&self) -> ToolDefinition;
    fn call(&self, args: &str) -> ToolFuture<'_>;
}

This allows capturing &self safely inside call.

---

132-181: Good serialization test; consider adding a prompt/toolbox test later

The tool_definition_serializes_expected_parameters test validates the JSON shape (type, required, property description, and enum), which is exactly where regressions are likely. As a possible follow‑up, you could add:

• a small test around tool_def_prompt to lock down formatting, and/or
• a smoke test for Toolbox::call_tools using a dummy tool.

Not necessary for this PR, but would increase confidence as the toolbox evolves.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between f668a2f and 8911e9c.

⛔ Files ignored due to path filters (1)

• Cargo.lock is excluded by !**/*.lock

📒 Files selected for processing (8)

• Cargo.toml (3 hunks)
• examples/simple.rs (1 hunks)
• examples/tool_call.rs (1 hunks)
• examples/utils/mod.rs (1 hunks)
• src/client/mod.rs (6 hunks)
• src/client/seq.rs (12 hunks)
• src/client/tools.rs (1 hunks)
• src/python.rs (3 hunks)

🧰 Additional context used

🧬 Code graph analysis (7)

examples/utils/mod.rs (2)

src/python.rs (1)

• modelsocket (413-419)

src/client/seq.rs (1)

• text (359-368)

examples/simple.rs (2)

src/client/mod.rs (5)

• connect (83-86)
• user (268-273)
• user (308-312)
• assistant (261-266)
• assistant (302-306)

examples/utils/mod.rs (1)

• print_stream (6-36)

src/client/mod.rs (2)

src/client/tools.rs (5)

• tools (18-18)
• tools (62-65)
• tool_def_prompt (61-67)
• tool_def_prompt (70-77)
• new (24-28)

src/client/seq.rs (1)

• new (32-48)

examples/tool_call.rs (4)

src/client/tools.rs (8)

• tools (18-18)
• tools (62-65)
• fmt (16-20)
• default (102-108)
• new (24-28)
• definition (80-80)
• call (50-59)
• call (81-81)

src/client/mod.rs (2)

• connect (83-86)
• new (60-81)

src/client/seq.rs (1)

• new (32-48)

examples/utils/mod.rs (1)

• print_stream (6-36)

src/client/seq.rs (1)

src/protocol.rs (2)

• cid (94-106)
• seq_id (108-120)

src/python.rs (1)

src/client/mod.rs (1)

• open (177-231)

src/client/tools.rs (2)

src/client/seq.rs (1)

• new (32-48)

examples/tool_call.rs (5)

• definition (66-85)
• definition (96-114)
• call (87-89)
• call (116-131)
• serde_json (119-119)

🔇 Additional comments (6)

src/python.rs (1)

319-344: GenStream iteration and error mapping for Python look correct

The updated __next__ correctly:

• uses stream.next().await on GenStream<Item = Result<GenChunk, ModelSocketError>>,
• maps any Err(ModelSocketError) into a PyErr via map_err,
• and preserves the stream between iterations by storing it back on self.

This keeps Python streaming aligned with the new Rust GenStream semantics.

src/client/seq.rs (1)

20-48: GenStream Result wrapping and sequence lifecycle look consistent

• Storing gen_streams as Sender<Result<GenChunk, ModelSocketError>> and updating GenStream to expose Item = Result<GenChunk, ModelSocketError> is consistent end‑to‑end:
  • on_text sends Ok(GenChunk),
  • generate wires up the mpsc channel correctly,
  • GenStream::text / text_and_tokens unwrap with chunk? and filter on hidden.
• Forking passes self.toolbox.clone() into the child Seq::new, so tools remain available after fork().
• close now notifies all outstanding commands and generation streams with Err(ModelSocketError::SeqClosed) before issuing the close command, which gives callers a clear signal instead of hanging.

This all looks correct and matches the new error surface.

Also applies to: 139-157, 247-272, 346-383

src/client/tools.rs (4)

11-13: Toolbox core structure looks solid

Storing tools as HashMap<String, Box<dyn Tool>> keyed by their definition name is a straightforward and flexible design, and it composes well with the Tool trait object.

---

15-21: Nice Debug impl focusing on tool names

The custom Debug implementation that only exposes the tool keys (names) is a good balance between debuggability and not dumping full tool internals.

---

84-109: Tool metadata structs line up well with JSON Schema‑style usage

ToolDefinition, ToolParameters, and the Default impl for ToolParameters map cleanly to an object‑typed parameter schema, and the serde attributes (rename = "type", skip_serializing_if) look correct for the intended JSON layout. This should work nicely with the serde_json::from_value patterns used in the examples.

---

111-121: Property and schema‑type design is flexible and matches typical tool schemas

ToolProperty with optional description and enum (via allowed_values) plus ToolSchemaType’s snake_case serialization should cover most tool parameter shapes you’ll need, while staying simple. This mirrors common LLM tool schema conventions well.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

CLI hidden/colors flags are ignored

You parse hidden and colors on Args but always call:

utils::print_stream(stream, true, true).await?;

This ignores the user’s choices. Consider:

-    utils::print_stream(stream, true, true).await?;
+    utils::print_stream(stream, args.hidden, args.colors).await?;

so the CLI options actually affect output.

Also applies to: 32-41

🤖 Prompt for AI Agents

In examples/simple.rs around lines 7-17 (and also apply same fix at lines
~32-41): the CLI boolean flags are defined with duplicate attribute tokens and
their parsed values are never used because print_stream is called with hardcoded
true/true. Fix by removing the duplicated 'long' attributes (keep a single
#[clap(long, default_value_t = false)] for hidden and #[clap(long,
default_value_t = true)] for colors) so clap parses them correctly, and change
the print_stream calls to pass the parsed args.hidden and args.colors instead of
hardcoded true, true; apply the same replacement for the other call site at
lines 32-41.

[coderabbitai]

Actionable comments posted: 0

♻️ Duplicate comments (2)

examples/simple.rs (1)

13-16: Remove duplicated long in clap attributes

hidden and colors use #[clap(long, long, ...)]; the second long is redundant and may cause the derive macro to error. Keep a single long:

-    #[clap(long, long, default_value_t = false)]
+    #[clap(long, default_value_t = false)]
-    #[clap(long, long, default_value_t = true)]
+    #[clap(long, default_value_t = true)]

examples/tool_call.rs (1)

20-23: Fix duplicated long in clap attributes

Same as in examples/simple.rs, hidden and colors use #[clap(long, long, ...)]; drop the extra long:

-    #[clap(long, long, default_value_t = false)]
+    #[clap(long, default_value_t = false)]
-    #[clap(long, long, default_value_t = true)]
+    #[clap(long, default_value_t = true)]

🧹 Nitpick comments (2)

Cargo.toml (1)

3-57: Confirm default client feature and tokio setup are intentional

Making client the default feature and keeping tokio optional while also adding a dev-only tokio is a meaningful API/dependency change for crate consumers (they now pull in client + async stack by default). Please confirm this is the intended public behavior and that you don’t want a lighter default profile.

examples/tool_call.rs (1)

95-105: Mark expression as required in CalculatorTool schema

WeatherTool declares "required": ["location"], but CalculatorTool doesn’t mark expression as required, even though call will error if it’s missing. Consider:

         let parameters = json!({
             "type": "object",
             "properties": {
                 "expression": {
                     "type": "string",
                     "description": "A mathematical expression to evaluate (e.g. 1 + 1)"
                 }
             },
+            "required": ["expression"]
         });

to keep the schema consistent and reduce tool-call failures.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8911e9c and d6f4322.

⛔ Files ignored due to path filters (1)

• Cargo.lock is excluded by !**/*.lock

📒 Files selected for processing (3)

• Cargo.toml (4 hunks)
• examples/simple.rs (1 hunks)
• examples/tool_call.rs (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (2)

examples/simple.rs (2)

src/client/mod.rs (5)

• connect (83-86)
• user (268-273)
• user (308-312)
• assistant (261-266)
• assistant (302-306)

examples/utils/mod.rs (1)

• print_stream (6-36)

examples/tool_call.rs (3)

src/client/tools.rs (7)

• tools (18-18)
• tools (62-65)
• default (102-108)
• new (24-28)
• definition (80-80)
• call (50-59)
• call (81-81)

src/client/seq.rs (1)

• new (32-48)

examples/utils/mod.rs (1)

• print_stream (6-36)

🔇 Additional comments (2)

examples/simple.rs (1)

19-41: Example flow and CLI usage look solid

Connection, sequence open, append, generate, and stream printing are wired cleanly, and print_stream now respects hidden/colors from CLI.

examples/tool_call.rs (1)

26-59: Toolbox wiring and tool implementations look good

Toolbox creation, tool registration, sequence open with OpenOpts, and the end-to-end tool-call example (weather + calculator) are clear and aligned with the client/tooling API.