diff --git a/Cargo.lock b/Cargo.lock index 75f74b4c1..c7fb3f350 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -585,11 +585,11 @@ dependencies = [ "brk_computer", "brk_core", "brk_indexer", + "brk_rmcp", "brk_vec", "color-eyre", "derive_deref", - "rmcp", - "schemars", + "schemars 1.0.0-rc.2", "serde", "serde_json", "serde_with", @@ -621,6 +621,48 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "brk_rmcp" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90d257949059da288c47b4f88ac80abcaa18afe65972d8b543369ee2099b37db" +dependencies = [ + "base64 0.22.1", + "brk_rmcp-macros", + "bytes", + "chrono", + "futures", + "http", + "http-body", + "http-body-util", + "paste", + "pin-project-lite", + "rand 0.9.1", + "schemars 1.0.0-rc.2", + "serde", + "serde_json", + "sse-stream", + "thiserror 2.0.12", + "tokio", + "tokio-stream", + "tokio-util", + "tower-service", + "tracing", + "uuid", +] + +[[package]] +name = "brk_rmcp-macros" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12869ae1128aad2e0777db2688e661a55d389ecfb863936145ea1b789e0c2d34" +dependencies = [ + "proc-macro2", + "quote", + "serde_json", + "syn 2.0.103", +] + [[package]] name = "brk_rolldown" version = "0.0.1" @@ -941,6 +983,7 @@ dependencies = [ "brk_interface", "brk_logger", "brk_parser", + "brk_rmcp", "brk_vec", "clap", "clap_derive", @@ -948,7 +991,6 @@ dependencies = [ "jiff", "log", "minreq", - "rmcp", "serde", "tokio", "tower-http", @@ -3402,44 +3444,6 @@ dependencies = [ "libc", ] -[[package]] -name = "rmcp" -version = "0.1.6" -dependencies = [ - "base64 0.22.1", - "bytes", - "chrono", - "futures", - "http", - "http-body", - "http-body-util", - "paste", - "pin-project-lite", - "rand 0.9.1", - "rmcp-macros", - "schemars", - "serde", - "serde_json", - "sse-stream", - "thiserror 2.0.12", - "tokio", - "tokio-stream", - "tokio-util", - "tower-service", - "tracing", - "uuid", -] - -[[package]] -name = "rmcp-macros" -version = "0.1.6" -dependencies = [ - "proc-macro2", - "quote", - "serde_json", - "syn 2.0.103", -] - [[package]] name = "rolldown-ariadne" version = "0.5.2" @@ -3539,6 +3543,18 @@ name = "schemars" version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + +[[package]] +name = "schemars" +version = "1.0.0-rc.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "284cf8058b5165c2e1aedb8f56f4b09a838a764d1c110178e1747d1f77a01ceb" dependencies = [ "dyn-clone", "ref-cast", @@ -3549,9 +3565,9 @@ dependencies = [ [[package]] name = "schemars_derive" -version = "0.9.0" +version = "1.0.0-rc.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5016d94c77c6d32f0b8e08b781f7dc8a90c2007d4e77472cc2807bc10a8438fe" +checksum = "3ce51e6bea03670e88edba2d724eed817b97f4ab017a1001bae5da325ce93fe7" dependencies = [ "proc-macro2", "quote", @@ -3713,7 +3729,7 @@ dependencies = [ "hex", "indexmap 1.9.3", "indexmap 2.9.0", - "schemars", + "schemars 0.9.0", "serde", "serde_derive", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index c5b5e9aa6..dd96ac4af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,9 +29,10 @@ brk_core = { version = "0.0.66", path = "crates/brk_core" } brk_exit = { version = "0.0.66", path = "crates/brk_exit" } brk_fetcher = { version = "0.0.66", path = "crates/brk_fetcher" } brk_indexer = { version = "0.0.66", path = "crates/brk_indexer" } +brk_interface = { version = "0.0.66", path = "crates/brk_interface" } brk_logger = { version = "0.0.66", path = "crates/brk_logger" } brk_parser = { version = "0.0.66", path = "crates/brk_parser" } -brk_interface = { version = "0.0.66", path = "crates/brk_interface" } +brk_rmcp = { version = "0.1.6", features = ["transport-streamable-http-server", "transport-worker"]} brk_server = { version = "0.0.66", path = "crates/brk_server" } brk_state = { version = "0.0.66", path = "crates/brk_state" } brk_store = { version = "0.0.66", path = "crates/brk_store" } @@ -46,10 +47,7 @@ jiff = "0.2.15" log = { version = "0.4.27" } minreq = { version = "2.13.4", features = ["https", "serde_json"] } rayon = "1.10.0" -rmcp = { path = "../rust-sdk/crates/rmcp", features = ["transport-streamable-http-server", "transport-worker"] } -schemars = "0.9.0" -# rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "main" , features = ["transport-streamable-http-server", "transport-worker"] } -# schemars = "0.8.0" +schemars = "1.0.0-rc.2" serde = { version = "1.0.219" } serde_bytes = "0.11.17" serde_derive = "1.0.219" diff --git a/crates/brk_cli/src/config.rs b/crates/brk_cli/src/config.rs index 762cf9f09..58fa8c0f6 100644 --- a/crates/brk_cli/src/config.rs +++ b/crates/brk_cli/src/config.rs @@ -87,7 +87,12 @@ pub struct Config { #[arg(long, value_name = "SECONDS")] delay: Option, - /// DEV: Activate to watch the selected website's folder for changes, default: false, saved + /// Activate the Model Context Protocol (MCP) endpoint to give LLMs access to BRK, default: false, saved + #[serde(default, deserialize_with = "default_on_error")] + #[arg(long, value_name = "BOOL")] + mcp: Option, + + /// DEV: Activate watching the selected website's folder for changes, default: false, saved #[serde(default, deserialize_with = "default_on_error")] #[arg(long, value_name = "BOOL")] watch: Option, @@ -171,6 +176,10 @@ impl Config { config_saved.check_collisions = Some(check_collisions); } + if let Some(mcp) = config_args.mcp.take() { + config_saved.mcp = Some(mcp); + } + if let Some(watch) = config_args.watch.take() { config_saved.watch = Some(watch); } @@ -356,6 +365,10 @@ impl Config { self.check_collisions.is_some_and(|b| b) } + pub fn mcp(&self) -> bool { + self.mcp.is_some_and(|b| b) + } + pub fn watch(&self) -> bool { self.watch.is_some_and(|b| b) } diff --git a/crates/brk_cli/src/run.rs b/crates/brk_cli/src/run.rs index 4fe8b4b90..398153356 100644 --- a/crates/brk_cli/src/run.rs +++ b/crates/brk_cli/src/run.rs @@ -57,8 +57,9 @@ pub fn run() -> color_eyre::Result<()> { let server = Server::new(served_indexer, served_computer, config.website())?; let watch = config.watch(); + let mcp = config.mcp(); let opt = Some(tokio::spawn(async move { - server.serve(watch).await.unwrap(); + server.serve(watch, mcp).await.unwrap(); })); sleep(Duration::from_secs(1)); diff --git a/crates/brk_interface/Cargo.toml b/crates/brk_interface/Cargo.toml index 108716c70..8d50deb21 100644 --- a/crates/brk_interface/Cargo.toml +++ b/crates/brk_interface/Cargo.toml @@ -14,7 +14,7 @@ brk_indexer = { workspace = true } brk_vec = { workspace = true } color-eyre = { workspace = true } derive_deref = { workspace = true } -rmcp = { workspace = true } +brk_rmcp = { workspace = true } schemars = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/brk_interface/src/format.rs b/crates/brk_interface/src/format.rs index 6f75cb719..8095f61c6 100644 --- a/crates/brk_interface/src/format.rs +++ b/crates/brk_interface/src/format.rs @@ -1,5 +1,5 @@ +use brk_rmcp::schemars::JsonSchema; use color_eyre::eyre::eyre; -use rmcp::schemars::JsonSchema; use serde::Deserialize; #[allow(clippy::upper_case_acronyms)] diff --git a/crates/brk_interface/src/index.rs b/crates/brk_interface/src/index.rs index f62a0438d..9e991f6ce 100644 --- a/crates/brk_interface/src/index.rs +++ b/crates/brk_interface/src/index.rs @@ -165,6 +165,7 @@ impl<'de> Deserialize<'de> for Index { { let str = String::deserialize(deserializer)?; if let Ok(index) = Index::try_from(str.as_str()) { + // dbg!(index); Ok(index) } else { Err(Error::custom("Bad index")) diff --git a/crates/brk_interface/src/lib.rs b/crates/brk_interface/src/lib.rs index 923fa82f4..23a2708b0 100644 --- a/crates/brk_interface/src/lib.rs +++ b/crates/brk_interface/src/lib.rs @@ -15,6 +15,7 @@ mod format; mod index; mod maybe_ids; mod output; +mod pagination; mod params; mod table; mod vecs; @@ -22,7 +23,8 @@ mod vecs; pub use format::Format; pub use index::Index; pub use output::{Output, Value}; -pub use params::{Pagination, Params, ParamsOpt}; +pub use pagination::{PaginatedIndexParam, PaginationParam}; +pub use params::{IdParam, Params, ParamsOpt}; pub use table::Tabled; use vecs::Vecs; @@ -192,21 +194,15 @@ impl<'a> Interface<'a> { &self.vecs.accepted_indexes } - pub fn get_vecids(&self, pagination: Pagination) -> &[&str] { + pub fn get_vecids(&self, pagination: PaginationParam) -> &[&str] { self.vecs.ids(pagination) } - pub fn get_indexes_to_vecids( - &self, - pagination: Pagination, - ) -> BTreeMap<&'static str, Vec<&str>> { - self.vecs.indexes_to_ids(pagination) + pub fn get_index_to_vecids(&self, paginated_index: PaginatedIndexParam) -> Vec<&str> { + self.vecs.index_to_ids(paginated_index) } - pub fn get_vecids_to_indexes( - &self, - pagination: Pagination, - ) -> BTreeMap<&str, Vec<&'static str>> { - self.vecs.ids_to_indexes(pagination) + pub fn get_vecid_to_indexes(&self, id: String) -> Option<&Vec<&'static str>> { + self.vecs.id_to_indexes(id) } } diff --git a/crates/brk_interface/src/maybe_ids.rs b/crates/brk_interface/src/maybe_ids.rs index 50a4b0766..f2f2bb29c 100644 --- a/crates/brk_interface/src/maybe_ids.rs +++ b/crates/brk_interface/src/maybe_ids.rs @@ -22,9 +22,17 @@ impl<'de> Deserialize<'de> for MaybeIds { where D: serde::Deserializer<'de>, { - let str = String::deserialize(deserializer)?; - Ok(MaybeIds( - str.split(",").map(|s| s.to_string()).collect::>(), - )) + let maybe_ids = match serde_json::Value::deserialize(deserializer)? { + serde_json::Value::String(str) => { + str.split(",").map(|s| s.to_string()).collect::>() + } + serde_json::Value::Array(vec) => vec + .into_iter() + .map(|s| s.as_str().unwrap().to_string()) + .collect::>(), + _ => return Err(serde::de::Error::custom("Bad ids format")), + }; + // dbg!(&maybe_ids); + Ok(MaybeIds(maybe_ids)) } } diff --git a/crates/brk_interface/src/output.rs b/crates/brk_interface/src/output.rs index 22028124c..4562ed66d 100644 --- a/crates/brk_interface/src/output.rs +++ b/crates/brk_interface/src/output.rs @@ -6,6 +6,7 @@ use tabled::Tabled as TabledTabled; use crate::Format; #[derive(Debug, Serialize)] +#[serde(untagged)] pub enum Output { Json(Value), CSV(String), diff --git a/crates/brk_interface/src/pagination.rs b/crates/brk_interface/src/pagination.rs new file mode 100644 index 000000000..cfbe22a5e --- /dev/null +++ b/crates/brk_interface/src/pagination.rs @@ -0,0 +1,31 @@ +use schemars::JsonSchema; +use serde::Deserialize; + +use crate::Index; + +#[derive(Debug, Default, Deserialize, JsonSchema)] +pub struct PaginationParam { + #[serde(alias = "p")] + #[schemars(description = "Pagination index")] + #[serde(default)] + pub page: usize, +} + +impl PaginationParam { + const PER_PAGE: usize = 1_000; + + pub fn start(&self, len: usize) -> usize { + (self.page * Self::PER_PAGE).clamp(0, len) + } + + pub fn end(&self, len: usize) -> usize { + ((self.page + 1) * Self::PER_PAGE).clamp(0, len) + } +} + +#[derive(Debug, Deserialize, JsonSchema)] +pub struct PaginatedIndexParam { + pub index: Index, + #[serde(flatten)] + pub pagination: PaginationParam, +} diff --git a/crates/brk_interface/src/params.rs b/crates/brk_interface/src/params.rs index 37a5d873b..196d5aab6 100644 --- a/crates/brk_interface/src/params.rs +++ b/crates/brk_interface/src/params.rs @@ -1,6 +1,6 @@ use std::ops::Deref; -use rmcp::schemars::{self, JsonSchema}; +use brk_rmcp::schemars::{self, JsonSchema}; use serde::{Deserialize, Deserializer}; use crate::{Format, Index, maybe_ids::MaybeIds}; @@ -52,6 +52,7 @@ pub struct ParamsOpt { count: Option, /// Format of the output + #[serde(default)] format: Option, } @@ -85,7 +86,7 @@ impl ParamsOpt { if let Some(c) = self.count { let c = c as i64; if let Some(f) = self.from { - if f.is_positive() || f.abs() > c { + if f >= 0 || f.abs() > c { return Some(f + c); } } else { @@ -101,6 +102,11 @@ impl ParamsOpt { } } +#[derive(Debug, Deserialize, JsonSchema)] +pub struct IdParam { + pub id: String, +} + fn de_unquote_i64<'de, D>(deserializer: D) -> Result, D::Error> where D: Deserializer<'de>, @@ -154,22 +160,3 @@ where } } } - -#[derive(Debug, Default, Deserialize, JsonSchema)] -pub struct Pagination { - #[serde(alias = "p")] - #[schemars(description = "Pagination index")] - pub page: usize, -} - -impl Pagination { - const PER_PAGE: usize = 1_000; - - pub fn start(&self, len: usize) -> usize { - (self.page * Self::PER_PAGE).clamp(0, len) - } - - pub fn end(&self, len: usize) -> usize { - ((self.page + 1) * Self::PER_PAGE).clamp(0, len) - } -} diff --git a/crates/brk_interface/src/vecs.rs b/crates/brk_interface/src/vecs.rs index 800f4c14d..1f4b65a68 100644 --- a/crates/brk_interface/src/vecs.rs +++ b/crates/brk_interface/src/vecs.rs @@ -5,7 +5,7 @@ use brk_indexer::Indexer; use brk_vec::AnyCollectableVec; use derive_deref::{Deref, DerefMut}; -use crate::params::Pagination; +use crate::pagination::{PaginatedIndexParam, PaginationParam}; use super::index::Index; @@ -19,8 +19,8 @@ pub struct Vecs<'a> { pub index_count: usize, pub id_count: usize, pub vec_count: usize, - serialized_id_to_indexes: BTreeMap<&'a str, Vec<&'static str>>, - serialized_indexes_to_ids: BTreeMap<&'static str, Vec<&'a str>>, + id_to_indexes: BTreeMap<&'a str, Vec<&'static str>>, + indexes_to_ids: BTreeMap>, } impl<'a> Vecs<'a> { @@ -40,17 +40,22 @@ impl<'a> Vecs<'a> { .for_each(|vec| this.insert(vec)); let mut ids = this.id_to_index_to_vec.keys().cloned().collect::>(); - ids.sort_unstable_by(|a, b| { - let len_cmp = a.len().cmp(&b.len()); - if len_cmp == std::cmp::Ordering::Equal { - a.cmp(b) - } else { - len_cmp - } - }); + + let sort_ids = |ids: &mut Vec<&str>| { + ids.sort_unstable_by(|a, b| { + let len_cmp = a.len().cmp(&b.len()); + if len_cmp == std::cmp::Ordering::Equal { + a.cmp(b) + } else { + len_cmp + } + }) + }; + + sort_ids(&mut ids); this.ids = ids; - this.id_count = this.index_to_id_to_vec.keys().count(); + this.id_count = this.id_to_index_to_vec.keys().count(); this.index_count = this.index_to_id_to_vec.keys().count(); this.vec_count = this .index_to_id_to_vec @@ -67,7 +72,7 @@ impl<'a> Vecs<'a> { .keys() .map(|i| (i.serialize_long(), i.possible_values())) .collect::>(); - this.serialized_id_to_indexes = this + this.id_to_indexes = this .id_to_index_to_vec .iter() .map(|(id, index_to_vec)| { @@ -80,16 +85,14 @@ impl<'a> Vecs<'a> { ) }) .collect(); - this.serialized_indexes_to_ids = this + this.indexes_to_ids = this .index_to_id_to_vec .iter() - .map(|(index, id_to_vec)| { - ( - index.serialize_long(), - id_to_vec.keys().cloned().collect::>(), - ) - }) + .map(|(index, id_to_vec)| (*index, id_to_vec.keys().cloned().collect::>())) .collect(); + this.indexes_to_ids + .values_mut() + .for_each(|ids| sort_ids(ids)); this } @@ -150,35 +153,28 @@ impl<'a> Vecs<'a> { } } - pub fn ids(&self, pagination: Pagination) -> &[&'_ str] { + pub fn ids(&self, pagination: PaginationParam) -> &[&'_ str] { let len = self.ids.len(); let start = pagination.start(len); let end = pagination.end(len); &self.ids[start..end] } - pub fn ids_to_indexes(&self, pagination: Pagination) -> BTreeMap<&'_ str, Vec<&'static str>> { - let len = self.serialized_id_to_indexes.len(); - let start = pagination.start(len); - let end = pagination.end(len); - self.serialized_id_to_indexes - .iter() - .skip(start) - .take(end) - .map(|(ids, indexes)| (*ids, indexes.clone())) - .collect() + pub fn id_to_indexes(&self, id: String) -> Option<&Vec<&'static str>> { + self.id_to_indexes.get(id.as_str()) } - pub fn indexes_to_ids(&self, pagination: Pagination) -> BTreeMap<&'static str, Vec<&'a str>> { - let len = self.serialized_indexes_to_ids.len(); + pub fn index_to_ids( + &self, + PaginatedIndexParam { index, pagination }: PaginatedIndexParam, + ) -> Vec<&'a str> { + let vec = self.indexes_to_ids.get(&index).unwrap(); + + let len = vec.len(); let start = pagination.start(len); let end = pagination.end(len); - self.serialized_indexes_to_ids - .iter() - .skip(start) - .take(end) - .map(|(index, ids)| (*index, ids.clone())) - .collect() + + vec.iter().skip(start).take(end).cloned().collect() } } diff --git a/crates/brk_logger/src/lib.rs b/crates/brk_logger/src/lib.rs index 6e6b31825..e1e90c748 100644 --- a/crates/brk_logger/src/lib.rs +++ b/crates/brk_logger/src/lib.rs @@ -25,10 +25,9 @@ pub fn init(path: Option<&Path>) { .unwrap() }); - Builder::from_env( - Env::default() - .default_filter_or("info,fjall=off,lsm_tree=off,rolldown=off,brk_rolldown=off"), - ) + Builder::from_env(Env::default().default_filter_or( + "info,fjall=off,lsm_tree=off,rolldown=off,brk_rolldown=off,rmcp=off,brk_rmcp=off,tracing=off", + )) .format(move |buf, record| { let date_time = Timestamp::now() .to_zoned(tz::TimeZone::system()) @@ -83,6 +82,7 @@ fn write( args: impl Display, ) -> Result<(), std::io::Error> { writeln!(buf, "{} {} {} {}", date_time, dash, level, args) + // Don't remove, used to know the target of unwanted logs // writeln!( // buf, // "{} {} {} {} {}", diff --git a/crates/brk_server/Cargo.toml b/crates/brk_server/Cargo.toml index e881c0003..6f0140ee6 100644 --- a/crates/brk_server/Cargo.toml +++ b/crates/brk_server/Cargo.toml @@ -26,7 +26,7 @@ color-eyre = { workspace = true } jiff = { workspace = true } log = { workspace = true } minreq = { workspace = true } -rmcp = { workspace = true } +brk_rmcp = { workspace = true } serde = { workspace = true } tokio = { workspace = true } tower-http = { version = "0.6.6", features = ["compression-full", "trace"] } diff --git a/crates/brk_server/README.md b/crates/brk_server/README.md index 6c0099bfc..dc669bf97 100644 --- a/crates/brk_server/README.md +++ b/crates/brk_server/README.md @@ -43,27 +43,30 @@ The API uses `brk_interface` and so inherites all of its features including form #### [`GET /api/vecs/index-count`](https://bitcoinresearchkit.org/api/vecs/index-count) -Count of all possible indexes +Get the count of all existing indexes. #### [`GET /api/vecs/id-count`](https://bitcoinresearchkit.org/api/vecs/id-count) -Count of all possible ids +Get the count of all existing vec ids. #### [`GET /api/vecs/variant-count`](https://bitcoinresearchkit.org/api/vecs/variant-count) -Count of all possible variants +Get the count of all existing vecs. \ +Equals to the sum of supported Indexes of each vec id. #### [`GET /api/vecs/indexes`](https://bitcoinresearchkit.org/api/vecs/indexes) -Get a list of all possible indexes +Get the list of all existing indexes. #### [`GET /api/vecs/accepted-indexes`](https://bitcoinresearchkit.org/api/vecs/accepted-indexes) -Get a list of possible indexes and all their accepted variants +Get an object which has all existing indexes as keys and a list of their accepted variants as values. #### [`GET /api/vecs/ids`](https://bitcoinresearchkit.org/api/vecs/ids) -A list of all possible vec ids +Get a paginated list of all existing vec ids. \ +There are up to 1,000 values per page. \ +If the `page` param is omitted, it will default to page `0`. #### [`GET /api/vecs/variants`](https://bitcoinresearchkit.org/api/vecs/variants) diff --git a/crates/brk_server/examples/main.rs b/crates/brk_server/examples/main.rs index b5f7480ac..c92854022 100644 --- a/crates/brk_server/examples/main.rs +++ b/crates/brk_server/examples/main.rs @@ -48,7 +48,7 @@ pub fn main() -> color_eyre::Result<()> { let server = Server::new(served_indexer, served_computer, Website::Default)?; let server = tokio::spawn(async move { - server.serve(true).await.unwrap(); + server.serve(true, true).await.unwrap(); }); if process { diff --git a/crates/brk_server/src/api/mod.rs b/crates/brk_server/src/api/mod.rs index c0ef0690c..99631973c 100644 --- a/crates/brk_server/src/api/mod.rs +++ b/crates/brk_server/src/api/mod.rs @@ -5,7 +5,7 @@ use axum::{ response::{IntoResponse, Redirect, Response}, routing::get, }; -use brk_interface::{Index, Pagination, Params, ParamsOpt}; +use brk_interface::{IdParam, Index, PaginatedIndexParam, PaginationParam, Params, ParamsOpt}; use super::AppState; @@ -18,7 +18,7 @@ pub trait ApiRoutes { fn add_api_routes(self) -> Self; } -const TO_SEPARATOR: &str = "-to-"; +const TO_SEPARATOR: &str = "_to_"; impl ApiRoutes for Router { fn add_api_routes(self) -> Self { @@ -56,29 +56,29 @@ impl ApiRoutes for Router { "/api/vecs/ids", get( async |State(app_state): State, - Query(pagination): Query| + Query(pagination): Query| -> Response { Json(app_state.interface.get_vecids(pagination)).into_response() }, ), ) .route( - "/api/vecs/indexes-to-ids", + "/api/vecs/index-to-ids", get( async |State(app_state): State, - Query(pagination): Query| + Query(paginated_index): Query| -> Response { - Json(app_state.interface.get_indexes_to_vecids(pagination)).into_response() + Json(app_state.interface.get_index_to_vecids(paginated_index)).into_response() }, ), ) .route( - "/api/vecs/ids-to-indexes", + "/api/vecs/id-to-indexes", get( async |State(app_state): State, - Query(pagination): Query| + Query(param): Query| -> Response { - Json(app_state.interface.get_vecids_to_indexes(pagination)).into_response() + Json(app_state.interface.get_vecid_to_indexes(param.id)).into_response() }, ), ) @@ -92,16 +92,18 @@ impl ApiRoutes for Router { Query(params_opt): Query, state: State| -> Response { - let variant = variant.replace("_", "-"); + let variant = variant.replace("-", "_"); let mut split = variant.split(TO_SEPARATOR); - let params = Params::from(( - ( - Index::try_from(split.next().unwrap()).unwrap(), - split.collect::>().join(TO_SEPARATOR), - ), - params_opt, - )); - interface::handler(headers, Query(params), state).await + + if let Ok(index) = Index::try_from(split.next().unwrap()) { + let params = Params::from(( + (index, split.collect::>().join(TO_SEPARATOR)), + params_opt, + )); + interface::handler(headers, Query(params), state).await + } else { + "Bad variant".into_response() + } }, ), ) diff --git a/crates/brk_server/src/lib.rs b/crates/brk_server/src/lib.rs index 94fec9d7e..b35a9ca0c 100644 --- a/crates/brk_server/src/lib.rs +++ b/crates/brk_server/src/lib.rs @@ -113,7 +113,7 @@ impl Server { })) } - pub async fn serve(self, watch: bool) -> color_eyre::Result<()> { + pub async fn serve(self, watch: bool, mcp: bool) -> color_eyre::Result<()> { let state = self.0; if let Some(websites_path) = state.websites_path.clone() { @@ -162,7 +162,7 @@ impl Server { let router = Router::new() .add_api_routes() .add_website_routes(state.website) - .add_mcp_routes(state.interface) + .add_mcp_routes(state.interface, mcp) .route("/version", get(Json(VERSION))) .with_state(state) .layer(compression_layer) diff --git a/crates/brk_server/src/mcp/api.rs b/crates/brk_server/src/mcp/api.rs index 7fa22eb30..e0dd28561 100644 --- a/crates/brk_server/src/mcp/api.rs +++ b/crates/brk_server/src/mcp/api.rs @@ -1,5 +1,5 @@ -use brk_interface::{Interface, Pagination, Params}; -use rmcp::{ +use brk_interface::{IdParam, Interface, PaginatedIndexParam, PaginationParam, Params}; +use brk_rmcp::{ Error as McpError, RoleServer, ServerHandler, model::{ CallToolResult, Content, Implementation, InitializeRequestParam, InitializeResult, @@ -8,6 +8,7 @@ use rmcp::{ service::RequestContext, tool, }; +use log::info; #[derive(Clone)] pub struct API { @@ -23,88 +24,111 @@ impl API { } #[tool(description = " -Get the count of all existing indexes +Get the count of all existing indexes. ")] async fn get_index_count(&self) -> Result { + info!("mcp: get_index_count"); Ok(CallToolResult::success(vec![ Content::json(self.interface.get_index_count()).unwrap(), ])) } #[tool(description = " -Get the count of all existing vec ids +Get the count of all existing vec ids. ")] async fn get_vecid_count(&self) -> Result { + info!("mcp: get_vecid_count"); Ok(CallToolResult::success(vec![ Content::json(self.interface.get_vecid_count()).unwrap(), ])) } #[tool(description = " -Get the count of all existing vecs +Get the count of all existing vecs. +Equals to the sum of supported Indexes of each vec id. ")] async fn get_variant_count(&self) -> Result { + info!("mcp: get_variant_count"); Ok(CallToolResult::success(vec![ Content::json(self.interface.get_vec_count()).unwrap(), ])) } #[tool(description = " -Get the list of all existing indexes +Get the list of all existing indexes. ")] async fn get_indexes(&self) -> Result { + info!("mcp: get_indexes"); Ok(CallToolResult::success(vec![ Content::json(self.interface.get_indexes()).unwrap(), ])) } #[tool(description = " -Get an object which has all existing indexes as keys and a list of their accepted variants as values +Get an object which has all existing indexes as keys and a list of their accepted variants as values. ")] async fn get_accepted_indexes(&self) -> Result { + info!("mcp: get_accepted_indexes"); Ok(CallToolResult::success(vec![ Content::json(self.interface.get_accepted_indexes()).unwrap(), ])) } #[tool(description = " -Get the list of all existing vec ids +Get a paginated list of all existing vec ids. +There are up to 1,000 values per page. +If the `page` param is omitted, it will default to page `0`. ")] async fn get_vecids( &self, - #[tool(aggr)] pagination: Pagination, + #[tool(aggr)] pagination: PaginationParam, ) -> Result { + info!("mcp: get_vecids"); Ok(CallToolResult::success(vec![ Content::json(self.interface.get_vecids(pagination)).unwrap(), ])) } #[tool(description = " -Get an object which has all existing indexes as keys and a list of ids of vecs which support that index as values +Get a paginated list of all vec ids which support a given index. +There are up to 1,000 values per page. +If the `page` param is omitted, it will default to page `0`. ")] - async fn get_indexes_to_vecids( + async fn get_index_to_vecids( &self, - #[tool(aggr)] pagination: Pagination, + #[tool(aggr)] paginated_index: PaginatedIndexParam, ) -> Result { + info!("mcp: get_index_to_vecids"); Ok(CallToolResult::success(vec![ - Content::json(self.interface.get_indexes_to_vecids(pagination)).unwrap(), + Content::json(self.interface.get_index_to_vecids(paginated_index)).unwrap(), ])) } #[tool(description = " -Get an object which has all existing vec ids as keys and a list of indexes supported by that vec id as values +Get a list of all indexes supported by a given vec id. +The list will be empty if the vec id isn't correct. ")] - async fn get_vecids_to_indexes( + async fn get_vecid_to_indexes( &self, - #[tool(aggr)] pagination: Pagination, + #[tool(aggr)] param: IdParam, ) -> Result { + info!("mcp: get_vecid_to_indexes"); Ok(CallToolResult::success(vec![ - Content::json(self.interface.get_vecids_to_indexes(pagination)).unwrap(), + Content::json(self.interface.get_vecid_to_indexes(param.id)).unwrap(), ])) } - #[tool(description = "Get one or multiple vecs depending on given parameters")] + #[tool(description = " +Get one or multiple vecs depending on given parameters. +If you'd like to request multiple vec ids, simply separate them with a ','. +To get the last value set `-1` to the `from` parameter. +The response's format will depend on the given parameters, it will be: +- A value: If requested only one vec and the given range returns one value (for example: `from=-1`) +- A list: If requested only one vec and the given range returns multiple values (for example: `from=-1000&count=100` or `from=-444&to=-333`) +- A matrix: When multiple vecs are requested, even if they each return one value. +")] fn get_vecs(&self, #[tool(aggr)] params: Params) -> Result { + info!("mcp: get_vecs"); Ok(CallToolResult::success(vec![ Content::json(self.interface.search_and_format(params).unwrap()).unwrap(), ])) @@ -114,6 +138,7 @@ Get an object which has all existing vec ids as keys and a list of indexes suppo Get the running version of the Bitcoin Research Kit ")] async fn get_version(&self) -> Result { + info!("mcp: get_version"); Ok(CallToolResult::success(vec![Content::text(format!( "v{VERSION}" ))])) @@ -124,14 +149,22 @@ Get the running version of the Bitcoin Research Kit impl ServerHandler for API { fn get_info(&self) -> ServerInfo { ServerInfo { - protocol_version: ProtocolVersion::V_2025_03_26, + protocol_version: ProtocolVersion::LATEST, capabilities: ServerCapabilities::builder().enable_tools().build(), server_info: Implementation::from_build_env(), instructions: Some( " -This server provides an interface to communicate with a running instance of the Bitcoin Research Kit (brk). -Multiple tools are at your disposal including ones to fetch all sorts of Bitcoin on-chain data. -Arrays are also called Vectors (or Vecs). +This server provides an interface to communicate with a running instance of the Bitcoin Research Kit (also called brk or BRK). + +Multiple tools are at your disposal including ones to fetch all sorts of Bitcoin on-chain metrics and market prices. + +If you're unsure which datasets are available, try out different tools before browsing the web. Each tool gives important information about BRK's capabilities. + +Vectors can also be called 'Vecs', 'Arrays' or 'Datasets', they can all be used interchangeably. + +An 'Index' (or indexes) is the timeframe of a dataset. + +'VecId' (or vecids) are the name of the dataset and what it represents. " .to_string(), ), diff --git a/crates/brk_server/src/mcp/mod.rs b/crates/brk_server/src/mcp/mod.rs index 133feb381..dc1e8de0e 100644 --- a/crates/brk_server/src/mcp/mod.rs +++ b/crates/brk_server/src/mcp/mod.rs @@ -1,23 +1,28 @@ use axum::Router; use brk_interface::Interface; -use rmcp::transport::{ +use brk_rmcp::transport::{ StreamableHttpServerConfig, streamable_http_server::{StreamableHttpService, session::local::LocalSessionManager}, }; mod api; use api::*; +use log::info; use crate::AppState; pub trait MCPRoutes { - fn add_mcp_routes(self, interface: &'static Interface<'static>) -> Self; + fn add_mcp_routes(self, interface: &'static Interface<'static>, mcp: bool) -> Self; } impl MCPRoutes for Router { - fn add_mcp_routes(self, interface: &'static Interface<'static>) -> Self { + fn add_mcp_routes(self, interface: &'static Interface<'static>, mcp: bool) -> Self { + if !mcp { + return self; + } + let config = StreamableHttpServerConfig { - // stateful_mode: false, + // stateful_mode: false, // breaks Claude ..Default::default() }; @@ -27,6 +32,8 @@ impl MCPRoutes for Router { config, ); + info!("Setting MCP..."); + self.nest_service("/mcp", service) } }