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, 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()
}
}
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 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_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;
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 {
content: String,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
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 fn metadata(&self) -> &HashMap<String, serde_json::Value> {
&self.metadata
}
}
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, 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(),
}
}
#[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 = serde_json::to_value(value).map_err(|e| {
Self::new(format!(
"{} (metadata serialization also failed: {e})",
self.message
))
})?;
match json {
serde_json::Value::Object(map) => {
self.metadata.extend(map);
Ok(self)
}
other => Err(Self::new(format!(
"{} (metadata must serialize to a JSON object, got {})",
self.message,
other_type_name(&other),
))),
}
}
#[must_use]
pub fn metadata(&self) -> &HashMap<String, serde_json::Value> {
&self.metadata
}
}
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("error_kind", 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 {}
}
}
#[doc(hidden)]
pub mod __private {
#[cfg(not(feature = "std"))]
pub use alloc::borrow::Cow;
#[cfg(not(feature = "std"))]
use alloc::{
format,
string::{String, ToString},
};
pub use core::{convert::Into, result::Result};
#[cfg(feature = "std")]
pub use std::borrow::Cow;
#[cfg(feature = "std")]
pub use std::sync::LazyLock as Lazy;
#[cfg(not(feature = "std"))]
pub use spin::LazyLock as Lazy;
use super::{Json, ToolError, ToolOutput};
#[cfg(feature = "std")]
pub fn log_description_render_error(tool: &str, err: &dyn core::fmt::Display) {
eprintln!(
"llm-tool: tool `{tool}` description template failed to render ({err}); \
falling back to static description"
);
}
#[cfg(not(feature = "std"))]
#[inline]
pub fn log_description_render_error(_tool: &str, _err: &dyn core::fmt::Display) {}
pub struct Wrap<T>(pub T);
impl Wrap<ToolOutput> {
pub fn __convert(self) -> Result<ToolOutput, ToolError> {
Ok(self.0)
}
}
impl Wrap<String> {
pub fn __convert(self) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput::new(self.0))
}
pub fn __convert_prompt(self) -> Result<super::PromptOutput, ToolError> {
Ok(super::PromptOutput::user(self.0))
}
pub fn __convert_resource(
self,
uri: &str,
mime_type: Option<&str>,
) -> Result<super::ResourceOutput, ToolError> {
Ok(super::ResourceOutput::text(uri, mime_type, self.0))
}
}
impl Wrap<&str> {
pub fn __convert_prompt(self) -> Result<super::PromptOutput, ToolError> {
Ok(super::PromptOutput::user(self.0))
}
pub fn __convert_resource(
self,
uri: &str,
mime_type: Option<&str>,
) -> Result<super::ResourceOutput, ToolError> {
Ok(super::ResourceOutput::text(uri, mime_type, self.0))
}
}
impl Wrap<super::PromptOutput> {
pub fn __convert_prompt(self) -> Result<super::PromptOutput, ToolError> {
Ok(self.0)
}
}
impl Wrap<super::ResourceOutput> {
pub fn __convert_resource(
self,
_uri: &str,
_mime_type: Option<&str>,
) -> Result<super::ResourceOutput, ToolError> {
Ok(self.0)
}
}
impl<T: serde::Serialize> Wrap<Json<T>> {
pub fn __convert(self) -> Result<ToolOutput, ToolError> {
Ok((self.0).into())
}
}
pub trait SerializeFallback {
fn __convert(self) -> Result<ToolOutput, ToolError>;
}
impl<T: serde::Serialize> SerializeFallback for Wrap<T> {
fn __convert(self) -> Result<ToolOutput, ToolError> {
let json_value = serde_json::to_value(&self.0)
.map_err(|e| ToolError::new(format!("serialization failed: {e}")))?;
let content = json_value.to_string();
match json_value {
serde_json::Value::Object(map) => Ok(ToolOutput {
content,
metadata: map.into_iter().collect(),
}),
_ => Ok(ToolOutput::new(content)),
}
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
pub name: String,
pub description: String,
pub parameter_schema: serde_json::Value,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptDefinition {
pub name: String,
#[serde(default, skip_serializing_if = "String::is_empty")]
pub description: String,
#[serde(default, skip_serializing_if = "alloc::vec::Vec::is_empty")]
pub arguments: alloc::vec::Vec<PromptArgumentDefinition>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptArgumentDefinition {
pub name: String,
#[serde(default, skip_serializing_if = "String::is_empty")]
pub description: String,
#[serde(default)]
pub required: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptOutputMessage {
pub role: alloc::borrow::Cow<'static, str>,
pub content: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptOutput {
pub messages: alloc::vec::Vec<PromptOutputMessage>,
}
impl PromptOutput {
pub fn user(content: impl Into<String>) -> Self {
Self {
messages: alloc::vec![PromptOutputMessage {
role: alloc::borrow::Cow::Borrowed("user"),
content: content.into(),
}],
}
}
}
impl From<String> for PromptOutput {
fn from(content: String) -> Self {
Self::user(content)
}
}
impl From<&str> for PromptOutput {
fn from(content: &str) -> Self {
Self::user(content)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ResourceDefinition {
#[serde(rename = "uriTemplate")]
pub uri_template: String,
pub name: String,
#[serde(default, skip_serializing_if = "String::is_empty")]
pub description: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub mime_type: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ResourceOutputContent {
#[serde(rename_all = "camelCase")]
Text {
uri: String,
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
text: String,
},
#[serde(rename_all = "camelCase")]
Blob {
uri: String,
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
blob: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceOutput {
pub contents: alloc::vec::Vec<ResourceOutputContent>,
}
impl ResourceOutput {
pub fn text(uri: impl Into<String>, mime_type: Option<&str>, text: impl Into<String>) -> Self {
Self {
contents: alloc::vec![ResourceOutputContent::Text {
uri: uri.into(),
mime_type: mime_type.map(ToString::to_string),
text: text.into(),
}],
}
}
pub fn blob(uri: impl Into<String>, mime_type: Option<&str>, blob: impl Into<String>) -> Self {
Self {
contents: alloc::vec![ResourceOutputContent::Blob {
uri: uri.into(),
mime_type: mime_type.map(ToString::to_string),
blob: blob.into(),
}],
}
}
}
#[cfg(all(test, feature = "std"))]
mod tests;