Merge pull request #376 from lmnr-ai/dev

tool span, search, UI fixes

app-server/Cargo.toml

@@ -1,76 +1,73 @@
 [package]
+edition = "2021"
 name = "app-server"
 version = "0.1.0"
-edition = "2021"
 
 [dependencies]
-env_logger = "0.10.0"
 actix-web = "4"
 anyhow = "1"
+env_logger = "0.10.0"
 futures-util = "0.3.28"
-tonic = "0.12.3"
+tonic = {version = "0.12.3", features = ["gzip"]}
 
 # workaround [AFIT](https://rust-lang.github.io/rfcs/3185-static-async-fn-in-trait.html)
 # while [the feature](https://github.com/rust-lang/rust/pull/115822#issuecomment-1718261458) 
 # is lacking full support.
 async-trait = "0.1"
 
-dotenv = "0.15"
-prost = "0.13"
-tokio = { version = "1.24", features = ["macros", "rt-multi-thread"] }
-tokio-stream = { version = "0.1", features = ["net"] }
-futures = "0.3"
-rayon = "1"
-enum_dispatch = "0.3.12"
-reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "json", "stream", "multipart"] }
-serde = "1.0"
-serde_json = "1.0.105"
-log = "0.4.20"
-lazy_static = "1.4.0"
-futures-core = "0.3.28"
-async-stream = "0.3.5"
-actix-service = "2.0.2"
-uuid = { version = "1.4.1", features = ["v4", "fast-rng", "macro-diagnostics", "serde"] }
 actix-multipart = "0.6.1"
+actix-service = "2.0.2"
 actix-web-httpauth = "0.8.1"
-rand = "0.8.5"
-itertools = "0.11.0"
-chrono = { version = "0.4.31", features = ["serde"] }
-moka = { version = "0.12.1", features = ["sync", "future"] }
-sqlx = { version = "0.8", features = ["runtime-tokio", "postgres", "uuid", "json", "chrono", "bigdecimal"] }
-thiserror = "2"
-json_value_merge = "2.0.0"
-serde-jsonlines = "0.5.0"
-regex = "1.10.3"
-csv = "1.3.0"
-fancy-regex = "0.13.0"
-url = "2.5.0"
-bimap = "0.6.3"
-dashmap = "5.5.3"
-reqwest-eventsource = "0.6.0"
-tiktoken-rs = "0.5.9"
-handlebars = { version = "6.2.0", features = ["script_helper"] }
-handlebars_misc_helpers = { version = "0.17.0", features = ["json"] }
-aws-sdk-bedrockruntime = "1.61.0"
+actix-ws = "0.3.0"
+async-stream = "0.3.5"
 aws-config = "1.5.10"
 aws-credential-types = "1.2.1"
-lmnr-baml = { git = "https://github.com/lmnr-ai/lmnr-baml", branch = "rust" }
-lapin = "2.5.0"
+aws-sdk-bedrockruntime = "1.61.0"
+aws-sdk-s3 = "1.63.0"
+base64 = "0.22.1"
+bimap = "0.6.3"
 bytes = "1.7.1"
+chrono = {version = "0.4.31", features = ["serde"]}
 clickhouse = {version = "0.12.2", features = ["rustls-tls", "uuid"]}
-time = "0.3.36"
-rustls = { version = "0.23.19", features = ["ring"] }
-serde_repr = "0.1.19"
+csv = "1.3.0"
+dashmap = "5.5.3"
+dotenv = "0.15"
+enum_dispatch = "0.3.12"
+fancy-regex = "0.13.0"
+futures = "0.3"
+futures-core = "0.3.28"
+handlebars = {version = "6.2.0", features = ["script_helper"]}
+handlebars_misc_helpers = {version = "0.17.0", features = ["json"]}
+itertools = "0.11.0"
+json_value_merge = "2.0.0"
+lapin = "2.5.0"
+lazy_static = "1.4.0"
+lmnr-baml = {git = "https://github.com/lmnr-ai/lmnr-baml", branch = "rust"}
+log = "0.4.20"
+moka = {version = "0.12.1", features = ["sync", "future"]}
 num_cpus = "1.16.0"
