Make Teleprompt.compile more typing friendly.

[TyTodd]

I replaced student: Module in Teleprompt.compile's signature with student: M, a TypeVar that tracks the type and returns that same type

M = TypeVar("M", bound=Module)
class Teleprompter:
...
    def compile(
        self,
        student: M,
        ...

I also applied this signature to all other compilers in the SDK.

It doesn't change the functionality of anything at all. My main reason for adding this is I (and others) often add custom functionality to dspy.Modules and its just convenient to have the type checker still recognize those functions and attributes on the optimized program after running compile.

For example, in this script the type checker will not highlight weather_function because the custom function is forgotten when the new program is returned from compile. With this new commit that is no longer the case.

class MyAgent(Module):
    def __init__(self):
        super().__init__()
        self.predict = dspy.ReAct("question -> answer", tools=[self.weather_function])

    def weather_function(self):
        return random.random()

optimizer = MIPROv2(metric=dspy.evaluate.answer_exact_match, auto="light")
compiled_agent = optimizer.compile(student=MyAgent(), trainset=trainset, teacher=teacher, valset=valset)
print(compiled_agent.weather_function())

[TomeHirata]

Thanks @TyTodd! Is it hard to have typing.cast() in your codebase after .compile()?

[Copilot]

Pull Request Overview

This pull request introduces improved type safety to the DSPy teleprompt system by adding generic type variables. The main changes include introducing TypeVar("M", bound=Module) to preserve the specific type of the student module through the compile method, replacing previous Module or Any return types.

• Added TypeVar imports and defined M = TypeVar("M", bound=Module) in multiple teleprompt files
• Updated compile method signatures to use generic type M instead of Module or Any for better type preservation
• Applied consistent code formatting (line breaks, spacing) across multiple files

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
dspy/teleprompt/teleprompt.py	Added TypeVar M and updated base Teleprompter.compile signature to use generic type
dspy/teleprompt/simba.py	Added TypeVar M, updated SIMBA.compile signature, and reformatted docstrings
dspy/teleprompt/mipro_optimizer_v2.py	Added TypeVar M, updated MIPROv2.compile signature, and reformatted long lines
dspy/teleprompt/grpo.py	Added TypeVar M, updated GRPO.compile signature, and reformatted long lines/assertions
dspy/teleprompt/bootstrap_finetune.py	Added TypeVar M and updated FinetuneTeleprompter.compile signature
dspy/teleprompt/bettertogether.py	Added TypeVar M, updated BetterTogether.compile signature, and reformatted code

Comments suppressed due to low confidence (3)

dspy/teleprompt/bettertogether.py:7

• The import from dspy.primitives.module import Module on line 7 is redundant since line 17 already imports from dspy.primitives import Module. One of these imports should be removed to avoid duplication.

from dspy.primitives.module import Module

dspy/teleprompt/bettertogether.py:90

• The return type of _run_strategies should be M instead of Module to maintain consistency with the generic type used in the compile method signature. This would ensure the type is preserved through internal method calls.

    def _run_strategies(self, parsed_strategy, student, trainset, valset_ratio) -> Module:

dspy/teleprompt/bettertogether.py:128

• The return type of _compile_prompt_optimizer should use the generic type M instead of Module for consistency with the type safety improvements introduced in this PR.

    def _compile_prompt_optimizer(self, student, trainset, valset_ratio) -> Module:

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Overridden method signature does not match call, where it is passed too many arguments. Overriding method method BootstrapFinetune.compile matches the call.
Overridden method signature does not match call, where it is passed too many arguments. Overriding method method BootstrapFinetune.compile matches the call.
Overridden method signature does not match call, where it is passed an argument named 'trainset'. Overriding method method BootstrapFinetune.compile matches the call.
Overridden method signature does not match call, where it is passed an argument named 'trainset'. Overriding method method BootstrapFinetune.compile matches the call.
Overridden method signature does not match call, where it is passed an argument named 'teacher'. Overriding method method BootstrapFinetune.compile matches the call.

[Copilot]

This method requires at least 3 positional arguments, whereas overridden Teleprompter.compile requires 2.

[Copilot]

This method requires at least 3 positional arguments, whereas overridden Teleprompter.compile requires 2.

Suggested change

	def compile(self, student: M, trainset: list[Example], teacher: Module | list[Module] | None = None) -> M:
	# TODO: Print statements can be converted to logger.info if we ensure
	# that the default DSPy logger logs info level messages in notebook
	# environments.
	def compile(self, student: M, trainset: list[Example], **kwargs) -> M:
	# TODO: Print statements can be converted to logger.info if we ensure
	# that the default DSPy logger logs info level messages in notebook
	# environments.
	teacher = kwargs.get('teacher', None)

[Copilot]

This method requires at least 3 positional arguments, whereas overridden Teleprompter.compile requires 2.

[TyTodd]

Thanks @TyTodd! Is it hard to have typing.cast() in your codebase after .compile()?

Thanks for your response! It's not really hard. It's more of a taste thing I guess.
Selfishly, I am building a library that extends DSPy and the type forwarding would just make things prettier on my end.

[TomeHirata]

Let's not include unrelated changes

[TyTodd]

Must be my formatter. Let me see if I can undo the other formatting changes.

[TomeHirata]

I see, using base class in type hint is very common but I also understand where you come from. @okhat @chenmoneygithub what do you think about using generic in the type hint of .compile?