Finish resolving type issues

Arize-ai · Jan 23, 2024 · f3e182d · f3e182d
1 parent 413c163
commit f3e182d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/...ion/openinference-instrumentation-boto/src/openinference/instrumentation/boto/__init__.py b/...ion/openinference-instrumentation-boto/src/openinference/instrumentation/boto/__init__.py
@@ -137,14 +137,14 @@ def instrumented_response(*args: Any, **kwargs: Any) -> Dict[str, Any]:
                             content,
                         )
 
-                return response
+                return response  # type: ignore
 
         return instrumented_response
 
     return _invocation_wrapper
 
 
-class BotoInstrumentor(BaseInstrumentor):
+class BotoInstrumentor(BaseInstrumentor):  # type: ignore
     __slots__ = ("_original_client_creator",)
 
     def instrumentation_dependencies(self) -> Collection[str]:

diff --git a/python/instrumentation/openinference-instrumentation-boto/tests/test_instrumentor.py b/python/instrumentation/openinference-instrumentation-boto/tests/test_instrumentor.py
@@ -40,7 +40,7 @@ def instrument(
     in_memory_span_exporter.clear()
 
 
-def test_invoke_client(in_memory_span_exporter: InMemorySpanExporter):
+def test_invoke_client(in_memory_span_exporter: InMemorySpanExporter) -> None:
     output = b'{"completion":" Hello!","stop_reason":"stop_sequence","stop":"\\n\\nHuman:"}'
     streaming_body = StreamingBody(io.BytesIO(output), len(output))
     mock_response = {
@@ -73,7 +73,7 @@ def test_invoke_client(in_memory_span_exporter: InMemorySpanExporter):
     assert len(spans) == 1
     span = spans[0]
     assert span.status.is_ok
-    attributes = dict(span.attributes)
+    attributes = dict(span.attributes or dict())
     assert attributes["llm.model_name"] == "anthropic.claude-v2"
     assert attributes["llm.prompts"] == "Human: hello there? Assistant:"
     assert attributes["message.content"] == " Hello!"