mcp: part 2

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)]

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"))

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)
     }
 }

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::<Vec<_>>(),
-        ))
+        let maybe_ids = match serde_json::Value::deserialize(deserializer)? {
+            serde_json::Value::String(str) => {
+                str.split(",").map(|s| s.to_string()).collect::<Vec<_>>()
+            }
+            serde_json::Value::Array(vec) => vec
+                .into_iter()
+                .map(|s| s.as_str().unwrap().to_string())
+                .collect::<Vec<_>>(),
+            _ => return Err(serde::de::Error::custom("Bad ids format")),
+        };
+        // dbg!(&maybe_ids);
+        Ok(MaybeIds(maybe_ids))
     }
 }

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),

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,
+}

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<usize>,
 
     /// Format of the output
+    #[serde(default)]
     format: Option<Format>,
 }
 
@@ -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<Option<i64>, 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)
-    }
-}

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<Index, Vec<&'a str>>,
 }
 
 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::<Vec<_>>();
-        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::<BTreeMap<_, _>>();
-        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::<Vec<_>>(),
-                )
-            })
+            .map(|(index, id_to_vec)| (*index, id_to_vec.keys().cloned().collect::<Vec<_>>()))
             .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()
     }
 }