fix(py/genkit): additional ty check fixes in core (#4243)

Author: yesudeep

py/bin/sanitize_schema_typing.py

@@ -340,7 +340,7 @@ def add_header(content: str) -> str:
 import sys # noqa
 
 if sys.version_info < (3, 11):  # noqa
-    from strenum import StrEnum  # noqa
+    from strenum import StrEnum  # type: ignore # noqa
 else: # noqa
     from enum import StrEnum  # noqa
 """

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

@@ -89,7 +89,7 @@
 from collections.abc import AsyncIterator, Callable
 from contextvars import ContextVar
 from functools import cached_property
-from typing import Any
+from typing import Any, Awaitable
 
 from pydantic import BaseModel, TypeAdapter
 
@@ -228,7 +228,7 @@ def __init__(
         input_spec = inspect.getfullargspec(metadata_fn if metadata_fn else fn)
         action_args, arg_types = extract_action_args_and_types(input_spec)
         n_action_args = len(action_args)
-        self._fn, self._afn = _make_tracing_wrappers(name, kind, span_metadata, n_action_args, fn)
+        self._fn, self._afn = _make_tracing_wrappers(name, kind, span_metadata or {}, n_action_args, fn)
         self._initialize_io_schemas(action_args, arg_types, input_spec)
 
     @property
@@ -248,7 +248,7 @@ def metadata(self) -> dict[str, Any]:
         return self._metadata
 
     @cached_property
-    def input_type(self) -> type | None:
+    def input_type(self) -> TypeAdapter[Any] | None:
         return self._input_type
 
     @cached_property
@@ -418,7 +418,7 @@ def stream(
             telemetry_labels=telemetry_labels,
             on_chunk=lambda c: stream.send(c),
         )
-        stream.set_close_future(resp)
+        stream.set_close_future(asyncio.create_task(resp))
 
         result_future: asyncio.Future[ActionResponse] = asyncio.Future()
         stream.closed.add_done_callback(lambda _: result_future.set_result(stream.closed.result().response))
@@ -483,7 +483,7 @@ class ActionMetadata(BaseModel):
 
 
 _SyncTracingWrapper = Callable[[Any | None, ActionRunContext], ActionResponse]
-_AsyncTracingWrapper = Callable[[Any | None, ActionRunContext], ActionResponse]
+_AsyncTracingWrapper = Callable[[Any | None, ActionRunContext], Awaitable[ActionResponse]]
 
 
 def _make_tracing_wrappers(
@@ -539,7 +539,7 @@ async def async_tracing_wrapper(input: Any | None, ctx: ActionRunContext) -> Act
                 ) from e
 
             record_output_metadata(span, output=output)
-            return ActionResponse(response=output, trace_id=trace_id)
+            return ActionResponse(response=output, traceId=trace_id)
 
     def sync_tracing_wrapper(input: Any | None, ctx: ActionRunContext) -> ActionResponse:
         """Wrap the function in a sync tracing wrapper.
@@ -580,6 +580,6 @@ def sync_tracing_wrapper(input: Any | None, ctx: ActionRunContext) -> ActionResp
                 ) from e
 
             record_output_metadata(span, output=output)
-            return ActionResponse(response=output, trace_id=trace_id)
+            return ActionResponse(response=output, traceId=trace_id)
 
     return sync_tracing_wrapper, async_tracing_wrapper

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

@@ -21,7 +21,7 @@
 from typing import Any
 
 if sys.version_info < (3, 11):
-    from strenum import StrEnum
+    from strenum import StrEnum  # type: ignore
 else:
     from enum import StrEnum
 

py/packages/genkit/src/genkit/core/environment.py

@@ -18,9 +18,10 @@
 
 import os
 import sys
+from typing import cast
 
 if sys.version_info < (3, 11):
-    from strenum import StrEnum
+    from strenum import StrEnum  # type: ignore
 else:
     from enum import StrEnum
 
@@ -64,8 +65,8 @@ def get_current_environment() -> GenkitEnvironment:
     """
     env = os.getenv(EnvVar.GENKIT_ENV)
     if env is None:
-        return GenkitEnvironment.PROD
+        return cast(GenkitEnvironment, GenkitEnvironment.PROD)
     try:
         return GenkitEnvironment(env)
     except ValueError:
-        return GenkitEnvironment.PROD
+        return cast(GenkitEnvironment, GenkitEnvironment.PROD)

py/packages/genkit/src/genkit/core/error.py

@@ -17,7 +17,7 @@
 """Base error classes and utilities for Genkit."""
 
 import traceback
-from typing import Any
+from typing import Any, cast
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -61,6 +61,8 @@ class HttpErrorWireFormat(BaseModel):
 class GenkitError(Exception):
     """Base error class for Genkit errors."""
 
+    status: StatusName
+
     def __init__(
         self,
         *,
@@ -81,12 +83,13 @@ def __init__(
             trace_id: A unique identifier for tracing the action execution.
             source: Optional source of the error.
         """
