fix(py/genkit): ty check fixes for genkit.ai

Author: yesudeep

py/packages/genkit/src/genkit/ai/_aio.py

@@ -20,11 +20,11 @@
 class while customizing it with any plugins.
 """
 
+import asyncio
 import uuid
-from asyncio import Future
 from collections.abc import AsyncIterator
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 from genkit.aio import Channel
 from genkit.blocks.document import Document
@@ -253,7 +253,7 @@ def generate_stream(
         messages: list[Message] | None = None,
         tools: list[str] | None = None,
         return_tool_requests: bool | None = None,
-        tool_choice: ToolChoice = None,
+        tool_choice: ToolChoice | None = None,
         config: GenerationCommonConfig | dict[str, Any] | None = None,
         max_turns: int | None = None,
         context: dict[str, Any] | None = None,
@@ -268,7 +268,7 @@ def generate_stream(
         timeout: float | None = None,
     ) -> tuple[
         AsyncIterator[GenerateResponseChunkWrapper],
-        Future[GenerateResponseWrapper],
+        asyncio.Future[GenerateResponseWrapper],
     ]:
         """Streams generated text or structured data using a language model.
 
@@ -351,7 +351,7 @@ def generate_stream(
             use=use,
             on_chunk=lambda c: stream.send(c),
         )
-        stream.set_close_future(resp)
+        stream.set_close_future(asyncio.create_task(resp))
 
         return stream, stream.closed
 
