Fixed mypy errors

py/bin/sanitize_schema_types.py

@@ -87,7 +87,7 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:  # noqa: N802
         type."""
         if self.is_rootmodel_class(node):
             # Filter out model_config assignments for RootModel classes
-            new_body = []
+            new_body: list[ast.stmt] = []
             for item in node.body:
                 if isinstance(item, ast.Assign):
                     targets = item.targets

py/packages/genkit/src/genkit/core/action.py

@@ -3,9 +3,8 @@
 
 import inspect
 import json
-from collections.abc import Callable
 from enum import Enum
-from typing import Any
+from typing import Any, Callable
 
 from genkit.core.tracing import tracer
 from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
@@ -62,7 +61,7 @@ def __init__(
         self,
         kind: ActionKind,
         name: str,
-        fn: Callable,
+        fn: Callable[..., Any],
         description: str | None = None,
         metadata: dict[ActionMetadataKey, Any] | None = None,
         span_metadata: dict[str, str] | None = None,
@@ -81,7 +80,7 @@ def __init__(
         self.kind: ActionKind = kind
         self.name = name
 
-        def tracing_wrapper(*args, **kwargs):
+        def tracing_wrapper(*args: Any, **kwargs: Any) -> ActionResponse:
             """Wraps the callable function in a tracing span and adds metadata
             to it."""
 
@@ -111,7 +110,7 @@ def tracing_wrapper(*args, **kwargs):
                 else:
                     span.set_attribute('genkit:output', json.dumps(output))
 
-                return ActionResponse(response=output, trace_id=trace_id)
+                return ActionResponse(response=output, traceId=trace_id)
 
         self.fn = tracing_wrapper
         self.description = description

py/packages/genkit/src/genkit/core/action_test.py

@@ -9,7 +9,10 @@
 def test_action_enum_behaves_like_str() -> None:
     """Ensure the ActionType behaves like a string and to ensure we're using the
     correct variants."""
-    assert ActionKind.CHATLLM == 'chat-llm'
+    assert isinstance(ActionKind.CHATLLM, str)
+    """Mypy treats enums as distinct types even when they inherit str.
+    Suppressed strict equality check for enums."""
+    assert ActionKind.CHATLLM == 'chat-llm' # type: ignore[comparison-overlap]
     assert ActionKind.CUSTOM == 'custom'
     assert ActionKind.EMBEDDER == 'embedder'
     assert ActionKind.EVALUATOR == 'evaluator'

py/packages/genkit/src/genkit/core/plugin_abc.py

@@ -6,11 +6,11 @@
 from __future__ import annotations
 
 import abc
-import typing
+from typing import TYPE_CHECKING, Any
 
 from genkit.core.schema_types import GenerateRequest, GenerateResponse
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from genkit.veneer import Genkit
 
 
@@ -38,7 +38,7 @@ def attach_to_veneer(self, veneer: Genkit) -> None:
         pass
 
     def _add_model_to_veneer(
-        self, veneer: Genkit, name: str, metadata: dict | None = None
+        self, veneer: Genkit, name: str, metadata: dict[str, Any] | None = None
     ) -> None:
         """
         Defines plugin's model in the Genkit Registry

py/packages/genkit/src/genkit/core/registry.py

@@ -35,6 +35,7 @@ def lookup_action(self, kind: ActionKind, name: str) -> Action | None:
         """
         if kind in self.actions and name in self.actions[kind]:
             return self.actions[kind][name]
+        return None
 
     def lookup_action_by_key(self, key: str) -> Action | None:
         """Lookup an action by its key.
@@ -60,5 +61,11 @@ def lookup_action_by_key(self, key: str) -> Action | None:
                 'Expected format: `<kind>/<name>`'
             )
             raise ValueError(msg)
-        kind, name = tokens
+        kind_str, name = tokens
+
+        try:
+            kind = ActionKind[kind_str]
+        except KeyError as err:
+            raise ValueError(f'Invalid action kind: `{kind_str}`') from err
+
         return self.lookup_action(kind, name)