use alloc::{
borrow::ToOwned,
boxed::Box,
format,
string::{String, ToString},
sync::Arc,
};
use core::any::{Any, TypeId};
use serde::{Deserialize, Serialize};
use crate::compat::{HashMap, OnceLock, RwLock, read_lock, write_lock};
#[derive(Clone)]
pub struct SharedState(Arc<RwLock<HashMap<String, serde_json::Value>>>);
impl SharedState {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn get_state(&self, key: &str, default: serde_json::Value) -> serde_json::Value {
match read_lock(&self.0) {
Ok(guard) => guard.get(key).cloned().unwrap_or(default),
Err(e) => {
tracing::warn!(key, error = %e, "SharedState::get_state: lock poisoned, returning default");
default
}
}
}
pub fn set_state(&self, key: &str, value: serde_json::Value) -> Result<(), ToolError> {
match write_lock(&self.0) {
Ok(mut guard) => {
guard.insert(key.to_owned(), value);
Ok(())
}
Err(e) => {
let msg = format!("SharedState::set_state: lock poisoned for key '{key}': {e}");
tracing::warn!("{msg}");
Err(ToolError::new(msg))
}
}
}
#[must_use]
pub fn remove_state(&self, key: &str) -> bool {
match write_lock(&self.0) {
Ok(mut guard) => guard.remove(key).is_some(),
Err(e) => {
tracing::warn!(key, error = %e, "SharedState::remove_state: lock poisoned");
false
}
}
}
pub fn clear_state(&self) {
match write_lock(&self.0) {
Ok(mut guard) => guard.clear(),
Err(e) => {
tracing::warn!(error = %e, "SharedState::clear_state: lock poisoned");
}
}
}
}
impl Default for SharedState {
fn default() -> Self {
Self(Arc::new(RwLock::new(HashMap::new())))
}
}
impl core::fmt::Debug for SharedState {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("SharedState").finish_non_exhaustive()
}
}
pub struct ToolContext {
conversation_id: Option<String>,
state: SharedState,
extensions: Arc<RwLock<HashMap<TypeId, Box<dyn Any + Send + Sync>>>>,
}
impl Clone for ToolContext {
fn clone(&self) -> Self {
Self {
conversation_id: self.conversation_id.clone(),
state: self.state.clone(),
extensions: Arc::clone(&self.extensions),
}
}
}
impl core::fmt::Debug for ToolContext {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let count = match read_lock(&self.extensions) {
Ok(g) => g.len(),
Err(e) => {
tracing::warn!(error = %e, "ToolContext::fmt: extensions lock poisoned");
0
}
};
f.debug_struct("ToolContext")
.field("conversation_id", &self.conversation_id)
.field("state", &self.state)
.field("extensions_count", &count)
.finish()
}
}
impl Default for ToolContext {
fn default() -> Self {
Self::new()
}
}
impl ToolContext {
#[must_use]
pub fn new() -> Self {
Self {
conversation_id: None,
state: SharedState::new(),
extensions: Arc::new(RwLock::new(HashMap::new())),
}
}
#[must_use]
pub fn with_conversation_id(mut self, conversation_id: impl Into<String>) -> Self {
self.conversation_id = Some(conversation_id.into());
self
}
#[must_use]
pub fn with_caller(&self, conversation_id: impl Into<String>) -> Self {
Self {
conversation_id: Some(conversation_id.into()),
state: self.state.clone(),
extensions: Arc::clone(&self.extensions),
}
}
#[must_use]
pub fn with_shared_state(mut self, state: SharedState) -> Self {
self.state = state;
self
}
#[must_use]
pub fn shared_state(&self) -> SharedState {
self.state.clone()
}
#[must_use]
pub fn conversation_id(&self) -> Option<&str> {
self.conversation_id.as_deref()
}
#[must_use]
pub fn get_state(&self, key: &str, default: serde_json::Value) -> serde_json::Value {
match read_lock(&self.state.0) {
Ok(guard) => guard.get(key).cloned().unwrap_or(default),
Err(e) => {
tracing::warn!(key, error = %e, "ToolContext::get_state: lock poisoned, returning default");
default
}
}
}
pub fn set_state(&self, key: &str, value: serde_json::Value) -> Result<(), ToolError> {
match write_lock(&self.state.0) {
Ok(mut guard) => {
guard.insert(key.to_owned(), value);
Ok(())
}
Err(e) => {
let msg = format!("ToolContext::set_state: lock poisoned for key '{key}': {e}");
tracing::warn!("{msg}");
Err(ToolError::new(msg))
}
}
}
pub fn set_ext<T: Send + Sync + 'static>(&self, value: T) -> Result<(), ToolError> {
match write_lock(&self.extensions) {
Ok(mut exts) => {
exts.insert(TypeId::of::<T>(), Box::new(value));
Ok(())
}
Err(e) => {
let msg = format!(
"ToolContext::set_ext: lock poisoned for type '{}': {e}",
core::any::type_name::<T>()
);
tracing::warn!("{msg}");
Err(ToolError::new(msg))
}
}
}
#[must_use]
pub fn get_ext<T: Clone + Send + Sync + 'static>(&self) -> Option<T> {
match read_lock(&self.extensions) {
Ok(exts) => exts
.get(&TypeId::of::<T>())
.and_then(|v| v.downcast_ref::<T>())
.cloned(),
Err(e) => {
tracing::warn!(error = %e, "ToolContext::get_ext: lock poisoned, returning None");
None
}
}
}
}
pub use llm_tool_macros::{llm_prompt, llm_resource, llm_tool};
pub use schemars::JsonSchema;
const fn other_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "bool",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolOutput {
pub(crate) content: String,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub(crate) metadata: HashMap<String, serde_json::Value>,
}
impl ToolOutput {
pub fn new(content: impl Into<String>) -> Self {
Self {
content: content.into(),
metadata: HashMap::new(),
}
}
pub fn json<T: serde::Serialize>(value: &T) -> Result<Self, ToolError> {
serde_json::to_string(value)
.map(Self::new)
.map_err(|e| ToolError::new(format!("serialization failed: {e}")))
}
pub fn from_metadata<T: serde::Serialize>(value: &T) -> Result<Self, ToolError> {
let json_value = serde_json::to_value(value)
.map_err(|e| ToolError::new(format!("metadata serialization failed: {e}")))?;
let content = json_value.to_string();
match json_value {
serde_json::Value::Object(map) => Ok(Self {
content,
metadata: map.into_iter().collect(),
}),
other => Err(ToolError::new(format!(
"metadata must serialize to a JSON object, got {}",
other_type_name(&other),
))),
}
}
#[must_use]
pub fn with_meta(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
self.metadata.insert(key.into(), value);
self
}
pub fn with_metadata<T: serde::Serialize>(mut self, value: &T) -> Result<Self, ToolError> {
let json = serde_json::to_value(value)
.map_err(|e| ToolError::new(format!("metadata serialization failed: {e}")))?;
match json {
serde_json::Value::Object(map) => {
self.metadata.extend(map);
Ok(self)
}
other => Err(ToolError::new(format!(
"metadata must serialize to a JSON object, got {}",
other_type_name(&other),
))),
}
}
#[must_use]
pub fn content(&self) -> &str {
&self.content
}
#[must_use]
pub fn into_content(self) -> String {
self.content
}
#[must_use]
pub const fn metadata(&self) -> &HashMap<String, serde_json::Value> {
&self.metadata
}
#[must_use]
pub fn truncated(mut self, policy: &TruncationPolicy) -> Self {
truncation::truncate_output(&mut self, policy);
self
}
pub fn truncate_detailed(
&mut self,
policy: &TruncationPolicy,
) -> Result<TruncationOutcome<'static>, SpillError> {
truncation::truncate_output_detailed(self, policy)
}
pub fn truncate_detailed_with_context(
&mut self,
policy: &TruncationPolicy,
ctx: &SpillContext<'_>,
) -> Result<TruncationOutcome<'static>, SpillError> {
truncation::truncate_output_detailed_with_context(self, policy, ctx)
}
}
impl core::fmt::Display for ToolOutput {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.write_str(&self.content)
}
}
impl From<String> for ToolOutput {
fn from(content: String) -> Self {
Self::new(content)
}
}
impl From<&str> for ToolOutput {
fn from(content: &str) -> Self {
Self::new(content)
}
}
impl From<i64> for ToolOutput {
fn from(value: i64) -> Self {
Self::new(value.to_string())
}
}
impl From<f64> for ToolOutput {
fn from(value: f64) -> Self {
Self::new(value.to_string())
}
}
impl From<bool> for ToolOutput {
fn from(value: bool) -> Self {
Self::new(value.to_string())
}
}
impl From<serde_json::Value> for ToolOutput {
fn from(value: serde_json::Value) -> Self {
Self::new(value.to_string())
}
}
pub struct Json<T>(pub T);
impl<T: serde::Serialize> From<Json<T>> for ToolOutput {
fn from(json: Json<T>) -> Self {
let json_value = serde_json::to_value(&json.0)
.expect("Json<T> serialization failed — this is a bug in the Serialize impl");
let content = json_value.to_string();
match json_value {
serde_json::Value::Object(map) => Self {
content,
metadata: map.into_iter().collect(),
},
_ => Self::new(content),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum RegistryItem {
Tool,
Prompt,
Resource,
}
impl core::fmt::Display for RegistryItem {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Tool => write!(f, "tool"),
Self::Prompt => write!(f, "prompt"),
Self::Resource => write!(f, "resource"),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolError {
pub message: String,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
metadata: HashMap<String, serde_json::Value>,
}
impl ToolError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
metadata: HashMap::new(),
}
}
pub const ERROR_KIND_KEY: &'static str = "error_kind";
pub const KIND_NOT_REGISTERED: &'static str = "not_registered";
pub const KIND_DESTRUCTIVE_BATCH: &'static str = "destructive_in_parallel_batch";
#[must_use]
pub fn not_found(kind: RegistryItem, name: &str) -> Self {
Self::new(format!("no {kind} named '{name}' is registered")).with_meta(
Self::ERROR_KIND_KEY,
serde_json::json!(Self::KIND_NOT_REGISTERED),
)
}
#[must_use]
pub fn destructive_in_parallel_batch(tool_name: &str) -> Self {
Self::new(format!(
"Tool '{tool_name}' has Destructive effect and cannot be executed in Parallel batch mode without explicit confirmation"
))
.with_meta(
Self::ERROR_KIND_KEY,
serde_json::json!(Self::KIND_DESTRUCTIVE_BATCH),
)
}
#[must_use]
pub fn with_meta(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
self.metadata.insert(key.into(), value);
self
}
pub fn with_metadata<T: serde::Serialize>(mut self, value: &T) -> Result<Self, Self> {
let json = match serde_json::to_value(value) {
Ok(j) => j,
Err(e) => {
self.message =
format!("{} (metadata serialization also failed: {e})", self.message);
return Err(self);
}
};
match json {
serde_json::Value::Object(map) => {
self.metadata.extend(map);
Ok(self)
}
other => {
self.message = format!(
"{} (metadata must serialize to a JSON object, got {})",
self.message,
other_type_name(&other),
);
Err(self)
}
}
}
#[must_use]
pub const fn metadata(&self) -> &HashMap<String, serde_json::Value> {
&self.metadata
}
#[must_use]
pub fn is_not_found(&self) -> bool {
self.metadata
.get(Self::ERROR_KIND_KEY)
.and_then(serde_json::Value::as_str)
== Some(Self::KIND_NOT_REGISTERED)
}
#[must_use]
pub fn is_destructive_batch_rejected(&self) -> bool {
self.metadata
.get(Self::ERROR_KIND_KEY)
.and_then(serde_json::Value::as_str)
== Some(Self::KIND_DESTRUCTIVE_BATCH)
}
}
impl core::fmt::Display for ToolError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{}", self.message)
}
}
impl core::error::Error for ToolError {}
impl From<String> for ToolError {
fn from(message: String) -> Self {
Self::new(message)
}
}
impl From<&str> for ToolError {
fn from(message: &str) -> Self {
Self::new(message)
}
}
#[cfg(feature = "std")]
impl From<std::io::Error> for ToolError {
fn from(e: std::io::Error) -> Self {
Self::new(e.to_string()).with_meta(
Self::ERROR_KIND_KEY,
serde_json::json!(format!("{:?}", e.kind())),
)
}
}
impl From<serde_json::Error> for ToolError {
fn from(e: serde_json::Error) -> Self {
Self::new(e.to_string())
.with_meta("category", serde_json::json!(format!("{:?}", e.classify())))
}
}
impl From<Box<dyn core::error::Error + Send + Sync>> for ToolError {
fn from(e: Box<dyn core::error::Error + Send + Sync>) -> Self {
Self::new(e.to_string())
}
}
impl From<core::convert::Infallible> for ToolError {
fn from(never: core::convert::Infallible) -> Self {
match never {}
}
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, schemars::JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum ToolEffect {
ReadOnly,
#[default]
Mutating,
Destructive,
}
impl ToolEffect {
#[must_use]
pub const fn is_read_only(self) -> bool {
matches!(self, Self::ReadOnly)
}
#[must_use]
pub const fn is_destructive(self) -> bool {
matches!(self, Self::Destructive)
}
}
#[non_exhaustive]
#[derive(Debug, Serialize)]
pub struct ToolDefinition {
pub name: String,
pub description: String,
pub parameter_schema: serde_json::Value,
pub effect: ToolEffect,
#[serde(skip)]
parameter_schema_json: OnceLock<Arc<str>>,
}
impl Clone for ToolDefinition {
fn clone(&self) -> Self {
Self {
name: self.name.clone(),
description: self.description.clone(),
parameter_schema: self.parameter_schema.clone(),
effect: self.effect,
parameter_schema_json: OnceLock::new(),
}
}
}
impl PartialEq for ToolDefinition {
fn eq(&self, other: &Self) -> bool {
self.name == other.name
&& self.description == other.description
&& self.parameter_schema == other.parameter_schema
&& self.effect == other.effect
}
}
impl Eq for ToolDefinition {}
impl<'de> Deserialize<'de> for ToolDefinition {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
struct RawDef {
name: String,
description: String,
parameter_schema: serde_json::Value,
#[serde(default)]
effect: Option<ToolEffect>,
#[serde(default)]
idempotent: Option<bool>,
}
let raw = RawDef::deserialize(deserializer)?;
let effect = match (raw.effect, raw.idempotent) {
(Some(eff), _) => eff,
(None, Some(true)) => ToolEffect::ReadOnly,
(None, Some(false) | None) => ToolEffect::Mutating,
};
Ok(Self {
name: raw.name,
description: raw.description,
parameter_schema: raw.parameter_schema,
effect,
parameter_schema_json: OnceLock::new(),
})
}
}
impl ToolDefinition {
#[must_use]
pub fn new(
name: impl Into<String>,
description: impl Into<String>,
parameter_schema: serde_json::Value,
) -> Self {
Self::with_effect(name, description, parameter_schema, ToolEffect::Mutating)
}
#[must_use]
pub fn with_effect(
name: impl Into<String>,
description: impl Into<String>,
parameter_schema: serde_json::Value,
effect: ToolEffect,
) -> Self {
Self {
name: name.into(),
description: description.into(),
parameter_schema,
effect,
parameter_schema_json: OnceLock::new(),
}
}
pub fn set_parameter_schema(&mut self, parameter_schema: serde_json::Value) {
self.parameter_schema = parameter_schema;
self.parameter_schema_json = OnceLock::new();
}
pub fn parameter_schema_mut(&mut self) -> &mut serde_json::Value {
self.parameter_schema_json = OnceLock::new();
&mut self.parameter_schema
}
#[must_use]
pub fn parameter_schema_json(&self) -> &str {
self.parameter_schema_json_arc()
}
#[must_use]
pub fn parameter_schema_json_arc(&self) -> &Arc<str> {
self.parameter_schema_json.get_or_init(|| {
match serde_json::to_string(&self.parameter_schema) {
Ok(s) => Arc::from(s),
Err(err) => {
tracing::warn!(error = %err, tool = %self.name, "Failed to serialize parameter_schema to JSON");
Arc::from("{}")
}
}
})
}
#[must_use]
pub const fn is_idempotent(&self) -> bool {
matches!(self.effect, ToolEffect::ReadOnly)
}
}
pub mod prompt;
pub mod quarantine;
pub mod resource;
pub mod truncation;
pub use prompt::*;
pub use quarantine::*;
pub use resource::*;
pub use truncation::*;
#[cfg(all(test, feature = "std"))]
mod tests;