updating new traces for llama_index, updating framework_vector_store_…

…mapping Signed-off-by: hansrajr <[email protected]>
monocle2ai · Sep 18, 2024 · bc1680c · bc1680c
1 parent 3da013a
commit bc1680c
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 69 deletions.
diff --git a/src/monocle_apptrace/wrap_common.py b/src/monocle_apptrace/wrap_common.py
@@ -29,7 +29,7 @@
     "haystack": "workflow.haystack"
 }
 
-frameworks_mapping = {
+framework_vector_store_mapping = {
     'langchain_core.retrievers': lambda instance: {
         'provider': instance.tags[1],
         'embedding_model': instance.tags[0],
@@ -297,8 +297,8 @@ def update_vectorstore_attributes(to_wrap, instance, span):
     """
     try:
         package = to_wrap.get('package')
-        if package in frameworks_mapping:
-            attributes = frameworks_mapping[package](instance)
+        if package in framework_vector_store_mapping:
+            attributes = framework_vector_store_mapping[package](instance)
             span._attributes.update({
                 TYPE: attributes['type'],
                 PROVIDER: attributes['provider'],

diff --git a/tests/llama_index_sample.py b/tests/llama_index_sample.py
@@ -44,81 +44,125 @@
 print(response)
 
 # {
-#     "trace_id": "0xbd54e5d0edcd96634fa8a02c25c27519",
-#     "start_time": "2024-04-15T23:27:54.806477Z",
-#     "end_time": "2024-04-15T23:27:57.182261Z",
-#     "duration_ms": "2376",
-#     "spans": [
+#     "name": "llamaindex.retrieve",
+#     "context": {
+#         "trace_id": "0x939aa2e13c3ce5b37c74b63dc7cfb163",
+#         "span_id": "0x4249f1d3557d62db",
+#         "trace_state": "[]"
+#     },
+#     "kind": "SpanKind.INTERNAL",
+#     "parent_id": "0x15eb14200cf48548",
+#     "start_time": "2024-09-18T09:11:10.380222Z",
+#     "end_time": "2024-09-18T09:11:11.159369Z",
+#     "status": {
+#         "status_code": "UNSET"
+#     },
+#     "attributes": {
+#         "type": "vector_store",
+#         "provider_name": "ChromaVectorStore",
+#         "embedding_model": "text-embedding-3-large"
+#     },
+#     "events": [
 #         {
-#             "span_name": "llamaindex.retrieve",
-#             "start_time": "2024-04-15T23:27:54.806773Z",
-#             "end_time": "2024-04-15T23:27:55.732604Z",
-#             "duration_ms": "926",
-#             "span_id": "0x030cf03872d4a092",
-#             "trace_id": "0xbd54e5d0edcd96634fa8a02c25c27519",
-#             "parent_id": "0xb4b14a8f14e7e770",
+#             "name": "context_input",
+#             "timestamp": "2024-09-18T09:11:10.380249Z",
 #             "attributes": {
-#             },
-#             "events": []
+#                 "question": "What did the author do growing up?"
+#             }
 #         },
 #         {
-#             "span_name": "llamaindex.openai",
-#             "start_time": "2024-04-15T23:27:55.740299Z",
-#             "end_time": "2024-04-15T23:27:57.181992Z",
-#             "duration_ms": "1442",
-#             "span_id": "0x225fbfb58481e58c",
-#             "trace_id": "0xbd54e5d0edcd96634fa8a02c25c27519",
-#             "parent_id": "0xb4b14a8f14e7e770",
+#             "name": "context_output",
+#             "timestamp": "2024-09-18T09:11:11.159351Z",
 #             "attributes": {
-#                 "model_name": "gpt-3.5-turbo-0125",
-#                 "provider_name": "openai.com",
-#             },
-#             "events": []
+#                 "response": "this is some sample text"
+#             }
+#         }
+#     ],
+#     "links": [],
+#     "resource": {
+#         "attributes": {
+#             "service.name": "llama_index_1"
+#         },
+#         "schema_url": ""
+#     }
+# },
+# {
+#     "name": "llamaindex.openai",
+#     "context": {
+#         "trace_id": "0x939aa2e13c3ce5b37c74b63dc7cfb163",
+#         "span_id": "0x32754f3f46059db0",
+#         "trace_state": "[]"
+#     },
+#     "kind": "SpanKind.INTERNAL",
+#     "parent_id": "0x15eb14200cf48548",
+#     "start_time": "2024-09-18T09:11:11.161538Z",
+#     "end_time": "2024-09-18T09:11:12.893143Z",
+#     "status": {
+#         "status_code": "UNSET"
+#     },
+#     "attributes": {
+#         "temperature": 0.1,
+#         "model_name": "gpt-4",
+#         "provider_name": "api.openai.com",
+#         "inference_endpoint": "https://api.openai.com/v1",
+#         "completion_tokens": 15,
+#         "prompt_tokens": 142,
+#         "total_tokens": 157
+#     },
+#     "events": [],
+#     "links": [],
+#     "resource": {
+#         "attributes": {
+#             "service.name": "llama_index_1"
+#         },
+#         "schema_url": ""
+#     }
+# },
+# {
+#     "name": "llamaindex.query",
+#     "context": {
+#         "trace_id": "0x939aa2e13c3ce5b37c74b63dc7cfb163",
+#         "span_id": "0x15eb14200cf48548",
+#         "trace_state": "[]"
+#     },
+#     "kind": "SpanKind.INTERNAL",
+#     "parent_id": null,
+#     "start_time": "2024-09-18T09:11:10.379910Z",
+#     "end_time": "2024-09-18T09:11:12.894191Z",
+#     "status": {
+#         "status_code": "UNSET"
+#     },
+#     "attributes": {
+#         "tags": [
+#             "text-embedding-3-large",
+#             "ChromaVectorStore"
+#         ],
+#         "workflow_name": "llama_index_1",
+#         "workflow_type": "workflow.llamaindex"
+#     },
+#     "events": [
+#         {
+#             "name": "input",
+#             "timestamp": "2024-09-18T09:11:10.379937Z",
+#             "attributes": {
+#                 "question": "What did the author do growing up?"
+#             }
 #         },
 #         {
-#             "span_name": "llamaindex.query",
-#             "start_time": "2024-04-15T23:27:54.806477Z",
-#             "end_time": "2024-04-15T23:27:57.182261Z",
-#             "duration_ms": "2376",
-#             "span_id": "0xb4b14a8f14e7e770",
-#             "trace_id": "0xbd54e5d0edcd96634fa8a02c25c27519",
-#             "parent_id": "None",
+#             "name": "output",
+#             "timestamp": "2024-09-18T09:11:12.894146Z",
 #             "attributes": {
-#                    "tags": [
-#                       "text-embedding-3-large",
-#                        "ChromaVectorStore"
-#                    ],
-#                   "type": "vector_store",
-#                   "provider_name": "ChromaVectorStore",
-#                   "embedding_model": "text-embedding-3-large",
-#                   "workflow_name": "llama_index_1",
-#                   "workflow_type": "workflow.llamaindex"
-#                   },
-#           "events": [
-#               {
-#                   "name": "input",
-#                   "timestamp": "2024-09-16T10:05:44.687175Z",
-#                   "attributes": {
-#                       "question": "What did the author do growing up?"
-#                   }
-#               },
-#               {
-#                   "name": "output",
-#                   "timestamp": "2024-09-16T10:05:47.345643Z",
-#                   "attributes": {
-#                       "response": "The context does not provide information about what the author did while growing up."
-#                       }
-#                   }
-#               ],
-#           "links": [],
-#           "resource": {
-#               "attributes": {
-#                   "service.name": "llama_index_1"
-#               },
-#               "schema_url": ""
-#              }
+#                 "response": "The context does not provide information about what the author did while growing up."
+#             }
 #         }
-#     ]
+#     ],
+#     "links": [],
+#     "resource": {
+#         "attributes": {
+#             "service.name": "llama_index_1"
+#         },
+#         "schema_url": ""
+#     }
 # }