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)