Spaces:
Sleeping
Sleeping
File size: 3,967 Bytes
61bcaab 0f5c565 61bcaab aa3a665 61bcaab aa3a665 61bcaab |
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 |
use std::time::Duration;
use rmcp::{serde_json, transport::StreamableHttpServer};
use rmcp::{Error as McpError, ServerHandler, model::*, schemars, tool};
use tracing_subscriber::{
layer::SubscriberExt,
util::SubscriberInitExt,
{self},
};
const BIND_ADDRESS: &str = "0.0.0.0:9090";
#[tokio::main]
async fn main() -> anyhow::Result<()> {
tracing_subscriber::registry()
.with(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "hf_file_mcp=debug".to_string().into()),
)
.with(tracing_subscriber::fmt::layer())
.init();
let ct = StreamableHttpServer::serve(BIND_ADDRESS.parse()?)
.await?
.with_service(HfFile::new);
tokio::signal::ctrl_c().await?;
ct.cancel();
Ok(())
}
#[derive(Clone)]
pub struct HfFile {
http_client: reqwest::Client,
}
#[tool(tool_box)]
impl HfFile {
pub fn new() -> Self {
let http_client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.build()
.expect("Failed to create HTTP client");
Self {
http_client: http_client,
}
}
#[tool(description = "List files in HF Hub")]
async fn list_repo_file(
&self,
#[tool(param)]
#[schemars(description = "repo")]
repository: String,
#[tool(param)]
#[schemars(description = "repo type (models, datasets, spaces.)")]
repo_type: String,
) -> Result<CallToolResult, McpError> {
let response = self
.http_client
.get(format!(
"https://huggingface.co/api/{repo_type}/{repository}"
))
.send()
.await
.unwrap();
if !response.status().is_success() {
return Err(McpError::internal_error("Failed to list files", None));
}
let body = response.json::<serde_json::Value>().await.unwrap();
let siblings = body.get("siblings").unwrap().as_array().unwrap();
let mut contents = Vec::new();
for sibling in siblings {
contents.push(sibling.get("rfilename").unwrap().as_str().unwrap());
}
Ok(CallToolResult::success(vec![Content::json(contents)?]))
}
#[tool(description = "Get file content in HF Hub")]
async fn get_file_content(
&self,
#[tool(param)]
#[schemars(description = "repo")]
repository: String,
#[tool(param)]
#[schemars(description = "repo type (models, datasets, spaces.)")]
repo_type: String,
#[tool(param)]
#[schemars(description = "file path")]
file_path: String,
) -> Result<CallToolResult, McpError> {
let response = self
.http_client
.get(format!(
"https://huggingface.co/{repo_type}/{repository}/resolve/main/{file_path}"
))
.send()
.await
.unwrap();
if !response.status().is_success() {
println!("Failed to get file content: {}", response.status());
println!(
"Failed to get file content: {}",
response.text().await.unwrap()
);
return Err(McpError::internal_error("Failed to get file content", None));
}
Ok(CallToolResult::success(vec![Content::text(
response.text().await.unwrap(),
)]))
}
}
#[tool(tool_box)]
impl ServerHandler for HfFile {
fn get_info(&self) -> ServerInfo {
ServerInfo {
protocol_version: ProtocolVersion::V_2025_03_26,
capabilities: ServerCapabilities::builder()
.enable_tools()
.build(),
server_info: Implementation::from_build_env(),
instructions: Some("This server provides a Hugging Face file tool that can list and get files from Hugging Face Hub.".to_string()),
}
}
}
|