Migrated python client to gRPC

Author: deven96

README.md

@@ -34,8 +34,8 @@ create_store(
     query_model="all-minilm-l6-v2",
 )
 ```
-- [`ahnlich-client-rs`](ahnlich/client/): Rust client for `ahnlich-db` and `ahnlich-ai` with support for connection pooling.
-- [`ahnlich-client-py`](sdk/ahnlich-client-py/): Python client for `ahnlich-db` and `ahnlich-ai` with support for connection pooling.
+- [`ahnlich-client-rs`](ahnlich/client/): Rust client for `ahnlich-db` and `ahnlich-ai`.
+- [`ahnlich-client-py`](sdk/ahnlich-client-py/): Python client for `ahnlich-db` and `ahnlich-ai`.
 
 - [`ahnlich-cli`](ahnlich/cli/): CLI for querying `ahnlich-db` and `ahnlich-ai`
 

TODO.md

@@ -3,20 +3,16 @@
 - [x] Fixing FIXMEs in DB test
 - [ ] Fixing AI tests to use new grpc methods
 - [x] Fixing ahnlich client tests to use new grpc methods
-- [ ] Fixing ahnlich client README documentation that still references creating connection pools and TCP stuff
+- [x] Fixing ahnlich client README documentation that still references creating connection pools and TCP stuff
 - [X] Fixing DSL to use grpc methods
 - [x] Renaming grpc_types to ahnlich_types
-- [ ] Fixing CLI to use grpc methods
+- [x] Fixing CLI to use grpc methods
 - [ ] Starting Python rewrite to use new grpc methods
 
 
-
-### REWRITE NOTES:
- - Client DB tests are flaky:
-    ```rust
-    thread 'grpc::db::test::test_simple_pipeline' panicked at /Users/davidonuh/Sandbox/rust/ahnlich/ahnlich/utils/src/server.rs:68:33:
-    Could not set up ahnlich-db with allocator_size
-
-    thread 'grpc::db::test::test_simple_pipeline' panicked at client/src/grpc/db.rs:584:14:
-    Could not initialize client: Tonic(tonic::transport::Error(Transport, ConnectError(ConnectError("tcp connect error", Os { code: 61, kind: ConnectionRefused, message: "Connection refused" }))))
-    ```
+## Python Rewrite
+- [ ] Create a blocking client that wraps around betterproto async client
+- [ ] Create CI step that checks that `grpc-update-client` does not produce any diffs so that our clients are always up to date.
+- [X] Update tests.
+- [ ] Fix Demo embed and demo tracing.
+- [ ] Migrate README to reflect grpc client.

ahnlich/client/src/ai.rs

@@ -413,7 +413,7 @@ mod test {
     }
 
     #[tokio::test]
-    async fn test_pool_commands_fail_if_server_not_exist() {
+    async fn test_ai_client_init_fail_if_server_not_exist() {
         let host = "127.0.0.1";
         let port = 1234;
         let ai_client = AiClient::new(format!("{host}:{port}")).await;

ahnlich/client/src/db.rs

@@ -612,7 +612,7 @@ mod test {
     }
 
     #[tokio::test]
-    async fn test_pool_commands_fail_if_server_not_exist() {
+    async fn test_db_client_init_fail_if_server_not_exist() {
         let host = "127.0.0.1";
         let port = 1234;
         let address = format!("{host}:{port}");

ahnlich/types/build.rs

@@ -4,7 +4,7 @@ use std::{
     path::PathBuf,
 };
 use walkdir::WalkDir;
-// TODO: this would serve as a stand in replacement for the types crate
+
 fn main() -> Result<()> {
     // Get the current package directory
     let manifest_dir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));

ahnlich/types/src/ai/mod.rs

@@ -1,6 +1,6 @@
+pub mod query;
 pub mod execution_provider;
 pub mod models;
-pub mod pipeline;
 pub mod preprocess;
-pub mod query;
 pub mod server;
+pub mod pipeline;

ahnlich/types/src/ai/models.rs

@@ -1,17 +1,6 @@
 // This file is @generated by prost-build.
-#[derive(
-    serde::Serialize,
-    serde::Deserialize,
-    Clone,
-    Copy,
-    Debug,
-    PartialEq,
-    Eq,
-    Hash,
-    PartialOrd,
-    Ord,
-    ::prost::Enumeration,
-)]
+#[derive(serde::Serialize, serde::Deserialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
 #[repr(i32)]
 pub enum AiModel {
     AllMiniLmL6V2 = 0,

ahnlich/types/src/ai/pipeline.rs

@@ -20,7 +20,9 @@ pub mod ai_query {
         #[prost(message, tag = "4")]
         CreatePredIndex(super::super::query::CreatePredIndex),
         #[prost(message, tag = "5")]
-        CreateNonLinearAlgorithmIndex(super::super::query::CreateNonLinearAlgorithmIndex),
+        CreateNonLinearAlgorithmIndex(
+            super::super::query::CreateNonLinearAlgorithmIndex,
+        ),
         #[prost(message, tag = "6")]
         DropPredIndex(super::super::query::DropPredIndex),
         #[prost(message, tag = "7")]

ahnlich/types/src/ai/query.rs

@@ -63,10 +63,7 @@ pub struct GetSimN {
     #[prost(uint64, tag = "4")]
     pub closest_n: u64,
     /// Algorithm to use for similarity search
-    #[prost(
-        enumeration = "super::super::algorithm::algorithms::Algorithm",
-        tag = "5"
-    )]
+    #[prost(enumeration = "super::super::algorithm::algorithms::Algorithm", tag = "5")]
     pub algorithm: i32,
     /// Preprocessing actions to apply to input before querying
     #[prost(enumeration = "super::preprocess::PreprocessAction", tag = "6")]

ahnlich/types/src/ai/server.rs

@@ -59,7 +59,8 @@ pub struct CreateIndex {
     #[prost(uint64, tag = "1")]
     pub created_indexes: u64,
 }
-#[derive(Eq, PartialOrd, Ord, Hash, Clone, PartialEq, ::prost::Message)]
+#[derive(Eq, PartialOrd, Ord, Hash)]
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct AiStoreInfo {
     #[prost(string, tag = "1")]
     pub name: ::prost::alloc::string::String,

ahnlich/types/src/algorithm/mod.rs

@@ -1,2 +1,2 @@
-pub mod algorithms;
 pub mod nonlinear;
+pub mod algorithms;

ahnlich/types/src/algorithm/nonlinear.rs

@@ -1,17 +1,6 @@
 // This file is @generated by prost-build.
-#[derive(
-    serde::Serialize,
-    serde::Deserialize,
-    Clone,
-    Copy,
-    Debug,
-    PartialEq,
-    Eq,
-    Hash,
-    PartialOrd,
-    Ord,
-    ::prost::Enumeration,
-)]
+#[derive(serde::Serialize, serde::Deserialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
 #[repr(i32)]
 pub enum NonLinearAlgorithm {
     KdTree = 0,

ahnlich/types/src/client.rs

@@ -1,5 +1,6 @@
 // This file is @generated by prost-build.
-#[derive(PartialOrd, Ord, Eq, Clone, PartialEq, ::prost::Message)]
+#[derive(PartialOrd, Ord, Eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ConnectedClient {
     #[prost(string, tag = "1")]
     pub address: ::prost::alloc::string::String,

ahnlich/types/src/db/mod.rs

@@ -1,3 +1,3 @@
+pub mod server;
 pub mod pipeline;
 pub mod query;
-pub mod server;

ahnlich/types/src/db/pipeline.rs

@@ -22,7 +22,9 @@ pub mod db_query {
         #[prost(message, tag = "5")]
         CreatePredIndex(super::super::query::CreatePredIndex),
         #[prost(message, tag = "6")]
-        CreateNonLinearAlgorithmIndex(super::super::query::CreateNonLinearAlgorithmIndex),
+        CreateNonLinearAlgorithmIndex(
+            super::super::query::CreateNonLinearAlgorithmIndex,
+        ),
         #[prost(message, tag = "7")]
         DropPredIndex(super::super::query::DropPredIndex),
         #[prost(message, tag = "8")]

ahnlich/types/src/db/query.rs

@@ -58,10 +58,7 @@ pub struct GetSimN {
     #[prost(uint64, tag = "3")]
     pub closest_n: u64,
     /// The algorithm to use for similarity computation.
-    #[prost(
-        enumeration = "super::super::algorithm::algorithms::Algorithm",
-        tag = "4"
-    )]
+    #[prost(enumeration = "super::super::algorithm::algorithms::Algorithm", tag = "4")]
     pub algorithm: i32,
     /// The predicate condition to apply.
     #[prost(message, optional, tag = "5")]

ahnlich/types/src/db/server.rs

@@ -54,10 +54,7 @@ pub struct CreateIndex {
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ServerResponse {
-    #[prost(
-        oneof = "server_response::Response",
-        tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10"
-    )]
+    #[prost(oneof = "server_response::Response", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10")]
     pub response: ::core::option::Option<server_response::Response>,
 }
 /// Nested message and enum types in `ServerResponse`.
@@ -86,7 +83,8 @@ pub mod server_response {
         CreateIndex(super::CreateIndex),
     }
 }
-#[derive(Hash, Eq, Ord, PartialOrd, Clone, PartialEq, ::prost::Message)]
+#[derive(Hash, Eq, Ord, PartialOrd)]
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct StoreInfo {
     #[prost(string, tag = "1")]
     pub name: ::prost::alloc::string::String,

ahnlich/types/src/keyval.rs

@@ -34,9 +34,12 @@ pub struct StoreEntry {
     #[prost(message, optional, tag = "2")]
     pub value: ::core::option::Option<StoreValue>,
 }
-#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message)]
+#[derive(serde::Serialize, serde::Deserialize)]
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct StoreValue {
     #[prost(map = "string, message", tag = "1")]
-    pub value:
-        ::std::collections::HashMap<::prost::alloc::string::String, super::metadata::MetadataValue>,
+    pub value: ::std::collections::HashMap<
+        ::prost::alloc::string::String,
+        super::metadata::MetadataValue,
+    >,
 }

ahnlich/types/src/lib.rs

@@ -1,13 +1,13 @@
-pub mod ai;
-pub mod algorithm;
-pub mod client;
-pub mod db;
-pub mod keyval;
+pub mod shared;
+pub mod utils;
 pub mod metadata;
 pub mod predicates;
+pub mod client;
+pub mod keyval;
 pub mod server_types;
 pub mod services;
-pub mod shared;
-pub mod similarity;
-pub mod utils;
+pub mod algorithm;
 pub mod version;
+pub mod db;
+pub mod ai;
+pub mod similarity;

ahnlich/types/src/metadata.rs

@@ -1,32 +1,14 @@
 // This file is @generated by prost-build.
-#[derive(
-    serde::Serialize,
-    serde::Deserialize,
-    PartialOrd,
-    Ord,
-    Hash,
-    Eq,
-    Clone,
-    PartialEq,
-    ::prost::Message,
-)]
+#[derive(serde::Serialize, serde::Deserialize, PartialOrd, Ord, Hash, Eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct MetadataValue {
     #[prost(oneof = "metadata_value::Value", tags = "2, 3")]
     pub value: ::core::option::Option<metadata_value::Value>,
 }
 /// Nested message and enum types in `MetadataValue`.
 pub mod metadata_value {
-    #[derive(
-        serde::Serialize,
-        serde::Deserialize,
-        PartialOrd,
-        Ord,
-        Hash,
-        Eq,
-        Clone,
-        PartialEq,
-        ::prost::Oneof,
-    )]
+    #[derive(serde::Serialize, serde::Deserialize, PartialOrd, Ord, Hash, Eq)]
+    #[derive(Clone, PartialEq, ::prost::Oneof)]
     pub enum Value {
         #[prost(string, tag = "2")]
         RawString(::prost::alloc::string::String),