@@ -73,24 +73,23 @@ def test_basic_completion(mock_client, mock_openai_response):
7373
7474 assert response == mock_openai_response
7575 assert mock_client .capture .call_count == 1
76-
76+
7777 call_args = mock_client .capture .call_args [1 ]
7878 props = call_args ["properties" ]
79-
79+
8080 assert call_args ["distinct_id" ] == "test-id"
8181 assert call_args ["event" ] == "$ai_generation"
8282 assert props ["$ai_provider" ] == "openai"
8383 assert props ["$ai_model" ] == "gpt-4"
8484 assert props ["$ai_input" ] == [{"role" : "user" , "content" : "Hello" }]
85- assert props ["$ai_output" ] == {
86- "choices" : [{"role" : "assistant" , "content" : "Test response" }]
87- }
85+ assert props ["$ai_output" ] == {"choices" : [{"role" : "assistant" , "content" : "Test response" }]}
8886 assert props ["$ai_input_tokens" ] == 20
8987 assert props ["$ai_output_tokens" ] == 10
9088 assert props ["$ai_http_status" ] == 200
9189 assert props ["foo" ] == "bar"
9290 assert isinstance (props ["$ai_latency" ], float )
9391
92+
9493def test_embeddings (mock_client , mock_embedding_response ):
9594 with patch ("openai.resources.embeddings.Embeddings.create" , return_value = mock_embedding_response ):
9695 client = OpenAI (api_key = "test-key" , posthog_client = mock_client )
@@ -103,10 +102,10 @@ def test_embeddings(mock_client, mock_embedding_response):
103102
104103 assert response == mock_embedding_response
105104 assert mock_client .capture .call_count == 1
106-
105+
107106 call_args = mock_client .capture .call_args [1 ]
108107 props = call_args ["properties" ]
109-
108+
110109 assert call_args ["distinct_id" ] == "test-id"
111110 assert call_args ["event" ] == "$ai_embedding"
112111 assert props ["$ai_provider" ] == "openai"
@@ -115,4 +114,4 @@ def test_embeddings(mock_client, mock_embedding_response):
115114 assert props ["$ai_input_tokens" ] == 10
116115 assert props ["$ai_http_status" ] == 200
117116 assert props ["foo" ] == "bar"
118- assert isinstance (props ["$ai_latency" ], float )
117+ assert isinstance (props ["$ai_latency" ], float )
0 commit comments