feat: add support for openTelemetry (#3576)

* feat: support otel Signed-off-by: Wei Zhang <[email protected]> * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes (attempt 2/3) * chore: drop download model span Signed-off-by: Wei Zhang <[email protected]> * chore: add copus=code to embed span Signed-off-by: Wei Zhang <[email protected]> * chore: update span names Signed-off-by: Wei Zhang <[email protected]> --------- Signed-off-by: Wei Zhang <[email protected]> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
TabbyML · Dec 17, 2024 · 6fa7c2d · 6fa7c2d
1 parent 7fffacb
commit 6fa7c2d
Show file tree

Hide file tree

Showing 8 changed files with 297 additions and 43 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -71,6 +71,11 @@ async-openai = "0.20"
 tracing-test = "0.2"
 clap = "4.3.0"
 ratelimit = "0.10"
+tracing-opentelemetry = "0.28.0"
+opentelemetry = { version = "0.27.0", features = ["trace", "metrics"] }
+opentelemetry_sdk = { version = "0.27.0", default-features = false, features = ["trace", "rt-tokio"] }
+opentelemetry-otlp = { version = "0.27.0" }
+opentelemetry-semantic-conventions = { version = "0.27.0", features = ["semconv_experimental"] }
 
 [workspace.dependencies.uuid]
 version = "1.3.3"

diff --git a/crates/http-api-bindings/src/embedding/openai.rs b/crates/http-api-bindings/src/embedding/openai.rs
@@ -5,6 +5,7 @@ use async_openai::{
 };
 use async_trait::async_trait;
 use tabby_inference::Embedding;
+use tracing::{info_span, Instrument};
 
 pub struct OpenAIEmbeddingEngine {
     client: async_openai::Client<OpenAIConfig>,
@@ -40,7 +41,12 @@ impl Embedding for OpenAIEmbeddingEngine {
             user: None,
             dimensions: None,
         };
-        let resp = self.client.embeddings().create(request).await?;
+        let resp = self
+            .client
+            .embeddings()
+            .create(request)
+            .instrument(info_span!("embedding", kind = "openai"))
+            .await?;
         let data = resp
             .data
             .into_iter()

diff --git a/crates/http-api-bindings/src/rate_limit.rs b/crates/http-api-bindings/src/rate_limit.rs
@@ -9,6 +9,7 @@ use futures::stream::BoxStream;
 use leaky_bucket::RateLimiter;
 use tabby_inference::{ChatCompletionStream, CompletionOptions, CompletionStream, Embedding};
 use tokio::time::Duration;
+use tracing::{info_span, Instrument};
 
 fn new_rate_limiter(rpm: u64) -> RateLimiter {
     let rps = (rpm as f64 / 60.0).ceil() as usize;
@@ -35,7 +36,10 @@ pub fn new_embedding(embedding: Box<dyn Embedding>, request_per_minute: u64) ->
 impl Embedding for RateLimitedEmbedding {
     async fn embed(&self, prompt: &str) -> anyhow::Result<Vec<f32>> {
         self.rate_limiter.acquire(1).await;
-        self.embedding.embed(prompt).await
+        self.embedding
+            .embed(prompt)
+            .instrument(info_span!("rate_limited_compute_embedding"))
+            .await
     }
 }
 

diff --git a/crates/tabby-index/src/code/mod.rs b/crates/tabby-index/src/code/mod.rs
@@ -11,7 +11,7 @@ use tabby_common::{
 };
 use tabby_inference::Embedding;
 use tokio::task::JoinHandle;
-use tracing::warn;
+use tracing::{info_span, warn, Instrument};
 
 use self::intelligence::SourceCode;
 use crate::{
@@ -126,7 +126,11 @@ async fn build_binarize_embedding_tokens(
     embedding: Arc<dyn Embedding>,
     body: &str,
 ) -> Result<Vec<String>> {
-    let embedding = match embedding.embed(body).await {
+    let embedding = match embedding
+        .embed(body)
+        .instrument(info_span!("index_compute_embedding", corpus = corpus::CODE))
+        .await
+    {
         Ok(x) => x,
         Err(err) => {
             bail!("Failed to embed chunk text: {}", err);

diff --git a/crates/tabby/Cargo.toml b/crates/tabby/Cargo.toml
@@ -38,6 +38,11 @@ strum = { workspace = true }
 strfmt = "0.2.4"
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
+tracing-opentelemetry.workspace = true
+opentelemetry.workspace = true
+opentelemetry_sdk.workspace = true
+opentelemetry-otlp.workspace = true
+opentelemetry-semantic-conventions.workspace = true
 tantivy = { workspace = true }
 anyhow = { workspace = true }
 sysinfo = "0.29.8"