synapto-interface 0.1.0-dev.2

#![doc = include_str!("tool.md")]

use crate::context::ContextRequest;
use crate::llm::LLMSafe;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

#[async_trait::async_trait]
pub trait Tool: Send + Sync + 'static {
    type Arguments: schemars::JsonSchema
        + serde::de::DeserializeOwned
        + LLMSafe
        + Send
        + Sync
        + 'static;
    const NAME: &'static str;
    const DESCRIPTION: &'static str;
    #[doc = " Evaluated dynamically every turn AFTER the ContextProviders have compiled the World State."]
    #[doc = " `compiled_context` is the JSON value generated by all ContextProviders that the LLM is about to see."]
    async fn is_available(
        &self,
        _ctx_request: &ContextRequest,
        _compiled_context: &serde_json::Value,
    ) -> Result<bool, String> {
        Ok(true)
    }
    #[doc = " Executes the tool. The result is serialized and fed back to the LLM."]
    async fn execute(
        &self,
        ctx_request: &ContextRequest,
        args: Self::Arguments,
    ) -> Result<serde_json::Value, String>;
}

#[async_trait::async_trait]
pub trait ErasedTool: Send + Sync + 'static {
    fn name(&self) -> &'static str;
    fn description(&self) -> &'static str;
    fn schema(&self) -> schemars::Schema;
    async fn erased_is_available(
        &self,
        ctx_request: &ContextRequest,
        compiled_context: &serde_json::Value,
    ) -> Result<bool, String>;
    async fn erased_execute(
        &self,
        ctx_request: &ContextRequest,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, String>;
}

#[async_trait::async_trait]
impl<T> ErasedTool for T
where
    T: Tool,
{
    fn name(&self) -> &'static str {
        <T as Tool>::NAME
    }
    fn description(&self) -> &'static str {
        <T as Tool>::DESCRIPTION
    }
    fn schema(&self) -> schemars::Schema {
        schemars::schema_for!(<T as Tool>::Arguments)
    }
    async fn erased_is_available(
        &self,
        ctx_request: &ContextRequest,
        compiled_context: &serde_json::Value,
    ) -> Result<bool, String> {
        <T as Tool>::is_available(self, ctx_request, compiled_context).await
    }
    async fn erased_execute(
        &self,
        ctx_request: &ContextRequest,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        let parsed_args = serde_json::from_value(args).map_err(|e| e.to_string())?;
        <T as Tool>::execute(self, ctx_request, parsed_args).await
    }
}

#[derive(Default)]
pub struct ToolRegistryBuilder {
    pub tools: std::sync::RwLock<std::collections::HashMap<String, std::sync::Arc<dyn ErasedTool>>>,
}

impl ToolRegistryBuilder {
    pub fn register<T>(&self, tool: T)
    where
        T: ErasedTool + 'static,
    {
        let tool_arc: std::sync::Arc<dyn ErasedTool> = std::sync::Arc::new(tool);
        self.register_erased(tool_arc);
    }
    pub fn register_erased(&self, tool: std::sync::Arc<dyn ErasedTool>) {
        self.tools
            .write()
            .unwrap_or_else(|e| panic!("Failed to acquire write lock on tools: {:?}", e))
            .insert(tool.name().to_string(), tool);
    }
    pub fn get(&self, name: &str) -> Option<std::sync::Arc<dyn ErasedTool>> {
        self.tools
            .read()
            .unwrap_or_else(|e| panic!("Failed to acquire read lock on tools: {:?}", e))
            .get(name)
            .cloned()
    }
    pub fn get_all(&self) -> Vec<std::sync::Arc<dyn ErasedTool>> {
        self.tools
            .read()
            .unwrap_or_else(|e| panic!("Failed to acquire read lock on tools: {:?}", e))
            .values()
            .cloned()
            .collect()
    }
}

#[derive(
    Serialize,
    Deserialize,
    JsonSchema,
    PartialEq,
    Eq,
    Debug,
    Clone,
    derive_more :: Display,
    derive_more :: From,
    derive_more :: Deref,
)]
pub struct ToolCallId(pub String);