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()
}
#[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(_) => false,
}
}
pub fn clear_state(&self) {
if let Ok(mut guard) = write_lock(&self.0) {
guard.clear();
}
}
}
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(_) => 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
}
}
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";
#[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 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)
}
}
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, PartialEq, Eq, 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, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PromptRole {
User,
Assistant,
System,
}
impl PromptRole {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::User => "user",
Self::Assistant => "assistant",
Self::System => "system",
}
}
}
impl core::fmt::Display for PromptRole {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptOutputMessage {
pub role: PromptRole,
pub content: String,
}
impl PromptOutputMessage {
pub fn user(content: impl Into<String>) -> Self {
Self {
role: PromptRole::User,
content: content.into(),
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: PromptRole::Assistant,
content: content.into(),
}
}
pub fn system(content: impl Into<String>) -> Self {
Self {
role: PromptRole::System,
content: content.into(),
}
}
}
#[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::user(content)],
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
messages: alloc::vec![PromptOutputMessage::assistant(content)],
}
}
pub fn system(content: impl Into<String>) -> Self {
Self {
messages: alloc::vec![PromptOutputMessage::system(content)],
}
}
}
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;