-
Notifications
You must be signed in to change notification settings - Fork 685
chore(py): re-organize file structure, api signatures, types #4879
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Draft
huangjeff5
wants to merge
29
commits into
main
Choose a base branch
from
jh-redesign
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+10,089
−32,100
Draft
Changes from all commits
Commits
Show all changes
29 commits
Select commit
Hold shift + click to select a range
203a940
chore(py): re-organize file structure, api signatures, types
huangjeff5 e488bc4
Clean up init
huangjeff5 638382f
APi cleanup - veneer types, kwargs, remove dead code
huangjeff5 f2a659e
fix build issues
huangjeff5 b99e551
fix failures
huangjeff5 b0afa50
Fix imports
huangjeff5 6c0c6b2
remove reranker / indexer
huangjeff5 915846a
major refactor
huangjeff5 0b34e55
fix preflight
huangjeff5 0e62ee0
updates
huangjeff5 2ffc0ec
more cleanup
huangjeff5 f6ce902
fix lint
huangjeff5 129a7a9
fixes
huangjeff5 c8f5793
fix
huangjeff5 586c646
fix failures
huangjeff5 04134c7
fix type checker issues
huangjeff5 1664f79
more typing fixes
huangjeff5 740f7f0
fix type checking for tools and samples
huangjeff5 c0c8d26
fix
huangjeff5 77b6516
more fixes
huangjeff5 3b42612
Clean up exports
huangjeff5 0e2b51a
fixes
huangjeff5 b1067e7
Delete py/docs/dev_ui_eventloop_model.md
huangjeff5 7f05120
Delete py/docs/python_docs_roadmap.md
huangjeff5 7e61230
address comments
huangjeff5 3146f93
ModelUsage
huangjeff5 1d5ba9e
fix preflight
huangjeff5 7465128
clean up realtime tracing
huangjeff5 00ca0d2
more cleanup
huangjeff5 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
The table of contents is too big for display.
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -54,6 +54,27 @@ | |
| class ClassTransformer(ast.NodeTransformer): | ||
| """AST transformer that modifies class definitions.""" | ||
|
|
||
| # Classes to exclude from the generated output because they have | ||
| # hand-written veneer types in the SDK. These wire types should not be | ||
| # exposed — the veneer types are the public API. | ||
| EXCLUDED_CLASSES: frozenset[str] = frozenset({ | ||
| # These classes have hand-written veneer types in the SDK. | ||
| # The veneer is the ONLY type — used by plugins and end users alike. | ||
| # Wire types are NOT exposed. | ||
| # DocumentData stays in _typing.py — it's the wire base used internally. | ||
| # Reranker/retriever/indexer types removed from SDK entirely. | ||
| 'RankedDocumentData', | ||
| 'RankedDocumentMetadata', | ||
| 'CommonRerankerOptions', | ||
| 'RerankerRequest', | ||
| 'RerankerResponse', | ||
| 'CommonRetrieverOptions', | ||
| 'RetrieverRequest', | ||
| 'RetrieverResponse', | ||
| # Note: ModelRequest, ModelResponse, ModelResponseChunk are NOT excluded | ||
| # because _model.py imports them as base classes for veneer types. | ||
| }) | ||
|
|
||
| def __init__(self, models_allowing_extra: set[str] | None = None) -> None: | ||
| """Initialize the ClassTransformer. | ||
|
|
||
|
|
@@ -261,8 +282,19 @@ def visit_ClassDef(self, node: ast.ClassDef) -> object: | |
| node: The ClassDef AST node to transform. | ||
|
|
||
| Returns: | ||
| The transformed ClassDef node. | ||
| The transformed ClassDef node, or None to remove it. | ||
| """ | ||
| # Rename classes to their Python-convention wire type names. | ||
| renamed_classes: dict[str, str] = { | ||
| 'Message': 'MessageData', # schema "Message" becomes Python "MessageData" wire type | ||
| } | ||
| if node.name in renamed_classes: | ||
| node.name = renamed_classes[node.name] | ||
|
|
||
| # Exclude classes that have hand-written veneer types. | ||
| if node.name in self.EXCLUDED_CLASSES: | ||
| return None | ||
|
|
||
| # First apply base class transformations recursively | ||
| node = cast(ast.ClassDef, super().generic_visit(node)) | ||
| new_body: list[ast.stmt | ast.Constant | ast.Assign] = [] | ||
|
|
@@ -277,7 +309,7 @@ def visit_ClassDef(self, node: ast.ClassDef) -> object: | |
| # Generate a more descriptive docstring based on class type | ||
| if self.is_rootmodel_class(node): | ||
| docstring = f'Root model for {node.name.lower().replace("_", " ")}.' | ||
| elif any(isinstance(base, ast.Name) and base.id == 'BaseModel' for base in node.bases): | ||
| elif any(isinstance(base, ast.Name) and base.id == 'GenkitModel' for base in node.bases): | ||
| docstring = f'Model for {node.name.lower().replace("_", " ")} data.' | ||
| elif any(isinstance(base, ast.Name) and base.id == 'Enum' for base in node.bases): | ||
| n = node.name.lower().replace('_', ' ') | ||
|
|
@@ -326,8 +358,8 @@ def visit_ClassDef(self, node: ast.ClassDef) -> object: | |
| self.modified = True | ||
| continue | ||
| new_body.append(stmt) | ||
| elif any(isinstance(base, ast.Name) and base.id == 'BaseModel' for base in node.bases): | ||
| # Add or update model_config for BaseModel classes | ||
| elif any(isinstance(base, ast.Name) and base.id == 'GenkitModel' for base in node.bases): | ||
| # Add or update model_config for GenkitModel classes | ||
| added_config = False | ||
| frozen = node.name == 'PathMetadata' | ||
| has_schema = self.has_schema_field(node) | ||
|
|
@@ -405,6 +437,11 @@ def visit_ClassDef(self, node: ast.ClassDef) -> object: | |
| if node.name == 'GenerateActionOutputConfig': | ||
| self._inject_schema_type_field(new_body) | ||
|
|
||
| # PYTHON EXTENSION: Inline wrapper types in ModelRequest for better DX. | ||
| # Plugin authors see `messages: list[MessageData]` instead of `messages: Messages`. | ||
| if node.name == 'ModelRequest': | ||
| self._inline_model_request_types(new_body) | ||
|
|
||
| node.body = cast(list[ast.stmt], new_body) | ||
| return node | ||
|
|
||
|
|
@@ -493,6 +530,66 @@ def _inject_schema_type_field(self, body: list[ast.stmt | ast.Constant | ast.Ass | |
| body.append(schema_type_field) | ||
| self.modified = True | ||
|
|
||
| def _inline_model_request_types(self, body: list[ast.stmt | ast.Constant | ast.Assign]) -> None: | ||
| """Inline wrapper types in ModelRequest for better developer experience. | ||
|
|
||
| Changes: | ||
| - messages: Messages -> messages: list[MessageData] | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. list[Message] and list[Document] if psosible no output I think |
||
| - tools: Tools | None -> tools: list[ToolDefinition] | None | ||
| - docs: Docs | None -> docs: list[DocumentData] | None | ||
| - output: OutputModel | None -> output: OutputConfig | None | ||
|
|
||
| This gives plugin authors a cleaner type signature in their IDE. | ||
| RootModel wrappers still exist for backward compatibility but ModelRequest | ||
| uses plain types directly. | ||
| """ | ||
| # Mapping from wrapper type name to inlined type | ||
| type_mappings: dict[str, tuple[str, str | None]] = { | ||
| # field_name: (inner_type, list_item_type or None if not a list) | ||
| 'messages': ('MessageData', 'list'), | ||
| 'tools': ('ToolDefinition', 'list'), | ||
| 'docs': ('DocumentData', 'list'), | ||
| 'output': ('OutputConfig', None), | ||
| } | ||
|
|
||
| for stmt in body: | ||
| if not isinstance(stmt, ast.AnnAssign): | ||
| continue | ||
| if not isinstance(stmt.target, ast.Name): | ||
| continue | ||
|
|
||
| field_name = stmt.target.id | ||
| if field_name not in type_mappings: | ||
| continue | ||
|
|
||
| inner_type, container = type_mappings[field_name] | ||
|
|
||
| # Build the new type annotation | ||
| if container == 'list': | ||
| # list[InnerType] | ||
| new_type = ast.Subscript( | ||
| value=ast.Name(id='list', ctx=ast.Load()), | ||
| slice=ast.Name(id=inner_type, ctx=ast.Load()), | ||
| ctx=ast.Load(), | ||
| ) | ||
| else: | ||
| # Just the inner type | ||
| new_type = ast.Name(id=inner_type, ctx=ast.Load()) | ||
|
|
||
| # Check if current annotation is Optional (X | None) | ||
| if isinstance(stmt.annotation, ast.BinOp) and isinstance(stmt.annotation.op, ast.BitOr): | ||
| # It's X | None, replace X with new_type | ||
| stmt.annotation = ast.BinOp( | ||
| left=new_type, | ||
| op=ast.BitOr(), | ||
| right=ast.Constant(value=None), | ||
| ) | ||
| else: | ||
| # Not optional, just replace | ||
| stmt.annotation = new_type | ||
|
|
||
| self.modified = True | ||
|
|
||
|
|
||
| def fix_field_defaults(content: str) -> str: | ||
| """Fix Field(None) and Field(None, ...) to use default=None for pyright compatibility. | ||
|
|
@@ -571,24 +668,111 @@ def add_header(content: str) -> str: | |
| # Ensure there's exactly one newline between header and content | ||
| # and future import is right after the header block's closing quotes. | ||
| future_import = 'from __future__ import annotations' | ||
| compat_import_block = """ | ||
|
|
||
| # Header imports - these go first after the future import | ||
| header_imports = """ | ||
| import sys | ||
| import warnings | ||
| from strenum import StrEnum | ||
| from typing import ClassVar | ||
|
|
||
| from genkit.core._compat import StrEnum | ||
| from pydantic.alias_generators import to_camel | ||
| """ | ||
|
|
||
| # Warnings filter - this goes AFTER all imports to avoid E402 | ||
| warnings_filter = """ | ||
| # Filter Pydantic warning about 'schema' field in OutputConfig shadowing BaseModel.schema(). | ||
| # This is intentional - the field name is required for wire protocol compatibility. | ||
| warnings.filterwarnings( | ||
| 'ignore', | ||
| message='Field name "schema" in "OutputConfig" shadows an attribute in parent', | ||
| category=UserWarning, | ||
| ) | ||
| """ | ||
|
|
||
| header_text = header.format(year=datetime.now().year) | ||
|
|
||
| # Remove existing future import and StrEnum import from content. | ||
| # Lines that are already in the header template and should not be duplicated. | ||
| lines_in_header = { | ||
| future_import, | ||
| 'from enum import StrEnum', | ||
| 'from strenum import StrEnum', | ||
| 'import sys', | ||
| 'import warnings', | ||
| 'from typing import ClassVar', | ||
| 'from pydantic.alias_generators import to_camel', | ||
| } | ||
|
|
||
| lines = content.splitlines() | ||
| filtered_lines = [ | ||
| line for line in lines if line.strip() != future_import and line.strip() != 'from enum import StrEnum' | ||
| ] | ||
| content_imports: list[str] = [] # imports from content that need to go before warnings.filterwarnings() | ||
| filtered_lines: list[str] = [] | ||
| in_docstring = False | ||
| skip_warnings_filterwarnings = False | ||
|
|
||
| for line in lines: | ||
| stripped = line.strip() | ||
|
|
||
| # Skip lines that are already in the header template | ||
| if stripped in lines_in_header: | ||
| continue | ||
|
|
||
| # Skip the module docstring (will be re-added by header) | ||
| if stripped.startswith('"""Schema types module') or stripped.startswith("'''Schema types module"): | ||
| in_docstring = True | ||
| continue | ||
| if in_docstring: | ||
| if stripped.endswith('"""') or stripped.endswith("'''"): | ||
| in_docstring = False | ||
| continue | ||
|
|
||
| # Skip standalone docstring lines that are just the closing quotes | ||
| if stripped in ('"""', "'''"): | ||
| continue | ||
|
|
||
| # Skip the string literal form of the docstring (from ast.unparse) | ||
| # This happens when ast.unparse converts a module docstring to a string expression | ||
| if stripped.startswith("'Schema types module") or stripped.startswith('"Schema types module'): | ||
| continue | ||
|
|
||
| # Skip warnings.filterwarnings call (may span multiple lines) | ||
| if stripped.startswith('warnings.filterwarnings('): | ||
| if not stripped.endswith(')'): | ||
| skip_warnings_filterwarnings = True | ||
| continue | ||
| if skip_warnings_filterwarnings: | ||
| if stripped.endswith(')'): | ||
| skip_warnings_filterwarnings = False | ||
| continue | ||
|
|
||
| # Collect imports separately - they need to go before warnings.filterwarnings() | ||
| # to avoid E402 "module level import not at top of file" | ||
| if stripped.startswith('from ') or stripped.startswith('import '): | ||
| content_imports.append(line) | ||
| continue | ||
|
|
||
| filtered_lines.append(line) | ||
|
|
||
| cleaned_content = '\n'.join(filtered_lines) | ||
|
|
||
| final_output = header_text + future_import + '\n' + compat_import_block + '\n\n' + cleaned_content | ||
| # Fix field type annotations: schema 'Message' was renamed to 'MessageData' | ||
| # but field references in other classes still say 'Message'. | ||
| import re | ||
|
|
||
| cleaned_content = re.sub(r'\bMessage\b(?!Data)', 'MessageData', cleaned_content) | ||
|
|
||
| # Assemble final output with imports BEFORE warnings.filterwarnings() to avoid E402. | ||
| # Order: header -> future import -> header imports -> content imports -> warnings filter -> content | ||
| content_imports_block = '\n'.join(content_imports) + '\n' if content_imports else '' | ||
| final_output = ( | ||
| header_text | ||
| + future_import | ||
| + '\n' | ||
| + header_imports | ||
| + content_imports_block | ||
| + warnings_filter | ||
| + '\n' | ||
| + cleaned_content | ||
| ) | ||
| if not final_output.endswith('\n'): | ||
| final_output += '\n' | ||
| return final_output | ||
|
|
@@ -721,14 +905,15 @@ def main() -> None: | |
| if len(sys.argv) != 2: | ||
| sys.exit(1) | ||
|
|
||
| typing_file = Path(sys.argv[1]) | ||
| typing_file = Path(sys.argv[1]).resolve() | ||
|
|
||
| # Derive genkit-schema.json path relative to the typing.py file | ||
| # typing.py is at: py/packages/genkit/src/genkit/core/typing.py | ||
| # Derive genkit-schema.json path relative to the _typing.py file | ||
| # _typing.py is at: py/packages/genkit/src/genkit/_core/_typing.py | ||
| # schema is at: genkit-tools/genkit-schema.json | ||
| # So we go up 6 directories from typing.py to reach repo root, then into genkit-tools/ | ||
| # Go up 6 directories from _typing.py to reach repo root (genkit-cleanup/), then into genkit-tools/ | ||
| # _core(1) -> genkit(2) -> src(3) -> genkit(4) -> packages(5) -> py(6) -> genkit-cleanup | ||
| schema_path = typing_file.parent | ||
| for _ in range(6): # Go up: core -> genkit -> src -> genkit -> packages -> py -> (repo root) | ||
| for _ in range(6): | ||
| schema_path = schema_path.parent | ||
| schema_path = schema_path / 'genkit-tools' / 'genkit-schema.json' | ||
|
|
||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.