Merge pull request #66 from beehyv/inference_endpoint

making change to capture inference endpoint for openai

src/monocle_apptrace/metamodel/maps/attributes/inference/langchain_entities.json

@@ -17,7 +17,7 @@
       },
       {
         "attribute": "inference_endpoint",
-        "accessor": "lambda instance,args: resolve_from_alias(instance.__dict__, ['azure_endpoint', 'api_base'])"
+        "accessor": "lambda instance,args: resolve_from_alias(instance.__dict__, ['azure_endpoint', 'api_base']) or args['inference_endpoint']"
       }
     ],
     [

src/monocle_apptrace/wrap_common.py

@@ -213,8 +213,8 @@ async def allm_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
         if 'haystack.components.retrievers' in to_wrap['package'] and 'haystack.retriever' in span.name:
             input_arg_text = get_attribute(DATA_INPUT_KEY)
             span.add_event(DATA_INPUT_KEY, {QUERY: input_arg_text})
-        provider_name = set_provider_name(instance)
-        instance_args = {"provider_name": provider_name}
+        provider_name, inference_endpoint = get_provider_name(instance)
+        instance_args = {"provider_name": provider_name, "inference_endpoint": inference_endpoint}
 
         process_span(to_wrap, span, instance, instance_args)
 
@@ -246,8 +246,8 @@ def llm_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
         if 'haystack.components.retrievers' in to_wrap['package'] and 'haystack.retriever' in span.name:
             input_arg_text = get_attribute(DATA_INPUT_KEY)
             span.add_event(DATA_INPUT_KEY, {QUERY: input_arg_text})
-        provider_name = set_provider_name(instance)
-        instance_args = {"provider_name": provider_name}
+        provider_name, inference_endpoint = get_provider_name(instance)
+        instance_args = {"provider_name": provider_name, "inference_endpoint": inference_endpoint}
 
         process_span(to_wrap, span, instance, instance_args)
 
@@ -289,12 +289,16 @@ def update_llm_endpoint(curr_span: Span, instance):
         )
 
 
-def set_provider_name(instance):
+def get_provider_name(instance):
     provider_url = ""
-
+    inference_endpoint = ""
     try:
         if isinstance(instance.client._client.base_url.host, str):
             provider_url = instance.client._client.base_url.host
+        if isinstance(instance.client._client.base_url, str):
+            inference_endpoint = instance.client._client.base_url
+        else:
+            inference_endpoint = str(instance.client._client.base_url)
     except:
         pass
 
@@ -309,7 +313,7 @@ def set_provider_name(instance):
             parsed_provider_url = urlparse(provider_url)
     except:
         pass
-    return parsed_provider_url.hostname or provider_url
+    return parsed_provider_url.hostname or provider_url,inference_endpoint
 
 
 def is_root_span(curr_span: Span) -> bool:

tests/langchain_chat_sample.py

@@ -278,9 +278,9 @@ def format_docs(docs):
 #         "entity.count": 2,
 #         "entity.1.type": "inference.azure_oai",
 #         "entity.1.provider_name": "api.openai.com",
+#         "entity.1.inference_endpoint": "https://api.openai.com/v1/",
 #         "entity.2.name": "gpt-3.5-turbo-0125",
-#         "entity.2.type": "model.llm",
-#         "entity.2.model_name": "gpt-3.5-turbo-0125"
+#         "entity.2.type": "model.llm.gpt-3.5-turbo-0125"
 #     },
 #     "events": [
 #         {

tests/langchain_custom_output_processor_test.py

@@ -170,10 +170,12 @@ def test_llm_chain(self, test_name, test_input_infra, test_output_infra, mock_po
             dataJson =  json.loads(dataBodyStr) # more asserts can be added on individual fields
 
             llm_vector_store_retriever_span = [x for x in dataJson["batch"] if 'langchain.task.VectorStoreRetriever' in x["name"]][0]
+            inference_span = [x for x in dataJson["batch"] if 'langchain.task.FakeListLLM' in x["name"]][0]
 
             assert llm_vector_store_retriever_span["attributes"]["span.type"] == "retrieval"
             assert llm_vector_store_retriever_span["attributes"]["entity.1.name"] == "FAISS"
             assert llm_vector_store_retriever_span["attributes"]["entity.1.type"] == "vectorstore.FAISS"
+            assert inference_span['attributes']["entity.1.inference_endpoint"] == "https://example.com/"
 
         finally:
             os.environ.pop(test_input_infra)