typecheck fixes

Author: keith-decker

util/opentelemetry-util-genai/src/opentelemetry/util/genai/data.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import Type
+from typing import Optional, Type
 
 
 @dataclass
@@ -44,7 +44,7 @@ def _to_part_dict(self):
 class ChatGeneration:
     content: str
     type: str
-    finish_reason: str = None
+    finish_reason: Optional[str] = None
 
 
 @dataclass

util/opentelemetry-util-genai/src/opentelemetry/util/genai/emitters.py

@@ -33,7 +33,7 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional, cast
 from uuid import UUID
 
 from opentelemetry import trace
@@ -71,15 +71,18 @@ class _SpanState:
     children: List[UUID] = field(default_factory=list)
 
 
-def _get_property_value(obj, property_name) -> object:
+def _get_property_value(obj: Any, property_name: str) -> Any:
     if isinstance(obj, dict):
-        return obj.get(property_name, None)
+        return cast(Any, obj.get(property_name, None))
 
-    return getattr(obj, property_name, None)
+    return cast(Any, getattr(obj, property_name, None))
 
 
 def _message_to_log_record(
-    message: Message, provider_name, framework, capture_content: bool
+    message: Message,
+    provider_name: Optional[str],
+    framework: Optional[str],
+    capture_content: bool,
 ) -> Optional[LogRecord]:
     content = _get_property_value(message, "content")
     message_type = _get_property_value(message, "type")
@@ -88,7 +91,7 @@ def _message_to_log_record(
     if content and capture_content:
         body = {"type": message_type, "content": content}
 
-    attributes = {
+    attributes: Dict[str, Any] = {
         # TODO: add below to opentelemetry.semconv._incubating.attributes.gen_ai_attributes
         "gen_ai.framework": framework,
         # TODO: Convert below to constant once opentelemetry.semconv._incubating.attributes.gen_ai_attributes is available
@@ -107,9 +110,9 @@ def _message_to_log_record(
 
 def _chat_generation_to_log_record(
     chat_generation: ChatGeneration,
-    index,
-    provider_name,
-    framework,
+    index: int,
+    provider_name: Optional[str],
+    framework: Optional[str],
     capture_content: bool,
 ) -> Optional[LogRecord]:
     if not chat_generation:

util/opentelemetry-util-genai/tests/test_utils.py

@@ -37,9 +37,9 @@ def telemetry_setup():
     trace.set_tracer_provider(trace.NoOpTracerProvider())
 
 
-def test_llm_start_and_stop_creates_span(telemetry_setup):
-    memory_exporter = telemetry_setup
-
+def test_llm_start_and_stop_creates_span(
+    telemetry_setup: InMemorySpanExporter,
+):
     run_id = uuid4()
     message = Message(content="hello world", type="Human", name="message name")
     chat_generation = ChatGeneration(content="hello back", type="AI")
@@ -53,7 +53,7 @@ def test_llm_start_and_stop_creates_span(telemetry_setup):
     )
 
     # Get the spans that were created
-    spans = memory_exporter.get_finished_spans()
+    spans = telemetry_setup.get_finished_spans()
 
     # Verify span was created
     assert len(spans) == 1