diff --git a/py/core/main/services/ingestion_service.py b/py/core/main/services/ingestion_service.py index fccc7a33c..0958e2f06 100644 --- a/py/core/main/services/ingestion_service.py +++ b/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"] diff --git a/py/core/main/services/retrieval_service.py b/py/core/main/services/retrieval_service.py index 29332cfde..4c6bb1cf3 100644 --- a/py/core/main/services/retrieval_service.py +++ b/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"}, ) ) diff --git a/py/core/parsers/media/audio_parser.py b/py/core/parsers/media/audio_parser.py index 7d5f9f1d4..aa39662a4 100644 --- a/py/core/parsers/media/audio_parser.py +++ b/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, ) diff --git a/py/core/parsers/media/img_parser.py b/py/core/parsers/media/img_parser.py index 1bf64ee73..670c40493 100644 --- a/py/core/parsers/media/img_parser.py +++ b/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: diff --git a/py/core/parsers/media/pdf_parser.py b/py/core/parsers/media/pdf_parser.py index dc4ca8eec..71a3caf6e 100644 --- a/py/core/parsers/media/pdf_parser.py +++ b/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: diff --git a/py/core/providers/embeddings/litellm.py b/py/core/providers/embeddings/litellm.py index 7322d7de4..ce24eea5a 100644 --- a/py/core/providers/embeddings/litellm.py +++ b/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]]: diff --git a/py/core/providers/llm/litellm.py b/py/core/providers/llm/litellm.py index 44d467c2a..735802e3c 100644 --- a/py/core/providers/llm/litellm.py +++ b/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}"