brk_rmcp 0.1.6

Rust SDK for Model Context Protocol
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

use rmcp::{Error as McpError, model::*, schemars, tool};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub enum ChatRole {
    System,
    User,
    Assistant,
    Tool,
}

#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub struct ChatMessage {
    pub role: ChatRole,
    pub content: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub struct ChatRequest {
    pub system: Option<String>,
    pub messages: Vec<ChatMessage>,
}

#[derive(Clone, Default)]
pub struct Demo;

#[tool(tool_box)]
impl Demo {
    pub fn new() -> Self {
        Self
    }

    #[tool(description = "LLM")]
    async fn chat(
        &self,
        #[tool(aggr)] chat_request: ChatRequest,
    ) -> Result<CallToolResult, McpError> {
        let content = Content::json(chat_request)?;
        Ok(CallToolResult::success(vec![content]))
    }
}

#[test]
fn test_complex_schema() {
    let attr = Demo::chat_tool_attr();
    let input_schema = attr.input_schema;
    let enum_number = input_schema
        .get("definitions")
        .unwrap()
        .as_object()
        .unwrap()
        .get("ChatRole")
        .unwrap()
        .as_object()
        .unwrap()
        .get("enum")
        .unwrap()
        .as_array()
        .unwrap()
        .len();
    assert_eq!(enum_number, 4);
    println!("{}", serde_json::to_string_pretty(&input_schema).unwrap());
}