fix(py/genkit): additional ty check fixes in genkit.blocks (#4246)

Author: yesudeep

py/packages/genkit/src/genkit/blocks/prompt.py

@@ -22,9 +22,9 @@
 generation and management across different parts of the application.
 """
 
+import asyncio
 import os
 import weakref
-from asyncio import Future
 from collections.abc import AsyncIterator, Callable
 from pathlib import Path
 from typing import Any
@@ -240,7 +240,7 @@ def stream(
         timeout: float | None = None,
     ) -> tuple[
         AsyncIterator[GenerateResponseChunkWrapper],
-        Future[GenerateResponseWrapper],
+        asyncio.Future[GenerateResponseWrapper],
     ]:
         """Streams the prompt with the given input and configuration.
 
@@ -262,7 +262,7 @@ def stream(
             context=context,
             on_chunk=lambda c: stream.send(c),
         )
-        stream.set_close_future(resp)
+        stream.set_close_future(asyncio.create_task(resp))
 
         return (stream, stream.closed)
 

py/packages/genkit/src/genkit/blocks/resource.py

@@ -25,7 +25,7 @@
 import inspect
 import re
 from collections.abc import Awaitable, Callable
-from typing import Any, Protocol, TypedDict
+from typing import Any, Protocol, TypedDict, cast
 
 from pydantic import BaseModel
 
@@ -88,6 +88,12 @@ def __call__(self, input: ResourceInput, ctx: ActionRunContext) -> Awaitable[Res
         ...
 
 
+class MatchableAction(Protocol):
+    """Protocol for actions that have a matches method."""
+
+    matches: Callable[[ResourceInput], bool]
+
+
 ResourceArgument = Action | str
 
 
@@ -136,9 +142,9 @@ async def lookup_resource_by_name(registry: Registry, name: str) -> Action:
         ValueError: If the resource cannot be found.
     """
     resource = (
-        await registry.resolve_action(ActionKind.RESOURCE, name)
-        or await registry.resolve_action(ActionKind.RESOURCE, f'/resource/{name}')
-        or await registry.resolve_action(ActionKind.RESOURCE, f'/dynamic-action-provider/{name}')
+        await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), name)
+        or await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), f'/resource/{name}')
+        or await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), f'/dynamic-action-provider/{name}')
     )
     if not resource:
         raise ValueError(f'Resource {name} not found')
@@ -161,7 +167,7 @@ def define_resource(registry: Registry, opts: ResourceOptions, fn: ResourceFn) -
     """
     action = dynamic_resource(opts, fn)
 
-    action.matches = create_matcher(opts.get('uri'), opts.get('template'))
+    cast(MatchableAction, action).matches = create_matcher(opts.get('uri'), opts.get('template'))
 
     # Mark as not dynamic since it's being registered
     action.metadata['dynamic'] = False
@@ -279,7 +285,7 @@ async def wrapped_fn(input_data: ResourceInput, ctx: ActionRunContext) -> Resour
 
     act = Action(
         name=name,
-        kind=ActionKind.RESOURCE,
+        kind=cast(ActionKind, ActionKind.RESOURCE),
         fn=wrapped_fn,
         metadata={
             'resource': {
@@ -291,7 +297,7 @@ async def wrapped_fn(input_data: ResourceInput, ctx: ActionRunContext) -> Resour
         description=opts.get('description'),
         span_metadata={'genkit:metadata:resource:uri': uri},
     )
-    act.matches = matcher
+    cast(MatchableAction, act).matches = matcher
     return act
 
 
@@ -385,23 +391,27 @@ async def find_matching_resource(
     """
     if dynamic_resources:
         for action in dynamic_resources:
-            if hasattr(action, 'matches') and action.matches(input_data):
+            if hasattr(action, 'matches') and cast(MatchableAction, action).matches(input_data):
                 return action
 
     # Try exact match in registry
-    resource = await registry.resolve_action(ActionKind.RESOURCE, input_data.uri)
+    resource = await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), input_data.uri)
     if resource:
         return resource
 
     # Iterate all resources to check for matches (e.g. templates)
     # This is less efficient but necessary for template matching if not optimized
-    resources = registry.get_actions_by_kind(ActionKind.RESOURCE) if hasattr(registry, 'get_actions_by_kind') else {}
+    resources = (
+        registry.get_actions_by_kind(cast(ActionKind, ActionKind.RESOURCE))
+        if hasattr(registry, 'get_actions_by_kind')
+        else {}
+    )
     if not resources and hasattr(registry, '_entries'):
         # Fallback for compatibility if registry instance is old (unlikely in this context)
-        resources = registry._entries.get(ActionKind.RESOURCE, {})
+        resources = registry._entries.get(cast(ActionKind, ActionKind.RESOURCE), {})
 
     for action in resources.values():
-        if hasattr(action, 'matches') and action.matches(input_data):
+        if hasattr(action, 'matches') and cast(MatchableAction, action).matches(input_data):
             return action
 
     return None

py/packages/genkit/src/genkit/blocks/retriever.py

@@ -22,8 +22,9 @@
 to accomplish a task.
 """
 
+import inspect
 from collections.abc import Callable
-from typing import Any, Generic, TypeVar
+from typing import Any, Awaitable, Generic, TypeVar, cast
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -34,8 +35,8 @@
 from genkit.core.typing import DocumentData, RetrieverResponse
 
 T = TypeVar('T')
-# type RetrieverFn[T] = Callable[[Document, T], RetrieverResponse]
-RetrieverFn = Callable[[Document, T], RetrieverResponse]
+# type RetrieverFn[T] = Callable[[Document, T], RetrieverResponse | Awaitable[RetrieverResponse]]
+RetrieverFn = Callable[[Document, T], RetrieverResponse | Awaitable[RetrieverResponse]]
 
 
 class Retriever(Generic[T]):
@@ -115,9 +116,8 @@ def retriever_action_metadata(
         retriever_metadata_dict['retriever']['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True)
 
     retriever_metadata_dict['retriever']['customOptions'] = options.config_schema if options.config_schema else None
-
     return ActionMetadata(
-        kind=ActionKind.RETRIEVER,
+        kind=cast(ActionKind, ActionKind.RETRIEVER),
         name=name,
         input_json_schema=to_json_schema(RetrieverRequest),
         output_json_schema=to_json_schema(RetrieverResponse),
@@ -191,7 +191,7 @@ def indexer_action_metadata(
     indexer_metadata_dict['indexer']['customOptions'] = options.config_schema if options.config_schema else None
 
     return ActionMetadata(
-        kind=ActionKind.INDEXER,
+        kind=cast(ActionKind, ActionKind.INDEXER),
         name=name,
         input_json_schema=to_json_schema(IndexerRequest),
         output_json_schema=to_json_schema(None),
@@ -222,18 +222,20 @@ async def wrapper(
         request: RetrieverRequest,
         ctx: Any,
     ) -> RetrieverResponse:
-        return await fn(request.query, request.options)
+        query = Document.from_document_data(request.query)
+        res = fn(query, request.options)
+        return await res if inspect.isawaitable(res) else res
 
     registry.register_action(
-        kind=ActionKind.RETRIEVER,
+        kind=cast(ActionKind, ActionKind.RETRIEVER),
         name=name,
         fn=wrapper,
         metadata=metadata.metadata,
         span_metadata=metadata.metadata,
     )
 
 
-IndexerFn = Callable[[list[Document], T], None]
+IndexerFn = Callable[[list[Document], T], None | Awaitable[None]]
 
 
 def define_indexer(
@@ -249,11 +251,13 @@ async def wrapper(
         request: IndexerRequest,
         ctx: Any,
     ) -> None:
-        docs = [Document.from_data(d) for d in request.documents]
-        await fn(docs, request.options)
+        docs = [Document.from_document_data(d) for d in request.documents]
+        res = fn(docs, request.options)
+        if inspect.isawaitable(res):
+            await res
 
     registry.register_action(
-        kind=ActionKind.INDEXER,
+        kind=cast(ActionKind, ActionKind.INDEXER),
         name=name,
         fn=wrapper,
         metadata=metadata.metadata,

py/packages/genkit/src/genkit/blocks/tools.py

@@ -21,10 +21,17 @@
 allowing for controlled interruptions and specific response formatting.
 """
 
-from typing import Any
+from typing import Any, Protocol, cast
 
 from genkit.core.action import ActionRunContext
-from genkit.core.typing import Metadata, Part, ToolRequestPart, ToolResponse
+from genkit.core.typing import Metadata, Part, ToolRequestPart, ToolResponse, ToolResponsePart
+
+
+class ToolRequestLike(Protocol):
+    """Protocol for objects that look like a ToolRequest."""
+
+    name: str
+    ref: str | None
 
 
 class ToolRunContext(ActionRunContext):
@@ -71,13 +78,13 @@ class ToolInterruptError(Exception):
     causing a hard failure.
     """
 
-    def __init__(self, metadata: dict[str, Any]):
+    def __init__(self, metadata: dict[str, Any] | None = None):
         """Initializes the ToolInterruptError.
 
         Args:
             metadata: Metadata associated with the interruption.
         """
-        self.metadata = metadata
+        self.metadata = metadata or {}
 
 
 def tool_response(
@@ -110,13 +117,18 @@ def tool_response(
     elif metadata:
         interrupt_metadata = metadata
 
+    tr = cast(ToolRequestLike, tool_request)
     return Part(
-        tool_response=ToolResponse(
-            name=tool_request.name,
-            ref=tool_request.ref,
-            output=response_data,
-        ),
-        metadata={
-            'interruptResponse': interrupt_metadata,
-        },
+        root=ToolResponsePart(
+            toolResponse=ToolResponse(
+                name=tr.name,
+                ref=tr.ref,
+                output=response_data,
+            ),
+            metadata=Metadata(
+                root={
+                    'interruptResponse': interrupt_metadata,
+                }
+            ),
+        )
     )

py/packages/genkit/src/genkit/core/schema.py

@@ -21,7 +21,7 @@
 from pydantic import TypeAdapter
 
 
-def to_json_schema(schema: type | dict[str, Any]) -> dict[str, Any]:
+def to_json_schema(schema: type | dict[str, Any] | None) -> dict[str, Any]:
     """Converts a Python type to a JSON schema.
 
     If the input `schema` is already a dictionary (assumed json schema), it is