-        self.status = status
-        if not self.status and isinstance(cause, GenkitError):
-            self.status = cause.status
-
-        if not self.status:
-            self.status = 'INTERNAL'
+        if status:
+            temp_status: StatusName = status
+        elif isinstance(cause, GenkitError):
+            temp_status = cause.status
+        else:
+            temp_status = 'INTERNAL'
+        self.status = temp_status
 
         source_prefix = f'{source}: ' if source else ''
         super().__init__(f'{source_prefix}{self.status}: {message}')
@@ -129,7 +132,7 @@ def to_serializable(self) -> GenkitReflectionApiErrorWireFormat:
         # This error type is used by 3P authors with the field "details",
         # but the actual Callable protocol value is "details"
         return GenkitReflectionApiErrorWireFormat(
-            details=self.details,
+            details=GenkitReflectionApiDetailsWireFormat(**self.details) if self.details else None,
             code=StatusCodes[self.status].value,
             message=repr(self.cause) if self.cause else self.original_message,
         )
@@ -201,7 +204,7 @@ def get_reflection_json(error: Any) -> GenkitReflectionApiErrorWireFormat:
     return GenkitReflectionApiErrorWireFormat(
         message=str(error),
         code=StatusCodes.INTERNAL.value,
-        details={'stack': get_error_stack(error)},
+        details=GenkitReflectionApiDetailsWireFormat(stack=get_error_stack(error)),
     )
 
 

py/packages/genkit/src/genkit/core/flows.py

