polish: making chat_buffers and chat_wrappers type safe.

Author: eternalcuriouslearner

instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/chat_buffers.py

@@ -12,41 +12,52 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
 
-class ToolCallBuffer:
-    def __init__(self, index, tool_call_id, function_name):
-        self.index = index
-        self.function_name = function_name
-        self.tool_call_id = tool_call_id
-        self.arguments = []
 
-    def append_arguments(self, arguments):
+class ToolCallBuffer:
+    def __init__(
+        self,
+        index: int,
+        tool_call_id: str | None,
+        function_name: str | None,
+    ) -> None:
+        self.index: int = index
+        self.function_name: str | None = function_name
+        self.tool_call_id: str | None = tool_call_id
+        self.arguments: list[str] = []
+
+    def append_arguments(self, arguments: str | None) -> None:
         if arguments is not None:
             self.arguments.append(arguments)
 
 
 class ChoiceBuffer:
-    def __init__(self, index):
-        self.index = index
-        self.finish_reason = None
-        self.text_content = []
-        self.tool_calls_buffers = []
+    def __init__(self, index: int) -> None:
+        self.index: int = index
+        self.finish_reason: str | None = None
+        self.text_content: list[str] = []
+        self.tool_calls_buffers: list[ToolCallBuffer | None] = []
 
-    def append_text_content(self, content):
+    def append_text_content(self, content: str) -> None:
         self.text_content.append(content)
 
-    def append_tool_call(self, tool_call):
+    def append_tool_call(self, tool_call: ChoiceDeltaToolCall) -> None:
         idx = tool_call.index
         for _ in range(len(self.tool_calls_buffers), idx + 1):
             self.tool_calls_buffers.append(None)
 
         function = tool_call.function
-        if not self.tool_calls_buffers[idx]:
-            self.tool_calls_buffers[idx] = ToolCallBuffer(
+        buffer = self.tool_calls_buffers[idx]
+        if buffer is None:
+            buffer = ToolCallBuffer(
                 idx,
                 tool_call.id,
                 function.name if function else None,
             )
+            self.tool_calls_buffers[idx] = buffer
 
         if function:
-            self.tool_calls_buffers[idx].append_arguments(function.arguments)
+            buffer.append_arguments(function.arguments)

instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/chat_wrappers.py

@@ -125,7 +125,7 @@ def _set_output_messages(self):
                 )
             if choice.tool_calls_buffers:
                 tool_calls = []
-                for tool_call in choice.tool_calls_buffers:
+                for tool_call in filter(None, choice.tool_calls_buffers):
                     arguments = None
                     arguments_str = "".join(tool_call.arguments)
                     if arguments_str:

instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/patch.py

@@ -770,7 +770,7 @@ def cleanup(self, error: Optional[BaseException] = None):
                 message["content"] = "".join(choice.text_content)
             if choice.tool_calls_buffers:
                 tool_calls = []
-                for tool_call in choice.tool_calls_buffers:
+                for tool_call in filter(None, choice.tool_calls_buffers):
                     function = {"name": tool_call.function_name}
                     if self.capture_content:
                         function["arguments"] = "".join(tool_call.arguments)