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,19 +389,20 @@ 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')
 
         return (
             await retrieve_action.arun(
                 RetrieverRequest(
-                    query=query,
+                    query=query,  # type: ignore[arg-type]
                     options=request_options if request_options else None,
                 )
             )
         ).response
 
+
     async def index(
         self,
         indexer: str | IndexerRef | None = None,
@@ -430,13 +431,16 @@ 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')
 
+        if documents is None:
+            raise ValueError('Documents must be specified for indexing.')
+
         await index_action.arun(
             IndexerRequest(
-                documents=documents,
+                documents=documents,  # type: ignore[arg-type]
                 options=req_options if req_options else None,
             )
         )
@@ -464,11 +468,14 @@ 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')
 
-        return (await embed_action.arun(EmbedRequest(input=documents, options=final_options))).response
+        if documents is None:
+            raise ValueError('Documents must be specified for embedding.')
+
+        return (await embed_action.arun(EmbedRequest(input=documents, options=final_options))).response  # type: ignore[arg-type]
 
     async def evaluate(
         self,
@@ -501,19 +508,22 @@ 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')
 
         if not eval_run_id:
             eval_run_id = str(uuid.uuid4())
 
+        if dataset is None:
+            raise ValueError('Dataset must be specified for evaluation.')
+
         return (
             await eval_action.arun(
                 EvalRequest(
                     dataset=dataset,
                     options=final_options,
-                    eval_run_id=eval_run_id,
+                    evalRunId=eval_run_id,
                 )
             )
         ).response

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
@@ -144,11 +144,11 @@ def wrapper(func: Callable) -> Callable:
             Returns:
                 The wrapped function that executes the flow.
             """
-            flow_name = name if name is not None else func.__name__
+            flow_name = name if name is not None else getattr(func, '__name__', 'unnamed_flow')
             flow_description = get_func_description(func, description)
             action = self.registry.register_action(
                 name=flow_name,
-                kind=ActionKind.FLOW,
+                kind=cast(ActionKind, ActionKind.FLOW),
                 fn=func,
                 description=flow_description,
                 span_metadata={'genkit:metadata:flow:name': flow_name},
@@ -257,7 +257,7 @@ def wrapper(func: Callable) -> Callable:
             Returns:
                 The wrapped function that executes the tool.
             """
-            tool_name = name if name is not None else func.__name__
+            tool_name = name if name is not None else getattr(func, '__name__', 'unnamed_tool')
             tool_description = get_func_description(func, description)
 
             input_spec = inspect.getfullargspec(func)
@@ -275,7 +275,7 @@ def tool_fn_wrapper(*args):
 
             action = self.registry.register_action(
                 name=tool_name,
-                kind=ActionKind.TOOL,
+                kind=cast(ActionKind, ActionKind.TOOL),
                 description=tool_description,
                 fn=tool_fn_wrapper,
                 metadata_fn=func,
@@ -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:
@@ -426,7 +426,7 @@ def define_reranker(
             name=name,
             fn=fn,
             options=RerankerOptions(
-                config_schema=reranker_meta['reranker'].get('customOptions'),
+                configSchema=reranker_meta['reranker'].get('customOptions'),
                 label=reranker_meta['reranker'].get('label'),
             ),
             description=reranker_description,
@@ -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:
@@ -544,13 +544,13 @@ async def eval_stepper_fn(req: EvalRequest) -> EvalResponse:
                                 logger.debug(traceback.format_exc())
                                 evaluation = Score(
                                     error=f'Evaluation of test case {datapoint.test_case_id} failed: \n{str(e)}',
-                                    status=EvalStatusEnum.FAIL,
+                                    status=cast(EvalStatusEnum, EvalStatusEnum.FAIL),
                                 )
                                 eval_responses.append(
                                     EvalFnResponse(
-                                        span_id=span_id,
-                                        trace_id=trace_id,
-                                        test_case_id=datapoint.test_case_id,
+                                        spanId=span_id,
+                                        traceId=trace_id,
+                                        testCaseId=datapoint.test_case_id,
                                         evaluation=evaluation,
                                     )
                                 )
@@ -566,11 +566,11 @@ async def eval_stepper_fn(req: EvalRequest) -> EvalResponse:
                             logger.debug(traceback.format_exc())
                             evaluation = Score(
                                 error=f'Evaluation of test case {datapoint.test_case_id} failed: \n{str(e)}',
-                                status=EvalStatusEnum.FAIL,
+                                status=cast(EvalStatusEnum, EvalStatusEnum.FAIL),
                             )
                             eval_responses.append(
                                 EvalFnResponse(
-                                    test_case_id=datapoint.test_case_id,
+                                    testCaseId=datapoint.test_case_id,
                                     evaluation=evaluation,
                                 )
                             )
@@ -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,

py/packages/genkit/src/genkit/blocks/evaluator.py

@@ -21,7 +21,7 @@
 model.
 """
 
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable, Coroutine
 from typing import Any, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -35,12 +35,11 @@
 T = TypeVar('T')
 
 # User-provided evaluator function that evaluates a single datapoint.
-# type EvaluatorFn[T] = Callable[[BaseDataPoint, T], EvalFnResponse]
-EvaluatorFn = Callable[[BaseDataPoint, T], EvalFnResponse]
+# Must be async (coroutine function).
+EvaluatorFn = Callable[[BaseDataPoint, T], Coroutine[Any, Any, EvalFnResponse]]
 
 # User-provided batch evaluator function that evaluates an EvaluationRequest
-# type BatchEvaluatorFn[T] = Callable[[EvalRequest, T], list[EvalFnResponse]]
-BatchEvaluatorFn = Callable[[EvalRequest, T], list[EvalFnResponse]]
+BatchEvaluatorFn = Callable[[EvalRequest, T], Coroutine[Any, Any, list[EvalFnResponse]]]
 
 
 class EvaluatorRef(BaseModel):
@@ -62,4 +61,4 @@ def evaluator_ref(name: str, config_schema: Any | None = None) -> EvaluatorRef:
     Returns:
         An EvaluatorRef instance.
     """
-    return EvaluatorRef(name=name, config_schema=config_schema)
+    return EvaluatorRef(name=name, configSchema=config_schema)