updating the context_output for vector store

Signed-off-by: hansrajr <[email protected]>
monocle2ai · Sep 30, 2024 · 19bc6f2 · 19bc6f2
1 parent e3309e6
commit 19bc6f2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git 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
@@ -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..."
 #             }
 #         }
 #     ],