synapto-interface 0.1.0-dev.8

#![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>;
}

#[doc = " Type-erased trait for tools registered dynamically at runtime."]
#[async_trait::async_trait]
pub trait ErasedTool: Send + Sync + 'static {
    #[doc = " Unique identifier name of the tool."]
    fn name(&self) -> &'static str;
    #[doc = " Human/LLM-readable description explaining the tool's capability."]
    fn description(&self) -> &'static str;
    #[doc = " JSON Schema describing expected tool call arguments."]
    fn schema(&self) -> schemars::Schema;
    #[doc = " Evaluated per turn to determine if tool is currently active/available."]
    async fn erased_is_available(
        &self,
        _ctx_request: &ContextRequest,
        _compiled_context: &serde_json::Value,
    ) -> Result<bool, String> {
        Ok(true)
    }
    #[doc = " Executes the tool with untyped JSON arguments and returns structured JSON output."]
    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
    }
}

/// Opaque handle wrapping a type-erased tool for dynamic registration.
#[derive(Clone)]
pub struct ToolHandle(std::sync::Arc<dyn ErasedTool>);

impl std::fmt::Debug for ToolHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ToolHandle")
            .field("name", &self.0.name())
            .finish()
    }
}

impl ToolHandle {
    pub fn new<T: ErasedTool + 'static>(tool: T) -> Self {
        Self(std::sync::Arc::new(tool))
    }

    pub fn from_arc(tool: std::sync::Arc<dyn ErasedTool>) -> Self {
        Self(tool)
    }

    pub fn into_inner(self) -> std::sync::Arc<dyn ErasedTool> {
        self.0
    }

    pub fn inner(&self) -> &std::sync::Arc<dyn ErasedTool> {
        &self.0
    }

    pub fn name(&self) -> &'static str {
        self.0.name()
    }

    pub fn description(&self) -> &'static str {
        self.0.description()
    }

    pub fn schema(&self) -> schemars::Schema {
        self.0.schema()
    }

    pub async fn erased_is_available(
        &self,
        ctx_request: &ContextRequest,
        compiled_context: &serde_json::Value,
    ) -> Result<bool, String> {
        self.0
            .erased_is_available(ctx_request, compiled_context)
            .await
    }

    pub async fn erased_execute(
        &self,
        ctx_request: &ContextRequest,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.0.erased_execute(ctx_request, args).await
    }
}

#[async_trait::async_trait]
impl ErasedTool for ToolHandle {
    fn name(&self) -> &'static str {
        self.0.name()
    }
    fn description(&self) -> &'static str {
        self.0.description()
    }
    fn schema(&self) -> schemars::Schema {
        self.0.schema()
    }
    async fn erased_is_available(
        &self,
        ctx_request: &ContextRequest,
        compiled_context: &serde_json::Value,
    ) -> Result<bool, String> {
        self.0
            .erased_is_available(ctx_request, compiled_context)
            .await
    }
    async fn erased_execute(
        &self,
        ctx_request: &ContextRequest,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        self.0.erased_execute(ctx_request, 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);