feat(py): Return pydantic instances when output schema passed in as P…

[huangjeff5]

This PR adds typed input and output support to the Python Genkit SDK, enabling static type inference and runtime Pydantic validation for prompts.

• Add input and output parameters to ai.prompt() with proper overloads for full type inference
• Store and propagate schema_type through the generate pipeline for runtime Pydantic validation
• When output.schema is a Pydantic model, response.output now returns a validated Pydantic instance instead of a raw dict
• Export Input alongside Output from the top-level genkit module
• Add lookup_schema_type() to registry for resolving named schemas to their Pydantic types

API changes

• ai.prompt('name', input=Input(schema=MyInput), output=Output(schema=MyOutput)) returns
• ExecutablePrompt[MyInput, MyOutput] with proper type inference
• Supports all combinations: neither typed, input-only, output-only, or both
• response.output returns a validated Pydantic instance instead of dict

[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 enhances the Python Genkit SDK by integrating robust type-checking and Pydantic validation for prompt inputs and outputs. These changes allow developers to define schemas for prompt data, enabling static type inference at design time and ensuring data integrity through runtime validation. The update aims to improve code reliability, maintainability, and overall developer experience by providing clearer data contracts and automated validation for AI prompt interactions.

Highlights

• Typed Input and Output Support: Introduced typed input and output parameters to the ai.prompt() function, enabling static type inference for prompt configurations.
• Pydantic Validation for Outputs: Implemented runtime Pydantic validation for prompt outputs. When an output schema is provided as a Pydantic model, response.output now returns a validated Pydantic instance instead of a raw dictionary.
• Schema Type Propagation: The schema_type (Pydantic model class) is now stored and propagated through the generate pipeline, facilitating runtime validation.
• Registry Enhancements: Added lookup_schema_type() to the registry to resolve named schemas to their corresponding Pydantic types, and updated define_schema to store these types.
• Module Export: The Input class is now exported from the top-level genkit module for easier access.

🧠 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.

Changelog

• py/packages/genkit/src/genkit/init.py
  • Exported Input alongside Output from the top-level genkit module.
• py/packages/genkit/src/genkit/ai/_aio.py
  • Removed the Input = _Input alias, directly importing Input from genkit.blocks.interfaces.
• py/packages/genkit/src/genkit/ai/_registry.py
  • Updated define_schema to accept and store Pydantic BaseModel types.
  • Added multiple @overload definitions for the prompt method to support typed input and output parameters, enabling static type inference.
  • Modified the prompt method to extract input_schema and output_schema and pass them to ExecutablePrompt.
• py/packages/genkit/src/genkit/blocks/generate.py
  • Modified the generate function to extract schema_type from the raw request and pass it to GenerateResponseWrapper for runtime validation.
• py/packages/genkit/src/genkit/blocks/model.py
  • Added _schema_type as a private attribute to GenerateResponseWrapper to store the Pydantic class for runtime validation.
  • Updated the __init__ method of GenerateResponseWrapper to accept and store schema_type.
  • Modified the output property to validate parsed output against _schema_type using model_validate if a Pydantic schema is provided.
• py/packages/genkit/src/genkit/blocks/prompt.py
  • Enhanced render function to handle output_schema as a string (named schema), a Pydantic type, or a raw dict, storing both JSON schema and the Pydantic schema_type.
  • Updated to_generate_action_options to also look up and store schema_type when output_schema is provided as a string or a Pydantic type.
  • Modified define_schema to pass the Pydantic schema type to registry.register_schema.
• py/packages/genkit/src/genkit/core/registry.py
  • Imported BaseModel from pydantic.
  • Added _schema_types_by_name dictionary to store Pydantic BaseModel types by name.
  • Updated register_schema to accept an optional schema_type and store it in _schema_types_by_name.
  • Added lookup_schema_type method to retrieve registered Pydantic BaseModel types by name.
• py/packages/genkit/src/genkit/core/typing.py
  • Added schema_type field to GenerateActionOutputConfig to store the original Pydantic type for runtime validation, excluded from JSON serialization.

Activity

• This pull request was opened by huangjeff5.

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 pull request introduces excellent support for typed inputs and outputs using Pydantic models in the Python SDK. The use of overloads for ai.prompt provides great type inference for developers. The propagation of schema_type for runtime validation is well-implemented, culminating in the GenerateResponseWrapper.output property which now returns validated Pydantic instances. The changes are comprehensive and significantly improve the developer experience. I have one suggestion regarding code duplication to enhance maintainability.

[gemini-code-assist]

Code Review

This pull request introduces typed input and output schemas for prompts, enabling static type checking and runtime Pydantic validation. The changes are well-structured, adding overloads for ai.prompt and propagating schema types through the generation pipeline. I've identified a potential logic issue in to_generate_action_options related to output_constrained handling and also noted some code duplication that could be refactored for better maintainability. Overall, this is a great enhancement to the SDK's type safety.