sugars_llm 0.2.0

LLM integration and AI agent builder utilities
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{AsyncTask, AsyncStream};
use crate::domain::chunk::ChatMessageChunk;
use crate::domain::message::MessageRole;
use std::future::Future;
use std::pin::Pin;


// Core tool traits
pub trait Tool: Send + Sync {
    fn name(&self) -> &str;
    fn description(&self) -> &str;
    fn parameters(&self) -> Value;
    fn execute(&self, args: Value) -> Pin<Box<dyn Future<Output = Result<Value, String>> + Send>>;
}

pub trait ToolEmbedding: Send + Sync {
    fn name(&self) -> &str;
    fn description(&self) -> &str;
    fn parameters(&self) -> Value;
    fn execute(&self, args: Value) -> Pin<Box<dyn Future<Output = Result<Value, String>> + Send>>;
    fn embedding(&self) -> &[f64];
}

pub enum ToolType {
    Direct(Box<dyn Tool>),
    Embedding(Box<dyn ToolEmbedding>),
}

pub struct ToolSet {
    tools: Vec<ToolType>,
}

pub struct ToolSetBuilder {
    tools: Vec<ToolType>,
}

impl ToolSet {
    // Direct creation
    pub fn from_tools(tools: Vec<ToolType>) -> Self {
        ToolSet { tools }
    }

    // Semantic builder entry
    pub fn tools() -> ToolSetBuilder {
        ToolSetBuilder {
            tools: Vec::new(),
        }
    }

    pub fn iter(&self) -> impl Iterator<Item = &ToolType> {
        self.tools.iter()
    }

    // Execute tool by name - returns AsyncStream<ChatMessageChunk>
    pub fn execute(&self, name: &str, args: Value) -> AsyncStream<ChatMessageChunk> {
        for tool in &self.tools {
            let tool_name = match tool {
                ToolType::Direct(t) => t.name(),
                ToolType::Embedding(t) => t.name(),
            };

            if tool_name == name {
                // Create a message chunk with tool execution result
                let chunk = ChatMessageChunk {
                    content: format!("Executing tool: {}", name),
                    role: MessageRole::Assistant,
                    is_final: true,
                    metadata: std::collections::HashMap::new(),
                };
                
                let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
                let _ = tx.send(chunk);
                return AsyncStream::new(rx);
            }
        }

        // Tool not found - return error chunk
        let error_chunk = ChatMessageChunk {
            content: format!("Tool '{}' not found", name),
            role: MessageRole::System,
            is_final: true,
            metadata: std::collections::HashMap::new(),
        };
        
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let _ = tx.send(error_chunk);
        AsyncStream::new(rx)
    }
}

impl ToolSetBuilder {
    pub fn add_tool(mut self, tool: ToolType) -> Self {
        self.tools.push(tool);
        self
    }

    pub fn simple<T: Tool + 'static>(mut self, tool: T) -> Self {
        self.tools.push(ToolType::Direct(Box::new(tool)));
        self
    }

    pub fn embedding<T: ToolEmbedding + 'static>(mut self, tool: T) -> Self {
        self.tools.push(ToolType::Embedding(Box::new(tool)));
        self
    }

    // Terminal method
    pub fn register(self) -> ToolSet {
        ToolSet {
            tools: self.tools,
        }
    }

    // Terminal method with execution
    pub fn on_call<F>(self, name: &str, args: Value, handler: F) -> AsyncTask<Value>
    where
        F: FnOnce(Result<Value, String>) -> Value + Send + 'static,
        Value: crate::async_task::NotResult,
    {
        let toolset = self.register();
        let name = name.to_string();
        AsyncTask::spawn(move || {
            // For now, just return error
            // TODO: Implement actual tool execution
            handler(Err(format!("Tool execution not implemented for '{}'", name)))
        })
    }
}