Fix tests and linting

getsentry · Apr 3, 2024 · d2a66df · d2a66df
1 parent 8084a18
commit d2a66df
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py
@@ -78,9 +78,9 @@ def _calculate_chat_completion_usage(
     messages, response, span, streaming_message_responses=None
 ):
     # type: (Iterable[ChatCompletionMessageParam], Any, Span, Optional[List[str]]) -> None
-    completion_tokens = 0
-    prompt_tokens = 0
-    total_tokens = 0
+    completion_tokens = 0  # type: Optional[int]
+    prompt_tokens = 0  # type: Optional[int]
+    total_tokens = 0  # type: Optional[int]
     if hasattr(response, "usage"):
         if hasattr(response.usage, "completion_tokens") and isinstance(
             response.usage.completion_tokens, int

diff --git a/tests/integrations/langchain/test_langchain.py b/tests/integrations/langchain/test_langchain.py
@@ -149,9 +149,10 @@ def test_langchain_agent(
             "You are very powerful"
             in chat_spans[0]["data"]["ai.input_messages"][0]["content"]
         )
-        assert "5" in chat_spans[0]["data"]["ai.responses"][0]
+        print("Responses are: ", chat_spans[0]["data"]["ai.responses"])
+        assert "5" in chat_spans[0]["data"]["ai.responses"]
         assert "word" in tool_exec_span["data"]["ai.input_messages"]
-        assert "5" in tool_exec_span["data"]["ai.responses"][0]
+        assert "5" in tool_exec_span["data"]["ai.responses"]
         assert (
             "You are very powerful"
             in chat_spans[1]["data"]["ai.input_messages"][0]["content"]