From 3da013aa74ee6d21d8b0503c186298354ca2a885 Mon Sep 17 00:00:00 2001 From: hansrajr Date: Wed, 18 Sep 2024 14:57:24 +0530 Subject: [PATCH] adding vector store attributes to correct retriever span for llama index Signed-off-by: hansrajr --- src/monocle_apptrace/haystack/wrap_pipeline.py | 2 ++ src/monocle_apptrace/utils.py | 2 +- src/monocle_apptrace/wrap_common.py | 11 +++++------ 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/monocle_apptrace/haystack/wrap_pipeline.py b/src/monocle_apptrace/haystack/wrap_pipeline.py index f75cb95..8e219f8 100644 --- a/src/monocle_apptrace/haystack/wrap_pipeline.py +++ b/src/monocle_apptrace/haystack/wrap_pipeline.py @@ -58,3 +58,5 @@ def get_embedding_model(instance): return text_embedder.model except: pass + + return None diff --git a/src/monocle_apptrace/utils.py b/src/monocle_apptrace/utils.py index edab403..2bfb8e9 100644 --- a/src/monocle_apptrace/utils.py +++ b/src/monocle_apptrace/utils.py @@ -90,4 +90,4 @@ def get_embedding_model() -> str: @return: The name of the embedding model, or 'unknown' if not set """ - return embedding_model_context.get('embedding_model', 'unknown') \ No newline at end of file + return embedding_model_context.get('embedding_model', 'unknown') diff --git a/src/monocle_apptrace/wrap_common.py b/src/monocle_apptrace/wrap_common.py index cec2953..692569f 100644 --- a/src/monocle_apptrace/wrap_common.py +++ b/src/monocle_apptrace/wrap_common.py @@ -35,9 +35,9 @@ 'embedding_model': instance.tags[0], 'type': VECTOR_STORE, }, - 'llama_index.core.base.base_query_engine': lambda instance: { - 'provider': type(instance.retriever._vector_store).__name__, - 'embedding_model': instance.retriever._embed_model.model_name, + 'llama_index.core.indices.base_retriever': lambda instance: { + 'provider': type(instance._vector_store).__name__, + 'embedding_model': instance._embed_model.model_name, 'type': VECTOR_STORE, }, 'haystack.components.retrievers': lambda instance: { @@ -215,7 +215,6 @@ def get_input_from_args(chain_args): return "" def update_span_from_llm_response(response, span: Span): - # extract token uasge from langchain openai if (response is not None and hasattr(response, "response_metadata")): response_metadata = response.response_metadata @@ -306,7 +305,7 @@ def update_vectorstore_attributes(to_wrap, instance, span): EMBEDDING_MODEL: attributes['embedding_model'] }) else: - print(f"Package '{package}' not recognized for vector store telemetry.") + logger.warning(f"Package '{package}' not recognized for vector store telemetry.") except Exception as e: - print(f"Error updating span attributes: {e}") \ No newline at end of file + logger.error(f"Error updating span attributes: {e}")