Spaces:
Sleeping
Sleeping
File size: 4,130 Bytes
61bcaab 0f5c565 61bcaab 92ec493 61bcaab aa3a665 61bcaab 92ec493 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 129 130 |
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> {
println!("Listing files in HF Hub {repo_type}/{repository}");
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> {
println!("Getting file content in HF Hub {file_path} for {repo_type}/{repository}");
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()),
}
}
}
|