Add Langfuse tracing across RAG ingestion pipeline with datapack-scoped project resolution

Instrument the RAG agent pipeline end-to-end with Langfuse spans/observations so ingestion,
chunking, embedding generation, vector writes, and pipeline state updates are traceable in a
single observability flow. This adds method-level @observe coverage to core pipeline layers
(, , , , ,
, , and active pipeline tracking DB methods), and introduces a
wrapped SQL executor in  to capture query metadata and row counts for
database operations.

In , add datapack-aware Langfuse key resolution from document URI and
, then bind the active public key context so nested
observations are routed to the correct Langfuse project. Also wrap top-level task processing in
an observation span and ensure structured success/error output is emitted to tracing. Improve
failure handling by updating document status to FAILED only when  is available.

Update dependencies by adding  to support the new tracing instrumentation.

Author: kushidhar-in

python/ai/rag_agent/db/active_pipeline_tracking.py

@@ -7,7 +7,7 @@
 from enum import Enum
 from db.connection_pool import ConnectionPool
 
-
+from langfuse import observe
 class PipelineStatus(Enum):
     """Enum for pipeline status values"""
     PROCESSING = 'PROCESSING'
@@ -38,6 +38,7 @@ def _return_connection(self, conn):
         """Return a connection to the pool."""
         self.pool.return_connection(conn)
 