+prost = "0.13"
+rand = "0.8.5"
+rayon = "1.10"
+regex = "1.10.3"
+reqwest = {version = "0.12", default-features = false, features = ["rustls-tls", "json", "stream", "multipart"]}
+reqwest-eventsource = "0.6.0"
+rustls = {version = "0.23.19", features = ["ring"]}
+serde = "1.0"
+serde-jsonlines = "0.5.0"
+serde_json = "1.0.105"
 sha3 = "0.10.8"
-aws-sdk-s3 = "1.63.0"
-base64 = "0.22.1"
 sodiumoxide = "0.2.7"
-actix-ws = "0.3.0"
-tokio-tungstenite = "*"
-prost-types = "0.13.4"
+sqlx = {version = "0.8", features = ["runtime-tokio", "postgres", "uuid", "json", "chrono", "bigdecimal"]}
+thiserror = "2"
+tiktoken-rs = "0.5.9"
+time = "0.3.36"
+tokio = {version = "1.24", features = ["macros", "rt-multi-thread"]}
+tokio-stream = {version = "0.1", features = ["net"]}
+tokio-tungstenite = "0.24"
+url = "2.5.0"
+uuid = {version = "1.4.1", features = ["v4", "fast-rng", "macro-diagnostics", "serde"]}
 zstd = "0.13.2"
 
-
 [build-dependencies]
 tonic-build = "0.12.3"

app-server/proto/semantic_search_grpc.proto

@@ -1,8 +1,6 @@
 syntax = "proto3";
 package semantic_search_grpc;
 
-import "google/protobuf/timestamp.proto";
-
 // When new model is added, don't forget to modify delete_collections function accordingly
 enum Model {
     GTE_BASE = 0;
@@ -38,15 +36,6 @@ message DeleteEmbeddingsResponse {
     string status = 1;
 }
 
-message DateRange {
-    string key = 1;
-    google.protobuf.Timestamp gte = 2;
-    google.protobuf.Timestamp lte = 3;
-}
-
-message DateRanges {
-    repeated DateRange date_ranges = 1;
-}
 
 message RequestPayload {
     map<string, string> payload = 1;
@@ -59,7 +48,6 @@ message QueryRequest {
     repeated RequestPayload payloads = 4;
     string collection_name = 5;
     Model model = 6;
-    optional DateRanges date_ranges = 7;
 }
 
 message QueryResponse {

app-server/src/api/v1/semantic_search.rs

@@ -100,8 +100,6 @@ pub async fn semantic_search(
             params.limit.unwrap_or(DEFAULT_LIMIT),
             params.threshold,
             payloads,
-            None,
-            false,
         )
         .await?;
 

app-server/src/api/v1/traces.rs

@@ -10,6 +10,7 @@ use crate::{
     db::{events::Event, project_api_keys::ProjectApiKey, spans::Span, DB},
     features::{is_feature_enabled, Feature},
     opentelemetry::opentelemetry::proto::collector::trace::v1::ExportTraceServiceRequest,
+    pipeline::runner::PipelineRunner,
     routes::types::ResponseResult,
     traces::{limits::get_workspace_limit_exceeded_by_project_id, producer::push_spans_to_queue},
 };
@@ -31,6 +32,7 @@ pub async fn process_traces(
     db: web::Data<DB>,
     clickhouse: web::Data<clickhouse::Client>,
     cache: web::Data<crate::cache::Cache>,
+    pipeline_runner: web::Data<PipelineRunner>,
 ) -> ResponseResult {
     let db = db.into_inner();
     let cache = cache.into_inner();
@@ -40,6 +42,8 @@ pub async fn process_traces(
     })?;
     let rabbitmq_connection = rabbitmq_connection.as_ref().clone();
 
+    let pipeline_runner = pipeline_runner.into_inner();
+
     if is_feature_enabled(Feature::UsageLimit) {
         let limits_exceeded = get_workspace_limit_exceeded_by_project_id(
             db.clone(),
@@ -61,6 +65,7 @@ pub async fn process_traces(
         db,
         clickhouse,
         cache,
+        pipeline_runner,
     )
     .await?;
     if response.partial_success.is_some() {