@@ -389,7 +389,7 @@ async def retrieve(
 
         request_options = {**(retriever_config or {}), **(options or {})}
 
-        retrieve_action = await self.registry.resolve_action(ActionKind.RETRIEVER, retriever_name)
+        retrieve_action = await self.registry.resolve_action(cast(ActionKind, ActionKind.RETRIEVER), retriever_name)
         if retrieve_action is None:
             raise ValueError(f'Retriever "{retriever_name}" not found')
 
@@ -430,7 +430,7 @@ async def index(
 
         req_options = {**(indexer_config or {}), **(options or {})}
 
-        index_action = await self.registry.resolve_action(ActionKind.INDEXER, indexer_name)
+        index_action = await self.registry.resolve_action(cast(ActionKind, ActionKind.INDEXER), indexer_name)
         if index_action is None:
             raise ValueError(f'Indexer "{indexer_name}" not found')
 
@@ -464,7 +464,7 @@ async def embed(
         # Merge options passed to embed() with config from EmbedderRef
         final_options = {**(embedder_config or {}), **(options or {})}
 
-        embed_action = await self.registry.resolve_action(ActionKind.EMBEDDER, embedder_name)
+        embed_action = await self.registry.resolve_action(cast(ActionKind, ActionKind.EMBEDDER), embedder_name)
         if embed_action is None:
             raise ValueError(f'Embedder "{embedder_name}" not found')
 
@@ -501,7 +501,7 @@ async def evaluate(
 
         final_options = {**(evaluator_config or {}), **(options or {})}
 
-        eval_action = await self.registry.resolve_action(ActionKind.EVALUATOR, evaluator_name)
+        eval_action = await self.registry.resolve_action(cast(ActionKind, ActionKind.EVALUATOR), evaluator_name)
         if eval_action is None:
             raise ValueError(f'Evaluator "{evaluator_name}" not found')
 

py/packages/genkit/src/genkit/ai/_base_async.py

@@ -246,6 +246,6 @@ def _make_reflection_server(registry: GenkitRegistry, spec: ServerSpec) -> uvico
     Returns:
         A uvicorn server instance.
     """
-    app = create_reflection_asgi_app(registry=registry)
+    app = create_reflection_asgi_app(registry=registry.registry)
     config = uvicorn.Config(app, host=spec.host, port=spec.port, loop='asyncio')
     return uvicorn.Server(config)

py/packages/genkit/src/genkit/ai/_registry.py

@@ -43,7 +43,7 @@
 import uuid
 from collections.abc import AsyncIterator, Callable
 from functools import wraps
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 if TYPE_CHECKING:
     from genkit.blocks.resource import ResourceFn, ResourceOptions
@@ -122,7 +122,7 @@ def __init__(self):
         """Initialize the Genkit registry."""
         self.registry: Registry = Registry()
 
-    def flow(self, name: str | None = None, description: str | None = None) -> Callable[[Callable], Callable]:
+    def flow(self, name: str | None = None, description: str | None = None) -> Action:
         """Decorator to register a function as a flow.
 
         Args:
@@ -235,7 +235,7 @@ def define_schema(self, name: str, schema: type) -> type:
         define_schema(self.registry, name, schema)
         return schema
 
-    def tool(self, name: str | None = None, description: str | None = None) -> Callable[[Callable], Callable]:
+    def tool(self, name: str | None = None, description: str | None = None) -> Action:
         """Decorator to register a function as a tool.
 
         Args:
@@ -315,10 +315,10 @@ def define_retriever(
         self,
         name: str,
         fn: RetrieverFn,
-        config_schema: BaseModel | dict[str, Any] | None = None,
+        config_schema: type[BaseModel] | dict[str, Any] | None = None,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
-    ) -> Callable[[Callable], Callable]:
+    ) -> Action:
         """Define a retriever action.
 
         Args:
@@ -339,7 +339,7 @@ def define_retriever(
         retriever_description = get_func_description(fn, description)
         return self.registry.register_action(
             name=name,
-            kind=ActionKind.RETRIEVER,
+            kind=cast(ActionKind, ActionKind.RETRIEVER),
             fn=fn,
             metadata=retriever_meta,
             description=retriever_description,
@@ -349,10 +349,10 @@ def define_indexer(
         self,
         name: str,
         fn: IndexerFn,
-        config_schema: BaseModel | dict[str, Any] | None = None,
+        config_schema: type[BaseModel] | dict[str, Any] | None = None,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
-    ) -> Callable[[Callable], Callable]:
+    ) -> Action:
         """Define an indexer action.
 
         Args:
@@ -374,7 +374,7 @@ def define_indexer(
         indexer_description = get_func_description(fn, description)
         return self.registry.register_action(
             name=name,
-            kind=ActionKind.INDEXER,
+            kind=cast(ActionKind, ActionKind.INDEXER),
             fn=fn,
             metadata=indexer_meta,
             description=indexer_description,
@@ -384,7 +384,7 @@ def define_reranker(
         self,
         name: str,
         fn: RerankerFn,
-        config_schema: BaseModel | dict[str, Any] | None = None,
+        config_schema: type[BaseModel] | dict[str, Any] | None = None,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
     ) -> Action:
@@ -482,7 +482,7 @@ def define_evaluator(
         definition: str,
         fn: EvaluatorFn,
         is_billed: bool = False,
-        config_schema: BaseModel | dict[str, Any] | None = None,
+        config_schema: type[BaseModel] | dict[str, Any] | None = None,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
     ) -> Action:
@@ -581,7 +581,7 @@ async def eval_stepper_fn(req: EvalRequest) -> EvalResponse:
 
         return self.registry.register_action(
             name=name,
-            kind=ActionKind.EVALUATOR,
+            kind=cast(ActionKind, ActionKind.EVALUATOR),
             fn=eval_stepper_fn,
             metadata=evaluator_meta,
             description=evaluator_description,
@@ -594,10 +594,10 @@ def define_batch_evaluator(
         definition: str,
         fn: BatchEvaluatorFn,
         is_billed: bool = False,
-        config_schema: BaseModel | dict[str, Any] | None = None,
+        config_schema: type[BaseModel] | dict[str, Any] | None = None,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
-    ) -> Callable[[Callable], Callable]:
+    ) -> Action:
         """Define a batch evaluator action.
 
         This action runs the callback function on the entire dataset.
@@ -627,7 +627,7 @@ def define_batch_evaluator(
         evaluator_description = get_func_description(fn, description)
         return self.registry.register_action(
             name=name,
-            kind=ActionKind.EVALUATOR,
+            kind=cast(ActionKind, ActionKind.EVALUATOR),
             fn=fn,
             metadata=evaluator_meta,
             description=evaluator_description,
@@ -666,7 +666,7 @@ def define_model(
         model_description = get_func_description(fn, description)
         return self.registry.register_action(
             name=name,
-            kind=ActionKind.MODEL,
+            kind=cast(ActionKind, ActionKind.MODEL),
             fn=fn,
             metadata=model_meta,
             description=model_description,
@@ -706,7 +706,7 @@ def define_embedder(
         embedder_description = get_func_description(fn, description)
         return self.registry.register_action(
             name=name,
-            kind=ActionKind.EMBEDDER,
+            kind=cast(ActionKind, ActionKind.EMBEDDER),
             fn=fn,
             metadata=embedder_meta,
             description=embedder_description,