+    @observe(name="Insert Pipeline Details / Active Pipeline Tracking", as_type="span")
     def insert_pipeline_details(
         self,
         document_id: str,
@@ -112,6 +113,7 @@ def insert_pipeline_details(
             if conn:
                 self._return_connection(conn)
 
+    @observe(name="Update Pipeline Status / Active Pipeline Tracking", as_type="span")
     def update_pipeline_status(
         self,
         pipeline_id: str,
@@ -164,6 +166,7 @@ def update_pipeline_status(
             if conn:
                 self._return_connection(conn)
 
+    @observe(name="Mark Pipeline Completed / Active Pipeline Tracking", as_type="span")
     def mark_pipeline_completed(self, pipeline_id: str) -> bool:
         """
         Mark a pipeline as completed.
@@ -206,6 +209,7 @@ def mark_pipeline_completed(self, pipeline_id: str) -> bool:
             if conn:
                 self._return_connection(conn)
 
+    @observe(name="Record Pipeline Error / Active Pipeline Tracking", as_type="span")
     def record_pipeline_error(
         self,
         pipeline_id: str,
@@ -256,6 +260,7 @@ def record_pipeline_error(
             if conn:
                 self._return_connection(conn)
 
+    @observe(name="Update Chunks Processed / Active Pipeline Tracking", as_type="span")
     def update_chunks_processed(self, pipeline_id: str, chunks_count: int) -> bool:
         """
         Update the number of chunks processed for a pipeline.
@@ -301,6 +306,7 @@ def update_chunks_processed(self, pipeline_id: str, chunks_count: int) -> bool:
             if conn:
                 self._return_connection(conn)
 
+    @observe(name="Update Embeddings Persisted / Active Pipeline Tracking", as_type="span")
     def update_embeddings_persisted(self, pipeline_id: str, embeddings_count: int) -> bool:
         """
         Update the number of embeddings persisted for a pipeline.

python/ai/rag_agent/db/yugabytedb_vector_store.py

@@ -6,6 +6,31 @@
 from psycopg import sql
 from db.connection_pool import ConnectionPool
 from db.active_pipeline_tracking import PipelineTracking
+from langfuse import observe, get_client
+
+
+@observe(name="Execute SQL / YugabyteDB Vector Store", as_type="span", capture_input=False, capture_output=False)
+def execute_sql(cur, query: str, params=None, many: bool = False) -> int:
+    if many:
+        cur.executemany(query, params)
+    elif params is not None:
+        cur.execute(query, params)
+    else:
+        cur.execute(query)
+
+    rowcount = cur.rowcount
+
+    get_client().update_current_span(
+        input={
+            "query": query.strip(),
+            "mode": "executemany" if many else "execute",
+            **({"batch_size": len(params)} if many and params is not None else {}),
+            **({"param_count": len(params)} if not many and params is not None else {}),
+        },
+        output={"rowcount": rowcount}
+    )
+
+    return rowcount
 
 
 class YugabyteDBVectorStore:
@@ -40,7 +65,8 @@ def _table_exists(self, conn, table_name, table_schema):
         cur = conn.cursor()
         try:
             # Query information_schema to check if table exists
-            cur.execute(
+            execute_sql(
+                cur,
                 """
                 SELECT EXISTS (
                     SELECT 1 FROM information_schema.tables
@@ -82,6 +108,7 @@ def _ensure_table_exists(self, table_name, schema):
             if conn:
                 self._return_connection(conn)
 
+    @observe(name="Insert Embeddings / YugabyteDB Vector Store", as_type="span")
     def insert_embeddings(
         self,
         document_id,
@@ -142,8 +169,12 @@ def insert_embeddings(
                      tenant_id, metadata_json)
                 )
                 if len(batch) >= batch_size:
-                    cur.executemany(insert_stmt, batch)
-                    rows_inserted = cur.rowcount
+                    rows_inserted = execute_sql(
+                        cur,
+                        insert_stmt,
+                        batch,
+                        many=True
+                    )
                     total_inserted += rows_inserted
                     self.pipeline_tracking.update_embeddings_persisted(
                         pipeline_id=pipeline_id,
@@ -159,8 +190,12 @@ def insert_embeddings(
 
             # Insert any remaining items in the batch
             if batch:
-                cur.executemany(insert_stmt, batch)
-                rows_inserted = cur.rowcount
+                rows_inserted = execute_sql(
+                    cur,
+                    insert_stmt,
+                    batch,
+                    many=True
+                )
                 total_inserted += rows_inserted
                 self.pipeline_tracking.update_embeddings_persisted(
                     pipeline_id=pipeline_id,
@@ -189,6 +224,7 @@ def insert_embeddings(
             if conn:
                 self._return_connection(conn)
 
+    @observe(name="Create Index / YugabyteDB Vector Store", as_type="span")
     def create_index(
         self,
         table_name,
@@ -223,7 +259,7 @@ def create_index(
                 table=sql.Identifier(table_name),
                 ops_class=sql.SQL(distance_metric),
             )
-            cur.execute(sql_create_index)
+            execute_sql(cur, sql_create_index)
 
             conn.commit()
             cur.close()

python/ai/rag_agent/embeddings/embed.py

@@ -2,6 +2,7 @@
 from langchain_openai import OpenAIEmbeddings
 from pdf_processing import PDFProcessor
 from html_processing import HTMLProcessor
+from langfuse import observe
 import logging
 import psycopg
 import os
@@ -118,6 +119,7 @@ def _generate_embeddings_for_text_files(
             f"empty/whitespace chunks, {yielded_count} embeddings yielded"
         )
 
+    @observe(name="Generate Embeddings for PDF Files / EmbeddingsGenerator", as_type="embedding")
     def _generate_embeddings_for_pdf_files(
         self,
         pipeline_id: int,
@@ -167,6 +169,7 @@ def _generate_embeddings_for_pdf_files(
             f"{chunk_count} total chunks, {yielded_count} embeddings yielded"
         )
 
+    @observe(name="Generate Embeddings for HTML Files / EmbeddingsGenerator", as_type="embedding")
     def _generate_embeddings_for_html_file(
         self,
         pipeline_id: int,
@@ -218,10 +221,12 @@ def _generate_embeddings_for_html_file(
             f"{chunk_count} total chunks, {yielded_count} embeddings yielded"
         )
 
+    @observe(name="Generate Embeddings for Video Files / EmbeddingsGenerator", as_type="embedding")
     def _generate_embeddings_for_video_files(self, file_location: str, chunk_args=None):
         """Generate embeddings for video files."""
         pass
 
+    @observe(name="Generate Embeddings / EmbeddingsGenerator", as_type="chain")
     def generate_embeddings(self, pipeline_id: int, file_location: str, chunk_args=None):
         """
         Generator that yields (chunk_text, embedding_vector) tuples.
@@ -275,6 +280,7 @@ def generate_embeddings(self, pipeline_id: int, file_location: str, chunk_args=N
         else:
             raise ValueError(f"Unsupported file type: {file_type}")
 
+    @observe(name="Generate User Prompt Embeddings / EmbeddingsGenerator", as_type="embedding")
     def generate_user_prompt_embeddings(self, user_prompt: str) -> list[float]:
         """Generate embeddings for user prompt."""
 

python/ai/rag_agent/embeddings/embedding_user_promt.py

@@ -2,6 +2,7 @@
 import logging
 import psycopg
 from langchain_openai import OpenAIEmbeddings
+from langfuse import observe
 
 
 class UserPromptEmbedder:
@@ -36,6 +37,7 @@ def __init__(self,
         )
         self.logger = logging.getLogger(__name__)
 
+    @observe(name="Embed Prompt / UserPromptEmbedder", as_type="embedding")
     def embed_prompt(self, prompt: str):
         """Get OpenAI embedding for a single prompt."""
         try:
@@ -45,6 +47,7 @@ def embed_prompt(self, prompt: str):
             self.logger.error(f"Failed to generate embedding for prompt: {e}")
             raise
 
+    @observe(name="Similarity Search / UserPromptEmbedder", as_type="retriever")
     def similarity_search(self, prompt: str):
         """
         Embed the prompt and perform a vector similarity search using HNSW index in the PG table.

python/ai/rag_agent/pdf_processing/process_pdf.py

@@ -9,6 +9,7 @@
 from langchain_community.document_loaders import UnstructuredPDFLoader
 from langchain_core.documents import Document
 from rag_pipeline.chunk import chunk_langchain_docs, DEFAULT_SPLITTER, DEFAULT_ARGS
+from langfuse import observe
 
 
 class PDFProcessor:
@@ -17,6 +18,7 @@ def __init__(self):
         # self.model = model
         pass
 
+    @observe(name="Load PDF from Local / PDFProcessor", as_type="span")
     def _load_pdf_from_local(self, file_path):
         logging.info(f"Loading PDF data from {file_path}")
         try:
@@ -118,6 +120,7 @@ def _load_pdf_from_local(self, file_path):
     #         logging.error(f"Failed to read file from S3: {e}")
     #         raise RuntimeError(f"Failed to read file from S3: {e}")
 
+    @observe(name="Load PDF from S3 / PDFProcessor", as_type="retriever")
     def _load_pdf_from_s3(self, file_path: str):
 
         logging.debug(f"Reading file from S3: {file_path}")
@@ -179,6 +182,7 @@ def _load_pdf_from_s3(self, file_path: str):
             logging.error(f"Failed to read file from S3: {e}")
             raise RuntimeError(f"Failed to read file from S3: {e}")
 
+    @observe(name="Process PDF Data / PDFProcessor", as_type="chain")
     def process_pdf_data(
         self, file_path: str, chunk_args: Dict[str, Any] = {}
     ) -> Generator[str, None, None]:

python/ai/rag_agent/rag_pipeline/chunk.py

@@ -7,6 +7,7 @@
     RecursiveCharacterTextSplitter,
     SpacyTextSplitter
 )
+from langfuse import observe
 import json
 import mimetypes
 
@@ -54,6 +55,7 @@ def get_splitter_for_filetype(file_location: str) -> tuple[str, str]:
     return splitter, args
 
 
+@observe(name="Chunk Text / chunk", as_type="span")
 def chunk(splitter, text, args):
     kwargs = json.loads(args)
 
@@ -63,6 +65,7 @@ def chunk(splitter, text, args):
         raise ValueError("Unknown splitter: {}".format(splitter))
 
 
+@observe(name="Chunk LangChain Docs / chunk", as_type="span")
 def chunk_langchain_docs(splitter, docs, args):
     kwargs = json.loads(args)