diff --git a/src/monocle_apptrace/metamodel/maps/attributes/inference/langchain_entities.json b/src/monocle_apptrace/metamodel/maps/attributes/inference/langchain_entities.json index 74510cb..473a694 100644 --- a/src/monocle_apptrace/metamodel/maps/attributes/inference/langchain_entities.json +++ b/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']" } ], [ diff --git a/src/monocle_apptrace/wrap_common.py b/src/monocle_apptrace/wrap_common.py index d25c67e..bbfe2ee 100644 --- a/src/monocle_apptrace/wrap_common.py +++ b/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: diff --git a/tests/langchain_chat_sample.py b/tests/langchain_chat_sample.py index 7a83aae..f2674be 100644 --- a/tests/langchain_chat_sample.py +++ b/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": [ # { diff --git a/tests/langchain_custom_output_processor_test.py b/tests/langchain_custom_output_processor_test.py index 549b3cb..cdf98ac 100644 --- a/tests/langchain_custom_output_processor_test.py +++ b/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)