From 19bc6f29f5c3f5171e9ad3bb36b20a1f90475109 Mon Sep 17 00:00:00 2001 From: hansrajr Date: Mon, 30 Sep 2024 15:58:04 +0530 Subject: [PATCH] updating the context_output for vector store Signed-off-by: hansrajr --- src/monocle_apptrace/wrap_common.py | 10 ++++++++-- tests/langchain_chat_sample.py | 20 +++++++++++++++++--- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/monocle_apptrace/wrap_common.py b/src/monocle_apptrace/wrap_common.py index b2432b0..f5e541d 100644 --- a/src/monocle_apptrace/wrap_common.py +++ b/src/monocle_apptrace/wrap_common.py @@ -62,12 +62,12 @@ def get_embedding_model_for_vectorstore(instance): }, 'llama_index.core.indices.base_retriever': lambda instance: { 'provider': type(instance._vector_store).__name__, - 'embedding_model': get_embedding_model_for_vectorstore(instance), # Now calls the updated function + 'embedding_model': get_embedding_model_for_vectorstore(instance), 'type': VECTOR_STORE, }, 'haystack.components.retrievers': lambda instance: { 'provider': instance.__dict__.get("document_store").__class__.__name__, - 'embedding_model': get_embedding_model(), # Call your custom embedding model logic here + 'embedding_model': get_embedding_model(), 'type': VECTOR_STORE, }, } @@ -282,6 +282,12 @@ def update_span_with_context_input(to_wrap, wrapped_args ,span: Span): def update_span_with_context_output(to_wrap, return_value ,span: Span): package_name: str = to_wrap.get('package') + if "langchain_core.retrievers" in package_name: + combined_output = " ".join([doc.page_content for doc in return_value if hasattr(doc, 'page_content')]) + if len(combined_output) > 100: + combined_output = combined_output[:100] + "..." + + span.add_event(CONTEXT_OUTPUT_KEY, {RESPONSE: combined_output}) if "llama_index.core.indices.base_retriever" in package_name: output_arg_text = return_value[0].text span.add_event(CONTEXT_OUTPUT_KEY, {RESPONSE:output_arg_text}) diff --git a/tests/langchain_chat_sample.py b/tests/langchain_chat_sample.py index 725d0ea..dc0884a 100644 --- a/tests/langchain_chat_sample.py +++ b/tests/langchain_chat_sample.py @@ -130,10 +130,17 @@ def format_docs(docs): # "events": [ # { # "name": "context_input", -# "timestamp": "2024-09-16T09:48:53.462202Z", +# "timestamp": "2024-09-30T09:49:04.602503Z", # "attributes": { # "question": "What is Task Decomposition?" # } +# }, +# { +# "name": "context_output", +# "timestamp": "2024-09-30T09:49:05.131481Z", +# "attributes": { +# "response": "Fig. 1. Overview of a LLM-powered autonomous agent system.\nComponent One: Planning#\nA complicated ta..." +# } # } # ], # "links": [], @@ -495,9 +502,16 @@ def format_docs(docs): # "events": [ # { # "name": "context_input", -# "timestamp": "2024-09-16T09:48:56.731819Z", +# "timestamp": "2024-09-30T09:49:07.279965Z", +# "attributes": { +# "question": "What are typical methods used for task decomposition?" +# } +# }, +# { +# "name": "context_output", +# "timestamp": "2024-09-30T09:49:07.809575Z", # "attributes": { -# "question": "What are some typical methods for task decomposition?" +# "response": "Tree of Thoughts (Yao et al. 2023) extends CoT by exploring multiple reasoning possibilities at each..." # } # } # ],