feat: add add_doc metrics

Signed-off-by: Keming <[email protected]>

docker/compose.yaml

@@ -31,6 +31,7 @@ services:
     image: "kemingy/spladepp"
     environment:
       - MOSEC_TIMEOUT=10000
+      - HF_TOKEN=hf_xxxxxxxxx # add your huggingface token here (required by https://huggingface.co/prithivida/Splade_PP_en_v1)
     ports:
       - "8083:8000"
 

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "httpx~=0.27",
     "cohere~=4.45",
     "prometheus-client~=0.20",
+    "numpy~=1.26",
 ]
 [project.optional-dependencies]
 dev = [

qtext/metrics.py

@@ -7,6 +7,9 @@
 doc_counter = Counter("add_doc", "Added documents", labelnames=labels)
 embedding_histogram = Histogram("embedding_latency_seconds", "Embedding cost time")
 sparse_histogram = Histogram("sparse_latency_seconds", "Sparse embedding cost time")
+add_doc_histogram = Histogram(
+    "add_doc_latency_seconds", "Add doc cost time", labelnames=labels
+)
 text_search_histogram = Histogram(
     "full_text_search_latency_seconds",
     "Full text search cost time",

qtext/pg_client.py

@@ -17,6 +17,7 @@
     sparse_search_histogram,
     text_search_histogram,
     vector_search_histogram,
+    add_doc_histogram,
 )
 from qtext.schema import DefaultTable, Querier
 from qtext.spec import AddNamespaceRequest, QueryDocRequest, SparseEmbedding
@@ -141,6 +142,7 @@ def add_doc(self, req):
             if primary_id is None or getattr(req, primary_id, None) is None:
                 attributes.remove(primary_id)
             placeholders = [getattr(req, key) for key in attributes]
+            start_time = perf_counter()
             self.conn.execute(
                 sql.SQL(
                     "INSERT INTO {table} ({fields}) VALUES ({placeholders})"
@@ -154,6 +156,7 @@ def add_doc(self, req):
                 placeholders,
             )
             self.conn.commit()
+            add_doc_histogram.labels(req.namespace).observe(perf_counter() - start_time)
             doc_counter.labels(req.namespace).inc()
         except psycopg.errors.Error as err:
             logger.info("pg client add doc error", exc_info=err)