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};
pub struct ToolContext {
conversation_id: Option<String>,
state: Arc<RwLock<HashMap<String, serde_json::Value>>>,
extensions: Arc<RwLock<HashMap<TypeId, Box<dyn Any + Send + Sync>>>>,
is_idle: bool,
}
impl ToolContext {
#[must_use]
pub fn new(conversation_id: Option<String>) -> Self {
Self {
conversation_id,
state: Arc::new(RwLock::new(HashMap::new())),
extensions: Arc::new(RwLock::new(HashMap::new())),
is_idle: false,
}
}
#[must_use]
pub fn with_shared_state(
conversation_id: Option<String>,
state: Arc<RwLock<HashMap<String, serde_json::Value>>>,
) -> Self {
Self {
conversation_id,
state,
extensions: Arc::new(RwLock::new(HashMap::new())),
is_idle: false,
}
}
#[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) {
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) {
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))
}
}
}
#[must_use]
pub const fn is_idle(&self) -> bool {
self.is_idle
}
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 text(&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};
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(),
}],
}
}
}
impl From<String> for ResourceOutput {
fn from(text: String) -> Self {
Self::text("", None, text)
}
}
impl From<&str> for ResourceOutput {
fn from(text: &str) -> Self {
Self::text("", None, text)
}
}