validating test case for inference endpoint for openai

Signed-off-by: hansrajr <[email protected]>
monocle2ai · Oct 28, 2024 · d4906d6 · d4906d6
1 parent 3f124ea
commit d4906d6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git 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
@@ -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)