Label R2R Langfuse traces

py/core/main/services/ingestion_service.py

@@ -322,6 +322,7 @@ async def augment_document_info(
                     model=self.config.ingestion.document_summary_model
                     or self.config.app.fast_llm
                 ),
+                metadata={"trace_name": "r2r-indexing"},
             )
 
             document_info.summary = response.choices[0].message.content  # type: ignore
@@ -331,6 +332,7 @@ async def augment_document_info(
 
             embedding = await self.providers.embedding.async_get_embedding(
                 text=document_info.summary,
+                metadata={"trace_name": "r2r-indexing"},
             )
             document_info.summary_embedding = embedding
         return
@@ -368,6 +370,7 @@ async def process_batch(
             # Retrieve embeddings in bulk
             vectors = await self.providers.embedding.async_get_embeddings(
                 texts,  # list of strings
+                metadata={"trace_name": "r2r-indexing"},
             )
             # Zip them back together
             results = []
@@ -757,6 +760,7 @@ async def _get_enriched_chunk_text(
                         ),
                         generation_config=chunk_enrichment_settings.generation_config
                         or GenerationConfig(model=self.config.app.fast_llm),
+                        metadata={"trace_name": "r2r-indexing"},
                     )
                 )
                 .choices[0]
@@ -776,7 +780,8 @@ async def _get_enriched_chunk_text(
 
         # Re-embed
         data = await self.providers.embedding.async_get_embedding(
-            updated_chunk_text
+            updated_chunk_text,
+            metadata={"trace_name": "r2r-indexing"},
         )
         chunk["metadata"]["original_text"] = chunk["text"]
 

py/core/main/services/retrieval_service.py

@@ -295,7 +295,8 @@ async def _basic_search(
         ):
             query_vector = (
                 await self.providers.completion_embedding.async_get_embedding(
-                    text=query
+                    text=query,
+                    metadata={"trace_name": "r2r-search"},
                 )
             )
 

py/core/parsers/media/audio_parser.py

@@ -52,10 +52,15 @@ async def ingest(  # type: ignore
                 temp_file_path = temp_file.name
 
             # Call Whisper transcription
+            metadata = {
+                **(kwargs.pop("metadata", None) or {}),
+                "trace_name": "r2r-indexing",
+            }
             response = await self.atranscription(
                 model=self.config.audio_transcription_model
                 or self.config.app.audio_lm,
                 file=open(temp_file_path, "rb"),
+                metadata=metadata,
                 **kwargs,
             )
 

py/core/parsers/media/img_parser.py

@@ -305,7 +305,9 @@ async def ingest(
                 ]
 
             response = await self.llm_provider.aget_completion(
-                messages=messages, generation_config=generation_config
+                messages=messages,
+                generation_config=generation_config,
+                metadata={"trace_name": "r2r-indexing"},
             )
 
             if not response.choices or not response.choices[0].message:

py/core/parsers/media/pdf_parser.py

@@ -202,6 +202,7 @@ async def process_page(self, image, page_num: int) -> dict[str, str]:
                     messages=messages,
                     generation_config=generation_config,
                     apply_timeout=True,
+                    metadata={"trace_name": "r2r-indexing"},
                 )
 
                 if response.choices and response.choices[0].message:

py/core/providers/embeddings/litellm.py

@@ -74,6 +74,10 @@ def _get_embedding_kwargs(self, **kwargs):
         if self.config.api_key:
             embedding_kwargs["api_key"] = self.config.api_key
         embedding_kwargs.update(kwargs)
+        embedding_kwargs["metadata"] = {
+            "trace_name": "r2r",
+            **(embedding_kwargs.get("metadata") or {}),
+        }
         return embedding_kwargs
 
     async def _execute_task(self, task: dict[str, Any]) -> list[list[float]]:

py/core/providers/llm/litellm.py

@@ -53,6 +53,10 @@ async def _execute_task(self, task: dict[str, Any]):
         args = self._get_base_args(generation_config)
         args["messages"] = messages
         args = {**args, **kwargs}
+        args["metadata"] = {
+            "trace_name": "r2r",
+            **(args.get("metadata") or {}),
+        }
 
         logger.debug(
             f"Executing LiteLLM task with generation_config={generation_config}"
@@ -68,6 +72,10 @@ def _execute_task_sync(self, task: dict[str, Any]):
         args = self._get_base_args(generation_config)
         args["messages"] = messages
         args = {**args, **kwargs}
+        args["metadata"] = {
+            "trace_name": "r2r",
+            **(args.get("metadata") or {}),
+        }
 
         logger.debug(
             f"Executing LiteLLM task with generation_config={generation_config}"