fix(py/genkit): fixing type check errors of mypy

Author: hendrixmar

py/plugins/compat-oai/tests/test_handler.py

@@ -19,7 +19,7 @@
 
 import pytest
 
-from genkit.ai import ActionRunContext
+from genkit.core.action import ActionRunContext
 from genkit.plugins.compat_oai.models import OpenAIModelHandler
 from genkit.plugins.compat_oai.models.model import OpenAIModel
 from genkit.plugins.compat_oai.models.model_info import (

py/plugins/compat-oai/tests/test_model.py

@@ -13,21 +13,30 @@
 # limitations under the License.
 #
 # SPDX-License-Identifier: Apache-2.0
+import sys  # noqa
+
+if sys.version_info < (3, 11):  # noqa
+    from typing_extensions import Self # noqa
+else:  # noqa
+    from typing import Self
+
 
 from unittest.mock import MagicMock
 
 import pytest
 
+
 from genkit.plugins.compat_oai.models import OpenAIModel
 from genkit.plugins.compat_oai.models.model_info import GPT_4
-from genkit.types import (
+from genkit.core.typing  import (
     GenerateResponse,
     GenerateResponseChunk,
     Role,
+    GenerateRequest
 )
 
 
-def test_get_messages(sample_request):
+def test_get_messages(sample_request: GenerateRequest) -> None:
     """Test _get_messages method.
     Ensures the method correctly converts GenerateRequest messages into OpenAI-compatible ChatMessage format.
     """
@@ -41,7 +50,7 @@ def test_get_messages(sample_request):
     assert messages[1]['content'] == 'Hello, world!'
 
 
-def test_get_openai_config(sample_request):
+def test_get_openai_config(sample_request: GenerateRequest) -> None:
     """
     Test _get_openai_request_config method.
     Ensures the method correctly constructs the OpenAI API configuration dictionary.
@@ -55,7 +64,7 @@ def test_get_openai_config(sample_request):
     assert isinstance(openai_config['messages'], list)
 
 
-def test_generate(sample_request):
+def test_generate(sample_request: GenerateRequest) -> None:
     """
     Test generate method calls OpenAI API and returns GenerateResponse.
     """
@@ -74,11 +83,13 @@ def test_generate(sample_request):
 
     mock_client.chat.completions.create.assert_called_once()
     assert isinstance(response, GenerateResponse)
+    assert response.message is not None
     assert response.message.role == Role.MODEL
+    assert response.message.content
     assert response.message.content[0].root.text == 'Hello, user!'
 
 
-def test_generate_stream(sample_request):
+def test_generate_stream(sample_request: GenerateRequest) -> None:
     """Test generate_stream method ensures it processes streamed responses correctly."""
     mock_client = MagicMock()
 
@@ -87,10 +98,10 @@ def __init__(self, data: list[str]) -> None:
             self._data = data
             self._current = 0
 
-        def __iter__(self):
+        def __iter__(self) -> Self:
             return self
 
-        def __next__(self):
+        def __next__(self) -> MagicMock:
             if self._current >= len(self._data):
                 raise StopIteration
 
@@ -112,7 +123,7 @@ def __next__(self):
     model = OpenAIModel(model=GPT_4, client=mock_client, registry=MagicMock())
     collected_chunks = []
 
-    def callback(chunk: GenerateResponseChunk):
+    def callback(chunk: GenerateResponseChunk) -> None:
         collected_chunks.append(chunk.content[0].root.text)
 
     model.generate_stream(sample_request, callback)

py/plugins/compat-oai/tests/test_tool_calling.py

@@ -22,7 +22,7 @@
 
 from genkit.plugins.compat_oai.models import OpenAIModel
 from genkit.plugins.compat_oai.models.model_info import GPT_4
-from genkit.types import GenerateRequest, GenerateResponseChunk, TextPart, ToolRequestPart
+from genkit.core.typing import GenerateRequest, GenerateResponseChunk, TextPart, ToolRequestPart
 
 
 def test_generate_with_tool_calls_executes_tools(sample_request: GenerateRequest) -> None:
@@ -143,4 +143,5 @@ def callback(chunk: GenerateResponseChunk):
     assert tool_part.tool_request.ref == 'tool123'
 
     accumulated_output = reduce(lambda res, tool_call: res + tool_call.tool_request.input, collected_chunks, '')
+    assert accumulated_output is not None
     assert json.loads(accumulated_output) == {'a': 1}