Improve Type Hints for Job-Specific Callback Pattern in Module Architecture

[thibaud-perrin]

📝 Description of the Issue

Currently, the job_specific_callback function and the run method in the module architecture use generic type hints, but they do not enforce or propagate the correct output type (OutputModelT) through the callback signature. This leads to situations where, in concrete module implementations (e.g., ExampleModule), the type checker (MyPy, Pyright, etc.) does not recognize that the callback parameter should specifically accept an AgentOutput instance.

The objective is to refactor the type hints and callback wrapping pattern so that:

• The callback signature is always correctly inferred and enforced in all concrete module implementations.
• The wrapper generated by job_specific_callback is properly typed, so that type checkers can guarantee type safety when passing callbacks to the run method.

🚀 Objective

• Ensure that in any subclass of BaseModule, the callback parameter of run and  _run_lifecycle is correctly typed to accept the module's specific output type.
• Refactor the job_specific_callback function to propagate type information, making it easier and safer to use in all modules.

🐞 Steps to Reproduce the Problem

1. Define a generic BaseModule and a job_specific_callback function as in the current codebase.
2. Test an implementation with a concrete module that specifies concrete types for input and output.
3. Attempt to pass a job-specific callback to the run method.
4. Observe that type checkers do not enforce the correct type for the callback parameter, potentially allowing runtime errors.

📊 Context (Environment)

• OS: Any
• Python version: 3.9+
• Type checker: MyPy, Pyright, or similar
• Libraries: typing, abc, functools, pydantic (for example models)

💰 Business Impact

Ensuring strict and correct type hints in the callback mechanism leads to:

• Fewer runtime bugs related to type mismatches
• Improved developer experience (better IDE support, auto-completion, and error detection)
• Easier onboarding for new developers
• Increased maintainability and robustness of the codebase

📝 Task List

• Refactor BaseModule.run and BaseModule._run_lifecycle to use Callable[[OutputModelT], Awaitable[None]] for the callback parameter.
• Refactor job_specific_callback to accept and return properly typed callables.
• Update all relevant docstrings and documentation.

💡 Potential Solution

Type-Safe Implementation Example

from abc import ABC, abstractmethod
from typing import (
    TypeVar,
    Generic,
    Callable,
    Awaitable,
    Any,
)
from functools import wraps
from pydantic import BaseModel

# Define type variables for generics
InputModelT = TypeVar("InputModelT")
OutputModelT = TypeVar("OutputModelT")
SetupModelT = TypeVar("SetupModelT")
SecretModelT = TypeVar("SecretModelT")
ConfigSetupModelT = TypeVar("ConfigSetupModelT")

class BaseModule(
    ABC,
    Generic[
        InputModelT,
        OutputModelT,
        SetupModelT,
        SecretModelT,
        ConfigSetupModelT,
    ],
):
    @abstractmethod
    async def run(
        self,
        input_data: InputModelT,
        setup_data: SetupModelT,
        callback: Callable[[OutputModelT], Awaitable[None]],
    ) -> None:
        """Run the module."""
        ...

    @staticmethod
    def job_specific_callback(
        callback: Callable[[str, OutputModelT], Awaitable[None]],
        job_id: str
    ) -> Callable[[OutputModelT], Awaitable[None]]:
        """
        Returns a wrapper that pre-fixes job_id in the call.
        This allows passing to `run` a callback that only expects OutputModelT.
        """
        @wraps(callback)
        def callback_wrapper(output_data: OutputModelT) -> Awaitable[None]:
            return callback(job_id, output_data)
        return callback_wrapper

class ArchetypeModule(
    BaseModule[
        InputModelT,
        OutputModelT,
        SetupModelT,
        SecretModelT,
        ConfigSetupModelT,
    ],
    ABC,
):
    """ArchetypeModule extends BaseModule for concrete module types."""
    ...

# Example Pydantic models
class AgentInput(BaseModel):
    query: str

class AgentOutput(BaseModel):
    answer: str

class AgentSetup(BaseModel):
    max_results: int

class AgentSecret(BaseModel):
    api_key: str

class ExampleModule(
    ArchetypeModule[
        AgentInput,
        AgentOutput,
        AgentSetup,
        AgentSecret,
        None,
    ]
):
    name = "ExampleArchetype"
    description = "This agent specializes in managing document creation based on user requirements."
    input_format = AgentInput
    output_format = AgentOutput
    setup_format = AgentSetup
    secret_format = AgentSecret

    async def run(
        self,
        input_data: AgentInput,
        setup_data: AgentSetup,
        callback: Callable[[AgentOutput], Awaitable[None]],
    ) -> None:
        # ... your async logic here ...
        result = AgentOutput(answer="Research completed")
        await callback(result)

# Usage Example

async def my_global_callback(job_id: str, output: AgentOutput) -> None:
    print(f"[{job_id}] got → {output.json()}")

async def main():
    module = ExampleModule()
    job_id = "job_123"
    cb = BaseModule.job_specific_callback(my_global_callback, job_id)
    await module.run(
        input_data=AgentInput(query="Why is the sky blue?"),
        setup_data=AgentSetup(max_results=5),
        callback=cb,
    )
# Run with: asyncio.run(main())

Key Explanations

• TypeVar("OutputModelT") is used to propagate the output type throughout the module and callback signatures.
• BaseModule.run and BaseModule._run_lifecycle now strictly accepts a Callable[[OutputModelT], Awaitable[None]], ensuring type safety.
• job_specific_callback wraps a callback expecting (str, OutputModelT) and returns one expecting only OutputModelT, with correct typing.
• In concrete modules (like ExampleModule), the type checker knows that callback expects an AgentOutput.

📈 Priority

This issue improves type safety and developer experience, reducing the risk of runtime errors and improving maintainability.