@@ -330,7 +330,7 @@ async def handle_standard_flow(
         routes=routes,
         middleware=[
             Middleware(
-                CORSMiddleware,
+                CORSMiddleware,  # type: ignore[arg-type]
                 allow_origins=['*'],
                 allow_methods=['*'],
                 allow_headers=['*'],

py/packages/genkit/src/genkit/core/plugin.py

@@ -21,6 +21,7 @@
 """
 
 import abc
+from typing import cast
 
 from genkit.core.action import Action, ActionMetadata
 from genkit.core.action.types import ActionKind
@@ -92,7 +93,7 @@ async def model(self, name: str) -> Action:
             ValueError: If the model is not found.
         """
         target = name if '/' in name else f'{self.name}/{name}'
-        action = await self.resolve(ActionKind.MODEL, target)
+        action = await self.resolve(cast(ActionKind, ActionKind.MODEL), target)
         if action is None:
             raise ValueError(f'Model not found: {target}')
         return action

py/packages/genkit/src/genkit/core/reflection.py

@@ -45,7 +45,7 @@
 import urllib.parse
 from collections.abc import AsyncGenerator, Callable
 from http.server import BaseHTTPRequestHandler
-from typing import Any
+from typing import Any, cast
 
 import structlog
 from starlette.applications import Starlette
@@ -58,7 +58,7 @@
 from genkit.aio.loop import run_async
 from genkit.codec import dump_dict, dump_json
 from genkit.core.action import Action
-from genkit.core.action.types import ActionKind
+from genkit.core.action.types import ActionKind, ActionResponse
 from genkit.core.constants import DEFAULT_GENKIT_VERSION
 from genkit.core.error import get_reflection_json
 from genkit.core.registry import Registry
@@ -77,7 +77,7 @@
 def _list_registered_actions(registry: Registry) -> dict[str, Action]:
     """Return all locally registered actions keyed as `/<kind>/<name>`."""
     registered: dict[str, Action] = {}
-    for kind in ActionKind:
+    for kind in ActionKind.__members__.values():
         for name, action in registry.get_actions_by_kind(kind).items():
             registered[f'/{kind.value}/{name}'] = action
     return registered
@@ -179,7 +179,7 @@ def log_message(self, format, *args):
                 message = format % args
                 address = self.address_string()
                 timestamp = self.log_date_time_string()
-                control_chars = self._control_char_table
+                control_chars = getattr(self, '_control_char_table', {})
                 logger.debug(f'{address} - - [{timestamp}] {message.translate(control_chars)}')
 
         def do_GET(self) -> None:  # noqa: N802
@@ -246,6 +246,10 @@ async def get_action():
                     return await registry.resolve_action_by_key(payload['key'])
 
                 action = run_async(loop, get_action)
+                if not action:
+                    self.send_response(404)
+                    self.end_headers()
+                    return
                 payload.get('input')
                 context = payload['context'] if 'context' in payload else {}
 
@@ -281,7 +285,7 @@ async def run_fn():
                                 context=context,
                             )
 
-                        output = run_async(loop, run_fn)
+                        output = cast(ActionResponse, run_async(loop, run_fn))
 
                         self.wfile.write(
                             bytes(
@@ -309,7 +313,7 @@ async def run_fn():
                         async def run_fn():
                             return await action.arun_raw(raw_input=payload.get('input'), context=context)
 
-                        output = run_async(loop, run_fn)
+                        output = cast(ActionResponse, run_async(loop, run_fn))
 
                         self.send_response(200)
                         self.send_header('x-genkit-version', DEFAULT_GENKIT_VERSION)
@@ -719,7 +723,7 @@ def wrapped_on_trace_start(tid):
         ],
         middleware=[
             Middleware(
-                CORSMiddleware,
+                CORSMiddleware,  # type: ignore[arg-type]
                 allow_origins=['*'],
                 allow_methods=['*'],
                 allow_headers=['*'],
@@ -728,5 +732,5 @@ def wrapped_on_trace_start(tid):
         on_startup=[on_app_startup] if on_app_startup else [],
         on_shutdown=[on_app_shutdown] if on_app_shutdown else [],
     )
-    app.active_actions = active_actions
+    setattr(app, 'active_actions', active_actions)
     return app

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

@@ -30,7 +30,7 @@
 import asyncio
 import threading
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 import structlog
 from dotpromptz.dotprompt import Dotprompt
@@ -383,7 +383,11 @@ async def resolve_action(self, kind: ActionKind, name: str) -> Action | None:
         # Skip if we're looking up a dynamic action provider itself to avoid recursion
         if kind != ActionKind.DYNAMIC_ACTION_PROVIDER:
             with self._lock:
-                providers = list(self._entries.get(ActionKind.DYNAMIC_ACTION_PROVIDER, {}).values())
+                if ActionKind.DYNAMIC_ACTION_PROVIDER in self._entries:
+                    providers_dict = self._entries[ActionKind.DYNAMIC_ACTION_PROVIDER]
+                else:
+                    providers_dict = {}
+                providers = list(providers_dict.values())
             for provider in providers:
                 try:
                     response = await provider.arun({'kind': kind, 'name': name})

py/packages/genkit/src/genkit/core/trace/default_exporter.py

@@ -51,12 +51,16 @@ def extract_span_data(span: ReadableSpan) -> dict[str, Any]:
     This function extracts the span data from a ReadableSpan object and returns
     a dictionary containing the span data.
     """
-    span_data = {'traceId': f'{span.context.trace_id}', 'spans': {}}
-    span_data['spans'][span.context.span_id] = {
-        'spanId': f'{span.context.span_id}',
+    span_data: dict[str, Any] = {'traceId': f'{span.context.trace_id}', 'spans': {}}
+    span_id = span.context.span_id
+    start_time = (span.start_time / 1000000) if span.start_time is not None else 0
+    end_time = (span.end_time / 1000000) if span.end_time is not None else 0
+
+    span_data['spans'][span_id] = {
+        'spanId': f'{span_id}',
         'traceId': f'{span.context.trace_id}',
-        'startTime': span.start_time / 1000000,
-        'endTime': span.end_time / 1000000,
+        'startTime': start_time,
+        'endTime': end_time,
         'attributes': {**span.attributes},
         'displayName': span.name,
         # "links": span.links,
@@ -75,8 +79,8 @@ def extract_span_data(span: ReadableSpan) -> dict[str, Any]:
             'version': 'v1',
         },
     }
-    if not span_data['spans'][span.context.span_id]['parentSpanId']:  # type: ignore
-        del span_data['spans'][span.context.span_id]['parentSpanId']  # type: ignore
+    if not span_data['spans'][span.context.span_id]['parentSpanId']:
+        del span_data['spans'][span.context.span_id]['parentSpanId']
 
     if not span.parent:
         span_data['displayName'] = span.name
@@ -127,7 +131,7 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
             for span in spans:
                 client.post(
                     urljoin(self.telemetry_server_url, self.telemetry_server_endpoint),
-                    data=json.dumps(extract_span_data(span)),
+                    json=extract_span_data(span),
                     headers={
                         'Content-Type': 'application/json',
                         'Accept': 'application/json',