Add a statistics endpoint (#20)

Matthijs van Otterdijk · web-flow · commit f827b1921283 · 2023-10-31T09:47:06.000+01:00
diff --git a/src/server.rs b/src/server.rs
@@ -97,6 +97,7 @@ enum ResourceSpec {
         commit: String,
         threshold: f32,
     },
+    GetStatistics,
 }
 
 #[derive(Debug, Error)]
@@ -149,6 +150,7 @@ fn uri_to_spec(uri: &Uri) -> Result<ResourceSpec, SpecParseError> {
         static ref RE_SEARCH: Regex = Regex::new(r"^/search(/?)$").unwrap();
         static ref RE_SIMILAR: Regex = Regex::new(r"^/similar(/?)$").unwrap();
         static ref RE_DUPLICATES: Regex = Regex::new(r"^/duplicates(/?)$").unwrap();
+        static ref RE_STATISTICS: Regex = Regex::new(r"^/statistics$").unwrap();
     }
     let path = uri.path();
 
@@ -239,6 +241,8 @@ fn uri_to_spec(uri: &Uri) -> Result<ResourceSpec, SpecParseError> {
             }
             _ => Err(SpecParseError::NoCommitIdOrDomain),
         }
+    } else if RE_STATISTICS.is_match(path) {
+        Ok(ResourceSpec::GetStatistics)
     } else {
         Err(SpecParseError::UnknownPath)
     }
@@ -659,6 +663,11 @@ impl Service {
                 let result = self.get_similar_documents(domain, commit, id, count).await;
                 string_response_or_error(result)
             }
+            Ok(ResourceSpec::GetStatistics) => {
+                let statistics = self.vector_store.statistics();
+                let json_string = serde_json::to_string_pretty(&statistics).map_err(|e| e.into());
+                json_response_or_error(json_string)
+            }
             Ok(_) => todo!(),
             Err(e) => Ok(Response::builder()
                 .status(StatusCode::NOT_FOUND)
@@ -801,6 +810,21 @@ fn string_response_or_error(
     }
 }
 
+fn json_response_or_error(
+    result: Result<String, ResponseError>,
+) -> Result<Response<Body>, Infallible> {
+    match result {
+        Ok(task_id) => Ok(Response::builder()
+            .header("Content-Type", "application/json")
+            .body(task_id.into())
+            .unwrap()),
+        Err(e) => Ok(Response::builder()
+            .status(400)
+            .body(e.to_string().into())
+            .unwrap()),
+    }
+}
+
 #[derive(Debug, Error)]
 enum AssignIndexError {
     #[error("io error: {0}")]
diff --git a/src/vectors.rs b/src/vectors.rs
@@ -12,6 +12,7 @@ use std::sync::atomic::{self, AtomicUsize};
 use std::sync::{Arc, Condvar, Mutex, RwLock, Weak};
 
 use lru::LruCache;
+use serde::Serialize;
 use urlencoding::encode;
 
 use crate::vecmath::{Embedding, EmbeddingBytes, EMBEDDING_BYTE_LENGTH, EMBEDDING_LENGTH};
@@ -413,7 +414,7 @@ impl PageArena {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Serialize, Debug, Clone, Copy, PartialEq, Eq)]
 pub struct VectorStoreStatistics {
     free: usize,
     loading: usize,

Original file line number	Diff line number	Diff line change
`@@ -12,6 +12,7 @@ use std::sync::atomic::{self, AtomicUsize};`
`12`	`12`	`use std::sync::{Arc, Condvar, Mutex, RwLock, Weak};`
`13`	`13`
`14`	`14`	`use lru::LruCache;`
	`15`	`+use serde::Serialize;`
`15`	`16`	`use urlencoding::encode;`
`16`	`17`
`17`	`18`	`use crate::vecmath::{Embedding, EmbeddingBytes, EMBEDDING_BYTE_LENGTH, EMBEDDING_LENGTH};`
`@@ -413,7 +414,7 @@ impl PageArena {`
`413`	`414`	`}`
`414`	`415`	`}`
`415`	`416`
`416`		`-#[derive(Debug, Clone, Copy, PartialEq, Eq)]`
	`417`	`+#[derive(Serialize, Debug, Clone, Copy, PartialEq, Eq)]`
`417`	`418`	`pub struct VectorStoreStatistics {`
`418`	`419`	`free: usize,`
`419`	`420`	`loading: usize,`