feat(py): Add type safety to Python SDK #4309

[huangjeff5]

Summary

This PR adds full type safety to the Python SDK to match what we have in JS. The core goal: when you call generate(), define_prompt(), or use a flow, the return types should be known at dev time so your IDE can autocomplete and catch errors before runtime.

What was broken

Type information was getting lost at key boundaries:

# Types lost - IDE shows `Any` or `object`
response = await ai.generate(output_schema=Recipe)
response.output.ingredients  # No autocomplete, no error checking

result = await recipe_prompt(input)
[result.output.name](http://result.output.name/)  # Also untyped

What this PR does

Generic Action class

Made Action generic so types flow through:

class Action(Generic[InputT, OutputT, ChunkT]):
    async def arun(self, input: InputT) -> ActionResponse[OutputT]: ...

Typed Output[T] for generate()

from genkit import Output

class Recipe(BaseModel):
    name: str
    ingredients: list[str]

response = await ai.generate(
    prompt="Give me a pasta recipe",
    output=Output(schema=Recipe),
)

# Now these are typed
[response.output.name](http://response.output.name/)         # str
response.output.ingredients  # list[str]
response.output.typo         # IDE error - caught before runtime

Typed generate_stream()

Same pattern works for streaming:

stream, future = ai.generate_stream(
    prompt="Give me a recipe",
    output=Output(schema=Recipe),
)

response = await future
[response.output.name](http://response.output.name/)  # Typed as str

Typed prompts with Input[T] and Output[T]

Prompts (including dotprompt files) get full type safety:

from genkit import Input, Output

class RecipeInput(BaseModel):
    dish: str
    servings: int

class Recipe(BaseModel):
    name: str
    ingredients: list[str]

recipe_prompt = ai.define_prompt(
    name="recipe",
    prompt="Create a recipe for {{dish}} serving {{servings}} people",
    input=Input(schema=RecipeInput),
    output=Output(schema=Recipe),
)

# Input is type-checked
response = await recipe_prompt(RecipeInput(dish="pizza", servings=4))

# Output is typed
[response.output.name](http://response.output.name/)  # str
response.output.ingredients  # list[str]

# Errors caught at dev time
await recipe_prompt("pizza")  # Wrong type - IDE flags this
response.output.wrong_field   # No such attribute - IDE flags this

Typed flows

Flows preserve types through the decorator:

@ai.flow()
async def analyze_document(doc: Document) -> Analysis:
    ...

analysis = await analyze_document(my_doc)
analysis.summary  # IDE knows this exists

Breaking change

output_schema removed from generate(). Use output=Output(schema=...) instead:

# Before
response = await ai.generate(output_schema=Recipe, output_format="json")

# After
response = await ai.generate(output=Output(schema=Recipe))

Other changes

• Added typed logger wrapper
• Type-safe registry lookups
• Reduced Any usage across codebase
• Added @override decorators where needed
• Zero pyright errors in strict mode

Testing

• Pyright verification tests in tests/typing/ - all pass
• Negative tests verify errors are caught
• Existing unit tests pass
• Manually verified autocomplete in VS Code

[gemini-code-assist]

Summary of Changes

Hello @huangjeff5, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly upgrades the Python SDK's type safety, providing a more robust and developer-friendly experience. By introducing generic Input and Output classes, and extensively applying modern Python typing features, the changes enable static analysis tools to provide richer feedback, improve code maintainability, and reduce runtime errors. The refactoring touches core components like action definitions, flow wrappers, and generation methods, ensuring that type information is preserved and utilized effectively across the library.

Highlights

• Enhanced Type Safety: Introduced generic Input and Output classes for generate and define_prompt methods, enabling compile-time type checking of input and output schemas. This significantly improves type inference and reduces the need for Any.
• Comprehensive Type Annotations: Extensively applied modern Python typing features such as TypeVar, Generic, overload, ClassVar, and override across the SDK. Core components like Action, FlowWrapper, and Channel are now generic, preserving type information throughout the library.
• Improved Developer Experience: A new typing-guide.md documentation file has been added to explain how to leverage the new type safety features for better IDE autocomplete, inline documentation, and early error detection, fostering a more robust development workflow.
• Pyright Compatibility & Typed Logging: Specific adjustments and pyright: ignore comments were added to ensure compatibility and suppress warnings from the Pyright static type checker. Additionally, a new genkit.core.logging module with a Logger Protocol provides type-safe logging throughout the SDK.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Ignored Files

• Ignored by pattern: .github/workflows/** (1)
  • .github/workflows/python.yml

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This is an excellent and extensive pull request that brings comprehensive type safety to the Python SDK. The introduction of generic types for core components like Action, ExecutablePrompt, and FlowWrapper, along with the new Input[T] and Output[T] classes, is a significant improvement for developer experience and code correctness. The changes are consistently applied across the codebase, including updates to samples and the addition of typing verification tests. The new typed logger and pyright configuration are also great additions. I have one minor suggestion for code clarity.