diff --git a/.github/workflows/python-samples.yml b/.github/workflows/python-samples.yml index ee279d281d..c3fdb97f64 100644 --- a/.github/workflows/python-samples.yml +++ b/.github/workflows/python-samples.yml @@ -59,17 +59,8 @@ jobs: - provider-google-genai-vertexai-hello - provider-google-genai-vertexai-image - provider-anthropic-hello - - provider-mistral-hello - - provider-xai-hello - - provider-deepseek-hello - - provider-cohere-hello - - provider-huggingface-hello - - provider-amazon-bedrock-hello - - provider-cloudflare-workers-ai-hello - provider-compat-oai-hello - - provider-microsoft-foundry-hello - - provider-checks-hello - - provider-observability-hello + - provider-vertex-ai-model-garden - framework-prompt-demo - framework-format-demo - framework-context-demo @@ -79,8 +70,6 @@ jobs: - framework-restaurant-demo - web-fastapi-bugbot - web-flask-hello - - web-multi-server - - web-short-n-long steps: - uses: actions/checkout@v5 @@ -124,9 +113,6 @@ jobs: matrix: include: - sample: provider-ollama-hello - # Use the smallest model for CI smoke test β€” full model - # list (gemma3:latest, mistral-nemo, llava, nomic-embed-text) - # would be ~10 GB and too slow for PR checks. models: "gemma3:1b" steps: - uses: actions/checkout@v5 diff --git a/py/CHANGELOG.md b/py/CHANGELOG.md deleted file mode 100644 index d5aac56258..0000000000 --- a/py/CHANGELOG.md +++ /dev/null @@ -1,227 +0,0 @@ -# Changelog - -All notable changes to the Genkit Python SDK will be documented in this file. - -The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), -and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - -## [Unreleased] - -## [0.5.0] - 2026-02-04 - -This is a major release with **178 commits** and **680 files changed** over **8 months** -since 0.4.0 (May 2025), representing the most significant update to the Genkit Python SDK to date. - -### Impact Summary - -| Category | Impact Level | Description | -|----------|-------------|-------------| -| **New Plugins** | 🟒 High | 7 new model providers and 3 new telemetry plugins | -| **Core Features** | 🟒 High | Rerankers, tool calling, background models, Dotprompt integration | -| **Type Safety** | 🟑 Medium | Comprehensive type checking with ty/pyrefly/pyright | -| **Breaking Changes** | 🟑 Medium | PluginV2 refactor requires migration | -| **Developer Experience** | 🟒 High | Hot reloading, improved samples, better docs | -| **Security** | 🟒 High | Ruff audit, PySentry scanning, SigV4 signing | -| **Performance** | 🟒 High | Per-event-loop HTTP caching, release pipeline 15x faster | - -### Added - -#### New Model Provider Plugins -- **`genkit-plugin-anthropic`**: Full Anthropic Claude model support (#3919) -- **`genkit-plugin-amazon-bedrock`**: AWS Bedrock integration for Claude, Titan, Llama models (#4389) -- **`genkit-plugin-microsoft-foundry`**: Azure OpenAI (Microsoft Foundry) support (#4383) -- **`genkit-plugin-cloudflare-workers-ai`**: Cloudflare Workers AI models (#4405) -- **`genkit-plugin-deepseek`**: DeepSeek models with structured output (#4051) -- **`genkit-plugin-xai`**: xAI (Grok) models with plugin config (#4001, #4289) -- **`genkit-plugin-mistral`**: Mistral AI models (Large, Small, Codestral, Pixtral) (#4406) -- **`genkit-plugin-huggingface`**: Hugging Face Inference API with 17+ providers (#4406) - -#### New Telemetry Plugins -- **`genkit-plugin-aws`**: AWS X-Ray telemetry with SigV4 signing (#4390, #4402) -- **`genkit-plugin-observability`**: Third-party telemetry (Sentry, Honeycomb, Datadog) -- **`genkit-plugin-google-cloud`**: GCP telemetry parity with JS/Go implementations (#4281) - -#### Core Framework Features -- **Agentive Tool Calling**: Define tools with `@ai.tool()` decorator for AI agents -- **Rerankers**: Initial reranker implementation for RAG pipelines (#4065) -- **Background Models**: Dynamic model discovery and background action support (#4327) -- **Resource Support**: Full resource management with MCP integration (#4204, #4048) -- **Evaluator Metrics**: ANSWER_RELEVANCY, FAITHFULNESS, MALICIOUSNESS metrics (#3806) -- **MCP Plugin**: Model Context Protocol integration with tests (#4054) -- **Retriever/Embedder References**: Reference support matching JS SDK (#3922, #3936) -- **Output Formats**: Array, enum, and JSONL formats for JS parity (#4230) -- **Pydantic Output**: Return Pydantic instances when output schema passed (#4413) - -#### Dotprompt Integration (via [google/dotprompt](https://github.com/google/dotprompt)) -- **Dotpromptz 0.1.5**: Upgraded to latest version with type-safe schema fields -- **Python 3.14 Support**: PyO3/maturin ABI compatibility for Rust-based Handlebars engine -- **Directory/File Prompt Loading**: Automatic prompt discovery matching JS SDK (#3955, #3971) -- **Handlebars Partials**: `define_partial` for template reuse (#4088) -- **Render System Prompts**: `render_system_prompt` and `render_user_prompt` methods (#3503, #3705) -- **Callable Support**: Prompts can now be used as callables (#4053) -- **Cycle Detection**: Partial resolution with cycle detection prevents infinite recursion -- **Path Traversal Hardening**: Security fix for CWE-22 vulnerability -- **Helper Parity**: Consistent Handlebars helper behavior across all runtimes -- **Release Pipeline**: Automated PyPI publishing, release time reduced from 30 min to 2 min - -#### Developer Experience -- **Hot Reloading**: [Watchdog](https://github.com/gorakhargosh/watchdog)-based autoreloading for all samples (#4268) -- **Security Scanning**: PySentry-rs integration in hooks and CI (#4273) -- **TODO Linting**: Automated issue creation for TODO comments (#4376) -- **Centralized Action Latency**: Built-in performance tracking (#4267) -- **Sample Improvements**: Preamble scripts, browser auto-open, rich tracebacks (#4375, #4373) -- **Release Automation**: `bin/release_check`, `bin/bump_version` scripts -- **Consistency Checks**: `bin/check_consistency` for package validation - -#### Type Safety Improvements -- **ty Integration**: Stricter, faster type checking from Astral (#4094) -- **pyrefly Integration**: Meta's type checker for additional coverage (#4316) -- **pyright Enforcement**: Full Microsoft type checking (#4310) -- **Comprehensive Fixes**: Zero type errors across all packages (#4249-4260, #4270) - -#### Documentation -- **Module Docstrings**: Terminology tables and ASCII data flow diagrams (#4322) -- **GEMINI.md Updates**: Test file naming, import guidelines, TODO format (#4381, #4393, #4397) -- **Sample Documentation**: Testing notes for all samples (#4294) -- **HTTP Client Guidelines**: Event loop binding best practices (#4430) -- **Roadmap**: Plugin API conformance analysis (#4431) - -#### Samples & Demos -- **New Samples**: tool-interrupt, short-n-long, media-models-demo, prompt samples -- **Run Script Standardization**: Central script for running samples with `genkit start` -- **Rich Tracebacks**: Improved error output in samples - -### Changed - -#### Breaking Changes -- **PluginV2 Refactor**: Major plugin architecture update - existing plugins may need migration (#4132) - - Plugins now use a standardized registration pattern - - Configuration options are more consistent across plugins -- **Async-First Architecture**: Removed sync base, fully async by default (#4244) -- **Embed API**: Refactored `embed/embed_many` for JS parity (#4269) - -#### Improvements -- **Python 3.14 Support**: Full compatibility with Python 3.14 (#3947) -- **Gemini 2.5/3.0 Upgrade**: Default models updated to Gemini 2.5/3.0 (#3771, #4277) -- **Dotpromptz 0.1.5**: Latest template engine with improved features (#4324) -- **PEP 8 Compliance**: All in-function imports moved to top-level (#4396-4400) -- **CI Consolidation**: Single workflow, every commit is release-worthy (#4410) -- **Reflection API**: Improved multi-runtime handling and health checks -- **Dev UI Defaults**: Better default configurations - -### Fixed - -#### Critical Fixes -- **Race Condition**: Dev server startup race condition resolved (#4225) -- **Thread Safety**: Per-event-loop HTTP client caching prevents event loop binding errors (#4419, #4429) -- **Infinite Recursion**: Cycle detection in Handlebars partial resolution (via Dotprompt) -- **Real-Time Telemetry**: Trace ID formatting and streaming fixes (#4285) -- **Structured Output**: DeepSeek model structured output generation (#4374) -- **JSON Schema**: None type handling per JSON Schema spec (#4247) -- **Windows Support**: File-safe timestamp format for runtime files (#3727) - -#### Model/Plugin Fixes -- **Gemini Models**: Various bug fixes (#4432) -- **TTS/Veo Models**: System prompt support in model config (#4411) -- **Google GenAI**: Model config and README updates (#4306, #4323) -- **Ollama**: Sample fixes and model server management (#4133, #4227) -- **Embedders**: Reflection health check fixes (#3969) -- **Complex Schemas**: Support for complex schemas in Gemini (#3049) - -#### Sample Fixes -- Extensive sample fixes across all demos (#4122-4418) -- System prompt fields added to all Gemini samples (#4391) -- Missing dependencies resolved (#4282) -- Consistent `genkit start` usage (#4226) -- GCloud auto-setup for Vertex AI samples (#4427) - -#### Type Errors -- Resolved all `ty`, `pyrefly`, and `pyright` type errors -- Re-enabled disabled tests and improved coverage -- Fixed evaluator plugin imports and StrEnum compatibility - -### Security - -- **Ruff Security Audit**: Addressed all security and code quality warnings (#4409) -- **SigV4 Signing**: AWS X-Ray OTLP exporter now uses proper AWS signatures (#4402) -- **Path Traversal Hardening**: CWE-22 vulnerability fix in Dotprompt (via google/dotprompt) -- **License Compliance**: Fixed license headers in all configuration files (#3930) -- **PySentry Integration**: Continuous security vulnerability scanning (#4273) - -### Performance - -- **Per-Event-Loop HTTP Client Caching**: Reuses HTTP connections within event loops, prevents connection overhead -- **Dotprompt Release Pipeline**: Reduced from 30 minutes to 2 minutes (15x faster) -- **CI Consolidation**: Single workflow, every commit is release-worthy (#4410) -- **ty Type Checker**: Faster type checking than pyright alone (#4094) - -### Deprecated - -- Sync API base classes are removed in favor of async-first architecture - -### Contributors - -This release includes contributions from **13 developers** across **188 PRs**. Thank you to everyone who contributed! - -| Contributor | PRs | Commits | Key Contributions | -|-------------|-----|---------|-------------------| -| [**@yesudeep**](https://github.com/yesudeep) | 91 | 93 | **Core**: async-first architecture (#4244), Genkit class methods (#4274), embed/embed_many API refactor (#4269), centralized action latency (#4267), array/enum/jsonl output formats (#4230). **Plugins**: AWS Bedrock (#4389), AWS X-Ray with SigV4 (#4390, #4402), Azure OpenAI (#4383), Cloudflare Workers AI (#4405), Mistral AI (#4406), Hugging Face (#4406), GCP telemetry (#4281). **Type Safety**: ty integration (#4094), pyrefly (#4316), pyright (#4310), comprehensive fixes (#4249-4270). **DevEx**: hot reloading (#4268), per-event-loop HTTP caching (#4419, #4429), PySentry security (#4273), TODO linting (#4376), CI consolidation (#4410), session/chat API (#4278, #4275), background models (#4327), docs (#4322, #4393, #4430). **Samples**: 20+ sample fixes and improvements (#4283, #4373, #4375, #4427). | -| [**@MengqinShen**](https://github.com/MengqinShen) (Elisa Shen) | 42 | 42 | **Core**: Resource support implementation (#4204). **Samples**: menu sample fixes (#4239, #4403), short-n-long (#4404), tool-interrupt (#4408), prompt sample (#4223, #4183), ollama-hello (#4133), genai-image (#4122, #4234), code-execution (#4134), anthropic sample (#4131). **Models**: Google GenAI model config (#4306), TTS/Veo model config (#4411), Gemini bug fixes (#4432), system prompt fields (#4391, #4418). **Docs**: README updates (#4323), multi-round flow logic (#4137). | -| [**@AbeJLazaro**](https://github.com/AbeJLazaro) | 11 | 8 | **Plugins**: Model Garden resolve/list actions (#3040), Ollama resolve action (#2972), type coverage and tests (#3011). **Fixes**: Gemini complex schema support (#3049), Firestore plugin naming (#3085), evaluator plugin requirements (#3166), optional dependencies setup (#3012). **Tests**: Model Garden tests (#3083). | -| [**@pavelgj**](https://github.com/pavelgj) | 10 | 7 | **Core**: Reflection API multi-runtime support (#3970), health check fixes (#3969). **Fixes**: Embedders reflection (#3969), Gemini version upgrades to 2.5 (#3909). | -| [**@zarinn3pal**](https://github.com/zarinn3pal) | 9 | 9 | **Plugins**: Anthropic (#3919), DeepSeek (#4051, structured output fix #4374), xAI/Grok (#4001, config #4289), ModelGarden (#2568). **Telemetry**: GCP telemetry for Firebase observability (#3826, #4386). **Samples**: OpenAI Compat tools (#3684). | -| [**@huangjeff5**](https://github.com/huangjeff5) | 7 | 7 | **Core**: PluginV2 refactor with new registration pattern (#4132), type safety improvements (#4310), Pydantic output instances (#4413), session/chat refactor (#4321). **Telemetry**: Real-time telemetry and trace ID formatting (#4285). | -| [**@hendrixmar**](https://github.com/hendrixmar) | 7 | 7 | **Evaluators**: ANSWER_RELEVANCY, FAITHFULNESS, MALICIOUSNESS metrics (#3806), ModelReference support (#3949, #3951). **Plugins**: OpenAI compat list_actions (#3240), resolve_method (#3055). **Dotprompt**: render_system_prompt (#3503), render_user_prompt (#3705). | -| [**@ssbushi**](https://github.com/ssbushi) | 6 | 2 | **Evaluators**: Simple evaluators plugin (#2835). **Docs**: MkDocs API reference updates (#2852), genkit-tools model optional (#3918). | -| [**@shrutip90**](https://github.com/shrutip90) | 1 | 1 | **Types**: ResourcePartSchema exports via genkit-tools (#3239). | -| [**@schlich**](https://github.com/schlich) | 1 | 1 | **Types**: Type annotations for ai module. | -| [**@ktsmadhav**](https://github.com/ktsmadhav) | 1 | 1 | **Fixes**: Windows support with file-safe timestamp format (#3727). | -| [**@junhyukhan**](https://github.com/junhyukhan) | 1 | 1 | **Docs**: Typo fixes. | -| [**@CorieW**](https://github.com/CorieW) | 1 | 1 | Community contribution. | - -**[google/dotprompt](https://github.com/google/dotprompt) Contributors** (Dotprompt Python integration): - -| Contributor | PRs | Commits | Key Contributions | -|-------------|-----|---------|-------------------| -| [**@yesudeep**](https://github.com/yesudeep) | 50+ | 100+ | **Rust Engine**: dotpromptz-handlebars with PyO3/maturin (#365), Python 3.14 ABI support, Rust Handlebars runtime. **Features**: Cycle detection in partial resolution, path traversal hardening (CWE-22), directory/file prompt loading (#3955, #3971), Handlebars partials (#4088), callable prompts (#4053). **Build**: Bazel rules_dart/rules_flutter, release pipeline 15x faster (30minβ†’2min), maturin wheel builds. **IDE**: Monaco syntax highlighting, CodeMirror 6 integration, Storybook demos. **Polyglot**: Python, Go, Dart, Rust, TypeScript implementations. | -| [**@MengqinShen**](https://github.com/MengqinShen) | 42 | 45 | **CI/CD**: GitHub Actions workflows for Python package publishing, automated release-please, dotpromptz PyPI releases (0.1.2-0.1.5), handlebarrz releases, wheel artifact management. | -| [**@Zereker**](https://github.com/Zereker) | 1 | 1 | **Go**: Closure fix preventing template sharing between instances. | - ---- - -## [0.4.0] - 2025-05-26 - -### Added - -- **Telemetry Plugins** - - `genkit-plugin-microsoft-foundry`: Azure Application Insights integration (consolidated with model access) - - `genkit-plugin-cf`: Generic OTLP export for Cloudflare and other backends - - `genkit-plugin-observability`: Unified presets for Sentry, Honeycomb, Datadog, Grafana Cloud, Axiom - -- **Model Provider Plugins** - - `genkit-plugin-mistral`: Mistral AI models (Large, Small, Codestral, Pixtral) - - `genkit-plugin-huggingface`: Hugging Face Inference API with 17+ inference providers - -- **Core Framework** - - Improved tracing and observability support - - Enhanced type safety across all modules - -### Changed - -- All plugins now share the same version number as the core framework -- Improved documentation and README files for all packages - -## [0.3.0] - 2025-04-08 - -### Added - -- Initial public release of Genkit Python SDK -- Core framework (`genkit`) -- Model plugins: Anthropic, Google GenAI, Ollama, Vertex AI, xAI, DeepSeek -- Telemetry plugins: AWS, Google Cloud, Firebase -- Utility plugins: Flask, MCP, Evaluators, Dev Local Vectorstore - -[Unreleased]: https://github.com/firebase/genkit/compare/genkit-python@0.5.0...HEAD -[0.5.0]: https://github.com/firebase/genkit/compare/genkit-python@0.4.0...genkit-python@0.5.0 -[0.4.0]: https://github.com/firebase/genkit/compare/genkit-python@0.3.0...genkit-python@0.4.0 -[0.3.0]: https://github.com/firebase/genkit/releases/tag/genkit-python@0.3.0 diff --git a/py/README.md b/py/README.md index 848e8c7b49..be7c5df2c9 100644 --- a/py/README.md +++ b/py/README.md @@ -41,13 +41,13 @@ and is loaded lazily by the `Registry`. ### Plugin Class Hierarchy -All plugins inherit from `genkit.core.plugin.Plugin` and implement three +All plugins inherit from `genkit._core.plugin.Plugin` and implement three abstract methods: ``` β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ Plugin (Abstract Base Class) β”‚ - β”‚ genkit.core.plugin.Plugin β”‚ + β”‚ genkit._core.plugin.Plugin β”‚ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”‚ β”‚ β”‚ name: str β”‚ @@ -182,9 +182,7 @@ registry uses a multi-step resolution algorithm: ### Writing a Custom Plugin ```python -from genkit.core.plugin import Plugin -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionKind +from genkit.plugin_api import Plugin, Action, ActionMetadata, ActionKind class MyPlugin(Plugin): @@ -305,7 +303,6 @@ and deciding whether to allow or reject it: py/ β”œβ”€β”€ packages/genkit/ # Core Genkit framework package β”œβ”€β”€ plugins/ # Official plugins -β”‚ β”œβ”€β”€ amazon-bedrock/ # Amazon Bedrock models + X-Ray telemetry (community) β”‚ β”œβ”€β”€ anthropic/ # Claude models β”‚ β”œβ”€β”€ azure/ # Azure AI telemetry (community) β”‚ β”œβ”€β”€ cloudflare-workers-ai/# Cloudflare Workers AI + OTLP telemetry (community) @@ -377,7 +374,7 @@ print(response.text) | **Context Provider** | Middleware that runs *before* a flow is called via HTTP. It reads the request (headers, body) and either provides auth info to the flow or rejects the request. | `api_key()`, `create_flows_asgi_app()` | | **Flow Server** | A built-in HTTP server that wraps your flows as API endpoints so `curl` (or any client) can call them. It's Genkit's simple way to deploy flows without a web framework. | `create_flows_asgi_app()` | | **Registry** | The internal directory of all defined flows, tools, models, and prompts. The Dev UI and CLI read it to discover what's available. | `ai.registry` | -| **Action** | The low-level building block behind flows, tools, models, and prompts. Everything you define becomes an "action" in the registry with input/output schemas and tracing. | `genkit.core.action` | +| **Action** | The low-level building block behind flows, tools, models, and prompts. Everything you define becomes an "action" in the registry with input/output schemas and tracing. | `genkit.plugin_api` | | **Middleware** | Functions that wrap around model calls to add behavior β€” logging, caching, safety checks, or modifying requests/responses. Runs at the model level, not HTTP level. | `ai.define_model(use=[...])` | | **Embedder** | A model that turns text into numbers (vectors) for similarity search. Used with vector stores for RAG (Retrieval-Augmented Generation). | `ai.embed()` | | **Retriever** | A component that searches a vector store and returns relevant documents for a query. Used in RAG pipelines. | `ai.retrieve()` | @@ -403,7 +400,6 @@ print(response.text) Some plugins are community-maintained and supported on a best-effort basis: -- **amazon-bedrock** - Amazon Bedrock models + AWS X-Ray telemetry - **azure** - Azure Monitor / Application Insights telemetry - **cloudflare-workers-ai** - Cloudflare Workers AI models + OTLP telemetry - **cohere** - Cohere command models + reranking @@ -427,7 +423,6 @@ examples of any feature: | `provider-xai-hello` | Model, Flow | Grok models | | `provider-mistral-hello` | Model, Flow | Mistral models | | `provider-huggingface-hello` | Model, Flow | HuggingFace Inference API | -| `provider-amazon-bedrock-hello` | Model, Flow, Telemetry | AWS Bedrock + X-Ray tracing | | `provider-cloudflare-workers-ai-hello` | Model, Flow, Telemetry | Cloudflare Workers AI | | `provider-microsoft-foundry-hello` | Model, Flow | Azure AI Foundry | | **Google Cloud** | | | @@ -470,7 +465,7 @@ uv run pytest . Run tests for a specific plugin: ```bash -uv run pytest plugins/amazon-bedrock/ +uv run pytest plugins/anthropic/ ``` ## Development diff --git a/py/bin/generate_schema_typing b/py/bin/generate_schema_typing index 49d0c60d82..ee87714351 100755 --- a/py/bin/generate_schema_typing +++ b/py/bin/generate_schema_typing @@ -38,7 +38,7 @@ while [[ $# -gt 0 ]]; do done TOP_DIR=$(git rev-parse --show-toplevel) -TYPING_FILE="${TOP_DIR}/py/packages/genkit/src/genkit/core/typing.py" +TYPING_FILE="${TOP_DIR}/py/packages/genkit/src/genkit/_core/_typing.py" # If in CI mode and the file exists, make a backup copy to compare later. BACKUP_FILE="" diff --git a/py/bin/sanitize_schema_typing.py b/py/bin/sanitize_schema_typing.py index 0d35b253df..6e710445b8 100644 --- a/py/bin/sanitize_schema_typing.py +++ b/py/bin/sanitize_schema_typing.py @@ -54,6 +54,37 @@ 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', + # ModelRequest is excluded β€” fully hand-written in _model.py with flat + # output fields (output_format, output_schema, etc.) instead of nested output. + 'ModelRequest', + # GenerateRequest excluded β€” wire type not needed; ModelRequest is the veneer. + 'GenerateRequest', + # OutputModel excluded β€” unused; OutputConfig is the type for model output. + 'OutputModel', + # GenerateResponse, Request, ModelResponse (wire) excluded β€” hand-written in _model.py. + # No references to rewrite; these were only referenced by each other. + 'GenerateResponse', + 'Request', + 'ModelResponse', + }) + def __init__(self, models_allowing_extra: set[str] | None = None) -> None: """Initialize the ClassTransformer. @@ -261,8 +292,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 or are unused. + 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 +319,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 +368,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) @@ -571,24 +613,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 +850,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' diff --git a/py/docs/assets/favicon.png b/py/docs/assets/favicon.png deleted file mode 100644 index c5a82e9bde..0000000000 Binary files a/py/docs/assets/favicon.png and /dev/null differ diff --git a/py/docs/assets/logo.png b/py/docs/assets/logo.png deleted file mode 100644 index cc41232300..0000000000 Binary files a/py/docs/assets/logo.png and /dev/null differ diff --git a/py/docs/dev_ui_eventloop_model.md b/py/docs/dev_ui_eventloop_model.md deleted file mode 100644 index a31890e16e..0000000000 --- a/py/docs/dev_ui_eventloop_model.md +++ /dev/null @@ -1,155 +0,0 @@ -# Dev UI + Event Loop Model - - -## Context - - -In Python async systems, an event loop is the runtime that drives `await` code. - - -For web apps, frameworks such as FastAPI typically own that loop for request handling. Genkit Dev UI reflection can also execute flows/actions, which means the same app can receive execution from: -- normal web requests (framework loop) -- Dev UI/reflection execution path (potentially another loop) - - -Some async clients (provider SDK clients, `httpx.AsyncClient`, etc.) are effectively tied to the loop where they were created. Reusing one instance from another loop can fail at runtime. - - -## Problem Statement - - -How do we let one Python app support both web-framework execution and Dev UI execution seamlessly, without: -- forcing framework-specific lifecycle wiring on app developers -- introducing hard-to-debug cross-loop runtime failures? - - -## Options Considered - - -### A) Single-event-loop architecture (current solution) - - -How it works: -- Reflection is forced onto the same event loop as app execution. -- Developers wire framework lifecycle so Genkit reflection starts/stops in-loop. - - -App code shape: - - -```python -@asynccontextmanager -async def lifespan(app: FastAPI): - await ai.start_reflection_same_loop() - try: - yield - finally: - await ai.stop_reflection() -``` - - -Pros: -- Eliminates cross-loop client reuse by construction. - - -Cons: -- Framework-specific lifecycle burden for app developers. -- More docs/support surface and framework adapter complexity. -- Harder "it just works" story across FastAPI/Flask/Quart/etc. - - -### B) Separate loops + loop-local client management - - -How it works: -- Reflection remains separate-loop in-process. -- Plugin/runtime clients are acquired per-event-loop through a loop-local getter. -- Action handlers use `get_client()` at call time. - - -Plugin code shape: - - -```python -from collections.abc import Callable -from genkit.core._loop_local import _loop_local_client - - -self._get_client: Callable[[], AsyncOpenAI] = _loop_local_client( - lambda: AsyncOpenAI(**self._params) -) - - -async def _run(req, ctx): - client = self._get_client() - return await call_model(client, req, ctx) -``` - - -Pros: -- No framework lifecycle wiring for most app developers. -- Fits current runtime topology with modest plugin changes. -- Incremental rollout; immediate correctness improvements for provider SDK use. - - -Cons: -- App-owned global async clients can still be a footgun across loops. -- Requires plugin author discipline and regression tests. - - -## A vs B (Why B is Better for Product DX) - - -If primary goal is seamless Dev UI + web framework integration, B is the better fit: -- Better default developer experience (less unrelated concepts for app developer). -- Lower integration friction across frameworks. -- Smaller incremental change than architectural rework. -- Correctness is handled where it matters most (plugin/runtime-managed clients). - - -A is stricter runtime-wise, but pushes integration burden onto users and framework-specific docs/support. - -That also means every framework requires its own lifecycle hook implementation and has to be bridged with a plugin or custom app developer code. - -## Remaining Footgun (Explicit) - - -Still risky app code: - - -```python -client = httpx.AsyncClient() # module-global, reused across loops -``` - - -Safer app code: - - -```python -async with httpx.AsyncClient() as client: - await client.post(...) -``` - - -Mitigation: -- Keep plugin internals loop-safe by default. -- Add concise docs for app-owned async clients. - - -## Helper Placement Decision - - -Question: where should the loop-local helper live? - - -Options: -- Plugin namespace (`genkit.plugins..utils`) -> duplicates logic, inconsistent usage. -- Public top-level API (`genkit.loop_local_client`) -> broad public contract, harder to evolve. -- Core internal utility (`genkit.core._loop_local`) -> shared implementation without expanding user API. - - -Recommendation: -- Keep helper in **core internal** (`genkit.core._loop_local`) for now. -- Use it across official plugins. -- Revisit public export only if app-level demand is clear and stable. - diff --git a/py/docs/index.md b/py/docs/index.md index a401ef5a37..3023418f48 100644 --- a/py/docs/index.md +++ b/py/docs/index.md @@ -3,60 +3,254 @@ !!! note Full Genkit documentation is available at [genkit.dev](https://genkit.dev/python/docs/get-started/) -## Core Classes +## genkit -::: genkit.ai.Genkit +::: genkit.Genkit -::: genkit.ai.GenkitRegistry +::: genkit.Plugin -::: genkit.ai.Plugin +::: genkit.Action -## Inputs and Outputs +::: genkit.Flow -::: genkit.ai.Input +::: genkit.ActionKind -::: genkit.ai.Output +::: genkit.ActionRunContext -::: genkit.ai.OutputOptions +::: genkit.ExecutablePrompt -## Actions +::: genkit.PromptGenerateOptions -::: genkit.ai.ActionKind +::: genkit.ResumeOptions -::: genkit.ai.ActionRunContext +::: genkit.ToolRunContext -::: genkit.ai.FlowWrapper +::: genkit.tool_response -## Prompts +::: genkit.StreamResponse -::: genkit.ai.ExecutablePrompt +::: genkit.ModelStreamResponse -::: genkit.ai.PromptGenerateOptions +::: genkit.GenkitError -::: genkit.ai.ResumeOptions +::: genkit.PublicError -## Response Types +::: genkit.ToolInterruptError -::: genkit.ai.GenerateResponseWrapper +::: genkit.Message -::: genkit.ai.GenerateStreamResponse +::: genkit.Part -## Tools +::: genkit.TextPart -::: genkit.ai.ToolRunContext +::: genkit.MediaPart -::: genkit.ai.tool_response +::: genkit.Media -## Documents +::: genkit.CustomPart -::: genkit.ai.Document +::: genkit.ReasoningPart -## Retriever +::: genkit.Role -::: genkit.ai.SimpleRetrieverOptions +::: genkit.Metadata -## Version Info +::: genkit.ToolRequest -::: genkit.ai.GENKIT_VERSION +::: genkit.ToolRequestPart -::: genkit.ai.GENKIT_CLIENT_HEADER +::: genkit.ToolResponse + +::: genkit.ToolResponsePart + +::: genkit.ToolDefinition + +::: genkit.ToolChoice + +::: genkit.Document + +::: genkit.DocumentPart + +::: genkit.EmbedderRef + +::: genkit.EmbedderOptions + +::: genkit.Embedding + +::: genkit.EmbedRequest + +::: genkit.EmbedResponse + +::: genkit.ModelRequest + +::: genkit.ModelResponse + +::: genkit.ModelResponseChunk + +::: genkit.ModelConfig + +::: genkit.ModelInfo + +::: genkit.ModelUsage + +::: genkit.Constrained + +::: genkit.Stage + +::: genkit.Supports + +::: genkit.FinishReason + +## genkit.model + +::: genkit.model.BackgroundAction + +::: genkit.model.ModelRequest + +::: genkit.model.ModelResponse + +::: genkit.model.ModelResponseChunk + +::: genkit.model.ModelUsage + +::: genkit.model.Candidate + +::: genkit.model.FinishReason + +::: genkit.model.GenerateActionOptions + +::: genkit.model.Error + +::: genkit.model.Operation + +::: genkit.model.ToolRequest + +::: genkit.model.ToolDefinition + +::: genkit.model.ToolResponse + +::: genkit.model.ModelInfo + +::: genkit.model.Supports + +::: genkit.model.Constrained + +::: genkit.model.Stage + +::: genkit.model.model_action_metadata + +::: genkit.model.model_ref + +::: genkit.model.ModelRef + +::: genkit.model.ModelConfig + +::: genkit.model.Message + +::: genkit.model.get_basic_usage_stats + +## genkit.embedder + +::: genkit.embedder.EmbedRequest + +::: genkit.embedder.EmbedResponse + +::: genkit.embedder.Embedding + +::: genkit.embedder.embedder_action_metadata + +::: genkit.embedder.embedder_ref + +::: genkit.embedder.EmbedderRef + +::: genkit.embedder.EmbedderSupports + +::: genkit.embedder.EmbedderOptions + +## genkit.plugin_api + +::: genkit.plugin_api.Plugin + +::: genkit.plugin_api.Action + +::: genkit.plugin_api.ActionMetadata + +::: genkit.plugin_api.ActionKind + +::: genkit.plugin_api.ActionRunContext + +::: genkit.plugin_api.StatusCodes + +::: genkit.plugin_api.StatusName + +::: genkit.plugin_api.GenkitError + +::: genkit.plugin_api.GENKIT_CLIENT_HEADER + +::: genkit.plugin_api.GENKIT_VERSION + +::: genkit.plugin_api.loop_local_client + +::: genkit.plugin_api.tracer + +::: genkit.plugin_api.add_custom_exporter + +::: genkit.plugin_api.AdjustingTraceExporter + +::: genkit.plugin_api.RedactedSpan + +::: genkit.plugin_api.to_display_path + +::: genkit.plugin_api.to_json_schema + +::: genkit.plugin_api.get_cached_client + +::: genkit.plugin_api.get_callable_json + +::: genkit.plugin_api.is_dev_environment + +::: genkit.plugin_api.model_action_metadata + +::: genkit.plugin_api.model_ref + +::: genkit.plugin_api.ModelRef + +::: genkit.plugin_api.embedder_action_metadata + +::: genkit.plugin_api.embedder_ref + +::: genkit.plugin_api.EmbedderRef + +::: genkit.plugin_api.evaluator_action_metadata + +::: genkit.plugin_api.evaluator_ref + +::: genkit.plugin_api.EvaluatorRef + +::: genkit.plugin_api.ContextProvider + +::: genkit.plugin_api.RequestData + +## genkit.evaluator + +::: genkit.evaluator.EvalRequest + +::: genkit.evaluator.EvalResponse + +::: genkit.evaluator.EvalFnResponse + +::: genkit.evaluator.Score + +::: genkit.evaluator.Details + +::: genkit.evaluator.BaseEvalDataPoint + +::: genkit.evaluator.BaseDataPoint + +::: genkit.evaluator.EvalStatusEnum + +::: genkit.evaluator.evaluator_action_metadata + +::: genkit.evaluator.evaluator_ref + +::: genkit.evaluator.EvaluatorRef diff --git a/py/docs/python_docs_roadmap.md b/py/docs/python_docs_roadmap.md deleted file mode 100644 index 5588f6bedd..0000000000 --- a/py/docs/python_docs_roadmap.md +++ /dev/null @@ -1,284 +0,0 @@ -# Python Documentation Roadmap for genkit.dev - -> **Source:** [genkit-ai/docsite](https://github.com/genkit-ai/docsite) -> **Docs path:** `src/content/docs/docs/` -> **Generated:** 2026-02-07 -> **Updated:** 2026-02-07 -> -> **Scope Exclusions:** -> - **Chat/Session API** β€” Deprecated, skip -> - **Agents / Multi-Agent** β€” Not yet in Python SDK, skip -> - **MCP** β€” Will come later, skip for now -> -> This roadmap tracks every genkit.dev documentation page and whether Python -> examples/tabs need to be added, updated, or are already complete. It also -> identifies features demonstrated in JS examples that should be covered by -> Python samples in the `firebase/genkit` repo. - ---- - -## Summary - -| Status | Count | Description | -|--------|-------|-------------| -| βœ… Complete | 5 | Python tab exists with full parity | -| πŸ”Ά Partial | 6 | Python tab exists but is incomplete or stale | -| ❌ Missing | 8 | No Python tab at all (JS/Go only or JS-only) | -| βž– N/A | 5 | Language-agnostic or meta pages | - ---- - -## 1. Core Documentation Pages - -### βœ… Python Tab Complete (verify accuracy) - -| Page | File | Languages | Notes | -|------|------|-----------|-------| -| **Models** | `models.mdx` | js, go, python | Python examples for: `generate()`, system prompts, model parameters, structured output, streaming, multimodal input. **Missing:** Generating Media, Middleware (retry/fallback). | -| **Tool Calling** | `tool-calling.mdx` | js, go, python | Python examples for: defining tools, using tools, interrupts (link), explicitly handling tool calls. **Missing:** `maxTurns`, dynamically defining tools at runtime. | -| **Flows** | `flows.mdx` | js, go, python | Python examples for: defining flows, input/output schemas, calling flows, streaming flows, deploying (Flask). **Missing:** Flow steps (`ai.run()`), durable streaming. | -| **Get Started** | `get-started.mdx` | js, go, python | Complete walkthrough for Python. | -| **Interrupts** | `interrupts.mdx` | js, go, python | Python examples for interrupt definition and resumption. | - -### πŸ”Ά Python Tab Exists but Incomplete - -| Page | File | Languages | What's Missing | -|------|------|-----------|----------------| -| **Models** | `models.mdx` | js, go, python | **Generating Media** section (Python SDK supports TTS, image gen via google-genai). **Middleware** section (retry/fallback β€” Python has `use=` support but no docs). | -| **Tool Calling** | `tool-calling.mdx` | js, go, python | **`maxTurns`** β€” Python supports `max_turns=`. **Dynamically defining tools at runtime** β€” needs investigation. **Streaming + Tool calling** β€” needs docs. | -| **Flows** | `flows.mdx` | js, go, python | **Flow steps** (`genkit.Run()` equivalent in Python). **Durable streaming** β€” needs investigation. | -| **RAG** | `rag.mdx` | js, go, python | Python tab may be stale. Verify: indexers, embedders, retrievers, simple retrievers, custom retrievers, rerankers sections. | -| **Context** | `context.mdx` | js only (rendered) | The rendered page shows JS-only. Python supports `context=` on `generate()` and flows. Needs Python examples for context in actions, context at runtime, context propagation. | -| **Dotprompt** | `dotprompt.mdx` | js, go | Python SDK has dotprompt support (`genkit.core.prompt`). Needs full Python tab with: creating prompt files, running prompts, model configuration, schemas, tool calling in prompts, multi-message prompts, partials, prompt variants, defining prompts in code. | - -### ❌ Python Tab Missing (needs to be added) - -| Page | File | Current Languages | Priority | What to Add | -|------|------|-------------------|----------|-------------| -| **Chat (Sessions)** | `chat.mdx` | js, go | **P1** | Python SDK has `ai.chat()` and `Session` with `session.chat()`. Need: session basics, stateful sessions with tools, multi-thread sessions. Session persistence is experimental β€” may tag as such. | -| **Agentic Patterns** | `agentic-patterns.mdx` | js, go | **P1** | Python supports all required primitives (flows, tools, interrupts). Need: sequential workflow, conditional routing, parallel execution, tool calling, iterative refinement, autonomous agent, stateful interactions. | -| **Multi-Agent** | `multi-agent.mdx` | js (likely) | **P2** | Need to verify Python support for agent-to-agent delegation. | -| **Durable Streaming** | `durable-streaming.mdx` | js (likely) | **P3** | Need to verify Python support. | -| **Client SDK** | `client.mdx` | js | **P3** | Client-side integration. May not apply to Python backend SDK directly. | -| **MCP Server** | `mcp-server.mdx` | js (likely) | **P2** | Python has MCP support via `genkit.plugins.mcp`. Needs Python examples. | -| **Model Context Protocol** | `model-context-protocol.mdx` | js (likely) | **P2** | Python MCP client integration. | -| **Evaluation** | `evaluation.mdx` | js (likely) | **P2** | Python has evaluator support (`evaluator-demo` sample exists). Needs Python examples. | - -### βž– Language-Agnostic / Meta Pages - -| Page | File | Notes | -|------|------|-------| -| **Overview** | `overview.mdx` | Conceptual overview, no code tabs needed | -| **API References** | `api-references.mdx` | Links to API docs | -| **API Stability** | `api-stability.mdx` | Policy document | -| **Error Types** | `error-types.mdx` | Reference | -| **Feedback** | `feedback.mdx` | Meta | -| **Develop with AI** | `develop-with-ai.mdx` | Meta/guide | -| **DevTools** | `devtools.mdx` | Dev UI documentation | -| **Local Observability** | `local-observability.mdx` | Observability setup | - ---- - -## 2. Subdirectory Pages (need individual audit) - -| Directory | Path | Known Pages | Python Status | -|-----------|------|-------------|---------------| -| **Deployment** | `deployment/` | Cloud Run, Firebase, etc. | πŸ”Ά Python Flask deployment exists; verify others | -| **Frameworks** | `frameworks/` | Express, etc. | ❌ Need Flask/FastAPI/Starlette Python examples | -| **Integrations** | `integrations/` | Various provider plugins | πŸ”Ά Some Python plugins documented; need audit | -| **Observability** | `observability/` | GCP, custom, etc. | πŸ”Ά Python GCP telemetry plugin exists | -| **Plugin Authoring** | `plugin-authoring/` | Writing plugins | ❌ Need Python plugin authoring guide | -| **Resources** | `resources/` | Additional resources | βž– Likely language-agnostic | -| **Tutorials** | `tutorials/` | Step-by-step guides | ❌ Need Python tutorials | - ---- - -## 3. Feature Parity: JS Examples β†’ Python Samples - -This section maps JS features documented on genkit.dev to their Python sample -coverage status. - -### `/docs/models` Features - -| Feature | JS Example | Python Sample Status | Python SDK Support | -|---------|-----------|---------------------|-------------------| -| Simple generation | `ai.generate('prompt')` | βœ… All hello samples have `say_hi` | βœ… | -| System prompts | `system: "..."` | βœ… Being added to all samples | βœ… | -| Multi-turn (messages) | `messages: [...]` | βœ… Being added to all samples | βœ… | -| Model parameters | `config: {...}` | βœ… `say_hi_with_config` in most samples | βœ… | -| Structured output | `output: { schema: ... }` | βœ… `generate_character` in most samples | βœ… | -| Streaming | `ai.generateStream()` | βœ… `streaming_flow` / `say_hi_stream` | βœ… | -| Streaming + structured | `generateStream() + output schema` | ❌ No dedicated sample | βœ… (need sample) | -| Multimodal input | `prompt: [{media: ...}, {text: ...}]` | βœ… `describe_image` in google-genai, anthropic, xai, msf | βœ… | -| Generating media (images) | `output: { format: 'media' }` | ❌ No dedicated sample | βœ… (google-genai supports it) | -| Generating media (TTS) | Text-to-speech | ❌ No dedicated sample | βœ… (google-genai supports it) | -| Middleware (retry) | `use: [retry({...})]` | ❌ No sample | πŸ”Ά Python has `use=` plumbing, but no retry/fallback middleware defined | -| Middleware (fallback) | `use: [fallback({...})]` | ❌ No sample | πŸ”Ά Same as above | - -### `/docs/tool-calling` Features - -| Feature | JS Example | Python Sample Status | Python SDK Support | -|---------|-----------|---------------------|-------------------| -| Define tools | `ai.defineTool()` | βœ… All samples with tools | βœ… | -| Use tools | `tools: [getWeather]` | βœ… `weather_flow` in most samples | βœ… | -| `maxTurns` | `maxTurns: 8` | βœ… 3 samples use `max_turns=2` | βœ… | -| Dynamic tools at runtime | `tool({...})` | ❌ No sample | ❓ Need investigation | -| Interrupts | `ctx.interrupt()` | βœ… `tool-interrupts`, `google-genai-hello`, `short-n-long` | βœ… | -| `returnToolRequests` | `returnToolRequests: true` | βœ… 1 sample (`google-genai-context-caching`) | βœ… | -| Streaming + tool calling | Stream with tools | ❌ No dedicated sample | βœ… (need sample) | - -### `/docs/interrupts` Features - -| Feature | JS Example | Python Sample Status | Python SDK Support | -|---------|-----------|---------------------|-------------------| -| `defineInterrupt()` | Dedicated interrupt definition | ❌ No equivalent sample | ❓ `define_interrupt` not found in Python SDK | -| `ctx.interrupt()` in tool | Tool-based interrupts | βœ… `tool-interrupts`, `google-genai-hello` | βœ… | -| Restartable interrupts | `restart` option | ❌ No sample | ❓ Need investigation | -| `response.interrupts` check | Interrupt loop | βœ… Demonstrated in `tool-interrupts` | βœ… | -| `resume: { respond: [...] }` | Resume generation | ❌ No sample using `resume` | ❓ Need investigation | - -### `/docs/context` Features - -| Feature | JS Example | Python Sample Status | Python SDK Support | -|---------|-----------|---------------------|-------------------| -| Context in flow | `{context}` destructured | ❌ No sample | βœ… (context available via ActionRunContext) | -| Context in tool | `{context}` in tool handler | ❌ No sample | βœ… (ToolContext has context) | -| Context in prompt file | `{{@auth.name}}` | ❌ No sample | βœ… (dotprompt supports @) | -| Provide context at runtime | `context: { auth: ... }` | ❌ No sample | βœ… (`context=` supported on `generate()`) | -| Context propagation | Auto-propagation to tools | ❌ No sample | βœ… | - -### `/docs/chat` Features - -| Feature | JS Example | Python Sample Status | Python SDK Support | -|---------|-----------|---------------------|-------------------| -| `ai.chat()` basic | Create chat, send messages | ❌ No sample | βœ… `ai.chat()` exists | -| Chat with system prompt | `ai.chat({ system: '...' })` | ❌ No sample | βœ… | -| Stateful sessions | Session with state management | ❌ No sample | βœ… `Session` class exists | -| Multi-thread sessions | Named chat threads | ❌ No sample | βœ… `session.chat('thread')` | -| Session persistence | Custom store implementation | ❌ No sample | πŸ”Ά Experimental | - -### `/docs/flows` Features - -| Feature | JS Example | Python Sample Status | Python SDK Support | -|---------|-----------|---------------------|-------------------| -| Define flows | `@ai.flow()` decorator | βœ… All samples | βœ… | -| Input/output schemas | Pydantic models | βœ… All samples | βœ… | -| Streaming flows | `ctx.send_chunk()` | βœ… Several samples | βœ… | -| Flow steps (`ai.run()`) | Named trace spans | ❌ No sample | ❓ Need investigation | -| Deploy with Flask | Flask integration | βœ… Documented on genkit.dev | βœ… | - -### `/docs/dotprompt` Features - -| Feature | JS Example | Python Sample Status | Python SDK Support | -|---------|-----------|---------------------|-------------------| -| Prompt files (.prompt) | `.prompt` file format | ❌ No sample | βœ… Dotprompt supported | -| Running prompts from code | `ai.prompt('name')` | ❌ No sample | βœ… | -| Input/output schemas | Picoschema/JSON Schema | ❌ No sample | βœ… | -| Tool calling in prompts | `tools: [...]` in frontmatter | ❌ No sample | βœ… | -| Multi-message prompts | `{{role "system"}}` | ❌ No sample | βœ… | -| Partials | `{{> partialName}}` | ❌ No sample | βœ… | -| Prompt variants | `.variant.prompt` files | ❌ No sample | βœ… | -| Defining prompts in code | `ai.define_prompt()` | ❌ No sample | βœ… | - -### `/docs/agentic-patterns` Features - -| Feature | JS Example | Python SDK Support | Sample Needed | -|---------|-----------|-------------|---------------| -| Sequential workflow | Chain of flows | βœ… | ❌ No sample | -| Conditional routing | If/else in flow | βœ… | ❌ No sample | -| Parallel execution | Multiple concurrent calls | βœ… (`asyncio.gather`) | ❌ No sample | -| Tool calling | Tools in generate | βœ… | βœ… Exists | -| Iterative refinement | Loop with evaluation | βœ… | ❌ No sample | -| Autonomous agent | Agent with tools loop | βœ… | ❌ No sample | -| Stateful interactions | Session-based | βœ… | ❌ No sample | - ---- - -## 4. Priority Action Items - -### P0: Critical (blocking feature parity) -1. **`models.mdx`** β€” Add Generating Media section for Python (images, TTS) -2. **`chat.mdx`** β€” Add Python tab with `ai.chat()` and `Session` examples -3. **`context.mdx`** β€” Add Python tab with context in flows, tools, and generate -4. **`dotprompt.mdx`** β€” Add Python tab with full dotprompt examples - -### P1: High Priority (important for developer experience) -5. **`agentic-patterns.mdx`** β€” Add Python tab for all agentic patterns -6. **`tool-calling.mdx`** β€” Add `max_turns` docs, streaming + tools -7. **`models.mdx`** β€” Add Middleware section for Python (investigate retry/fallback) -8. **`evaluation.mdx`** β€” Add Python tab for evaluation -9. **`mcp-server.mdx`** / **`model-context-protocol.mdx`** β€” Add Python MCP examples -10. **Python samples** β€” Add `streaming_structured_output` flow to hello samples - -### P2: Medium Priority (polish) -11. **`flows.mdx`** β€” Add flow steps docs for Python -12. **`multi-agent.mdx`** β€” Add Python tab if SDK supports agent delegation -13. **`frameworks/`** β€” Add Flask/FastAPI/Starlette deployment guides -14. **`plugin-authoring/`** β€” Add Python plugin authoring guide -15. **`interrupts.mdx`** β€” Verify Python section covers `defineInterrupt` equivalent and restartable interrupts - -### P3: Low Priority (nice to have) -16. **`durable-streaming.mdx`** β€” Investigate Python support -17. **`client.mdx`** β€” Determine if applicable to Python -18. **`tutorials/`** β€” Create Python-specific tutorials -19. **`deployment/`** β€” Add Python Cloud Run, etc. deployment guides - ---- - -## 5. Python Samples Gap Analysis - -### Samples needing `system_prompt` flow (in progress) -- [x] `google-genai-hello` -- [x] `compat-oai-hello` -- [x] `anthropic-hello` -- [x] `ollama-hello` -- [x] `amazon-bedrock-hello` -- [x] `deepseek-hello` -- [x] `xai-hello` -- [x] `cloudflare-workers-ai-hello` -- [ ] `microsoft-foundry-hello` -- [ ] `mistral-hello` -- [ ] `huggingface-hello` -- [ ] `google-genai-vertexai-hello` -- [ ] `short-n-long` -- [ ] `model-garden` - -### Samples needing `multi_turn_chat` flow (in progress) -- [x] `google-genai-hello` -- [x] `compat-oai-hello` -- [x] `anthropic-hello` -- [x] `ollama-hello` -- [x] `amazon-bedrock-hello` -- [x] `xai-hello` -- [x] `cloudflare-workers-ai-hello` -- [ ] `microsoft-foundry-hello` -- [ ] `google-genai-vertexai-hello` -- [ ] `short-n-long` -- [ ] `model-garden` - -### New standalone samples needed -- [x] ~~`dotprompt-hello`~~ β€” Covered by `prompt-demo` sample ⚠️ (P1 bug: recursion depth exceeded) -- [ ] ~~`chat-hello`~~ β€” Chat/Session API deprecated, skip -- [ ] ~~`agentic-patterns`~~ β€” Agents not yet in Python SDK, skip -- [ ] `context-demo` β€” Need dedicated context flows (context in generate, flows, tools, propagation, `ai.current_context()`) -- [x] ~~`streaming-structured-output`~~ β€” Covered by `google-genai-hello` / hello samples -- [x] ~~`media-generation`~~ β€” Covered by `media-models-demo` sample -- [ ] `middleware-demo` β€” Custom retry/fallback middleware using `use=` parameter -- [ ] `streaming-tools` β€” Streaming + tool calling flow -- [ ] `eval-pipeline` β€” End-to-end eval: dataset β†’ inference β†’ metrics β†’ results - ---- - -### Dotprompt sample gaps (in `prompt-demo`) -- [ ] Tool calling in prompts (`tools: [...]` in frontmatter) -- [ ] Multimodal prompts (`{{media url=photoUrl}}`) -- [ ] Defining prompts in code (`ai.define_prompt()`) -- [ ] Default input values (`default:` in frontmatter) - ---- - -## 6. Known Bugs - -| Sample | Bug | Severity | -|--------|-----|----------| -| `prompt-demo` | `Failed to load lazy action recipe.robot: maximum recursion depth exceeded` / same for `story` | **P0** β€” Blocks all prompt feature demos | diff --git a/py/docs/types.md b/py/docs/types.md index a9fdce9148..823b4c2259 100644 --- a/py/docs/types.md +++ b/py/docs/types.md @@ -1,115 +1,207 @@ # Types -This module provides all public types for Genkit applications. +Types exported from genkit, genkit.model, genkit.embedder, genkit.plugin_api, and genkit.evaluator. -## Errors +## genkit -::: genkit.types.GenkitError +::: genkit.Genkit -::: genkit.types.ToolInterruptError +::: genkit.Plugin -## Actions +::: genkit.Action -::: genkit.types.ActionRunContext +::: genkit.Flow -## Message and Part Types +::: genkit.ActionKind -::: genkit.core.typing.Message +::: genkit.ActionRunContext -::: genkit.blocks.model.MessageWrapper +::: genkit.ExecutablePrompt -::: genkit.core.typing.Part +::: genkit.PromptGenerateOptions -::: genkit.core.typing.TextPart +::: genkit.ResumeOptions -::: genkit.core.typing.MediaPart +::: genkit.ToolRunContext -::: genkit.core.typing.Media +::: genkit.StreamResponse -::: genkit.core.typing.CustomPart +::: genkit.ModelStreamResponse -::: genkit.core.typing.DataPart +::: genkit.GenkitError -::: genkit.core.typing.ReasoningPart +::: genkit.PublicError -::: genkit.core.typing.Role +::: genkit.ToolInterruptError -::: genkit.core.typing.Metadata +::: genkit.Message -## Tool Types +::: genkit.Part -::: genkit.core.typing.ToolRequest +::: genkit.TextPart -::: genkit.core.typing.ToolRequestPart +::: genkit.MediaPart -::: genkit.core.typing.ToolResponse +::: genkit.Media -::: genkit.core.typing.ToolResponsePart +::: genkit.CustomPart -::: genkit.core.typing.ToolDefinition +::: genkit.ReasoningPart -::: genkit.core.typing.ToolChoice +::: genkit.Role -## Document Types +::: genkit.Metadata -::: genkit.types.Document +::: genkit.ToolRequest -::: genkit.core.typing.DocumentData +::: genkit.ToolRequestPart -## Generation Types +::: genkit.ToolResponse -::: genkit.core.typing.GenerateRequest +::: genkit.ToolResponsePart -::: genkit.core.typing.GenerateResponse +::: genkit.ToolDefinition -::: genkit.blocks.model.GenerateResponseWrapper +::: genkit.ToolChoice -::: genkit.core.typing.GenerateResponseChunk +::: genkit.Document -::: genkit.core.typing.GenerateActionOptions +::: genkit.DocumentPart -::: genkit.core.typing.GenerationCommonConfig +::: genkit.EmbedderRef -::: genkit.core.typing.GenerationUsage +::: genkit.EmbedderOptions -::: genkit.core.typing.OutputConfig +::: genkit.Embedding -::: genkit.core.typing.FinishReason +::: genkit.EmbedRequest -## Embedding Types +::: genkit.EmbedResponse -::: genkit.core.typing.Embedding +::: genkit.ModelRequest -::: genkit.core.typing.EmbedRequest +::: genkit.ModelResponse -::: genkit.core.typing.EmbedResponse +::: genkit.ModelResponseChunk -## Retriever Types +::: genkit.ModelConfig -::: genkit.core.typing.RetrieverRequest +::: genkit.ModelInfo -::: genkit.core.typing.RetrieverResponse +::: genkit.ModelUsage -## Evaluation Types +::: genkit.Constrained -::: genkit.core.typing.BaseEvalDataPoint +::: genkit.Stage -::: genkit.core.typing.EvalRequest +::: genkit.Supports -::: genkit.core.typing.EvalResponse +::: genkit.FinishReason -::: genkit.core.typing.EvalFnResponse +## genkit.model -::: genkit.core.typing.EvalStatusEnum +::: genkit.model.BackgroundAction -::: genkit.core.typing.Score +::: genkit.model.ModelRequest -## Model Info (for Plugin Authors) +::: genkit.model.ModelResponse -::: genkit.core.typing.ModelInfo +::: genkit.model.ModelResponseChunk -::: genkit.core.typing.Supports +::: genkit.model.ModelUsage -::: genkit.core.typing.Constrained +::: genkit.model.Candidate -::: genkit.core.typing.Stage +::: genkit.model.FinishReason + +::: genkit.model.GenerateActionOptions + +::: genkit.model.Error + +::: genkit.model.Operation + +::: genkit.model.ToolRequest + +::: genkit.model.ToolDefinition + +::: genkit.model.ToolResponse + +::: genkit.model.ModelInfo + +::: genkit.model.Supports + +::: genkit.model.Constrained + +::: genkit.model.Stage + +::: genkit.model.ModelRef + +::: genkit.model.ModelConfig + +::: genkit.model.Message + +## genkit.embedder + +::: genkit.embedder.EmbedRequest + +::: genkit.embedder.EmbedResponse + +::: genkit.embedder.Embedding + +::: genkit.embedder.EmbedderRef + +::: genkit.embedder.EmbedderSupports + +::: genkit.embedder.EmbedderOptions + +## genkit.plugin_api + +::: genkit.plugin_api.Plugin + +::: genkit.plugin_api.Action + +::: genkit.plugin_api.ActionMetadata + +::: genkit.plugin_api.ActionKind + +::: genkit.plugin_api.ActionRunContext + +::: genkit.plugin_api.StatusCodes + +::: genkit.plugin_api.StatusName + +::: genkit.plugin_api.GenkitError + +::: genkit.plugin_api.AdjustingTraceExporter + +::: genkit.plugin_api.RedactedSpan + +::: genkit.plugin_api.ModelRef + +::: genkit.plugin_api.EmbedderRef + +::: genkit.plugin_api.EvaluatorRef + +::: genkit.plugin_api.ContextProvider + +::: genkit.plugin_api.RequestData + +## genkit.evaluator + +::: genkit.evaluator.EvalRequest + +::: genkit.evaluator.EvalResponse + +::: genkit.evaluator.EvalFnResponse + +::: genkit.evaluator.Score + +::: genkit.evaluator.Details + +::: genkit.evaluator.BaseEvalDataPoint + +::: genkit.evaluator.BaseDataPoint + +::: genkit.evaluator.EvalStatusEnum + +::: genkit.evaluator.EvaluatorRef diff --git a/py/engdoc/assets/favicon.png b/py/engdoc/assets/favicon.png deleted file mode 100644 index c64e8bd0a5..0000000000 Binary files a/py/engdoc/assets/favicon.png and /dev/null differ diff --git a/py/engdoc/assets/logo.png b/py/engdoc/assets/logo.png deleted file mode 100644 index 0be0d6ff58..0000000000 Binary files a/py/engdoc/assets/logo.png and /dev/null differ diff --git a/py/engdoc/compute-engine.mdx b/py/engdoc/compute-engine.mdx deleted file mode 100644 index 1f40de2f25..0000000000 --- a/py/engdoc/compute-engine.mdx +++ /dev/null @@ -1,116 +0,0 @@ - -

Deploy Genkit to Google Cloud Compute Engine

-A step-by-step guide to setting up a Google Cloud Platform VM, authenticating with GitHub, and running Genkit Python samples. This guide walks you through deploying Genkit on a Google Cloud Compute Engine instance to run the provider-google-genai-hello sample. - -## Create your VM Instance - -First, provision a virtual machine in the Google Cloud Console with the following specifications: - -``` -| Category | Configuration | -|---------------|----------------------------------------------------| -| Machine Name | genkit | -| Region / Zone | us-central1 / us-central1-a | -| Machine Type | E2 (General Purpose) | -| OS | Debian GNU/Linux 12 (bookworm) | -| Storage | 30 GB Balanced Persistent Disk | -| Firewall | Allow HTTP, HTTPS, and Load Balancer health checks | -``` - -## Access the VM via SSH - -Open your Google Cloud Shell and connect to your new instance: - -```bash -gcloud compute ssh genkit --zone us-central1-a -``` - -## Configure Git and GitHub Authentication - - - 1. Install Git: - - ```bash - sudo apt update && sudo apt install git -y - ``` - - 2. Set your Identity: - - ```bash - git config --global user.name "Your Name" - git config --global user.email "youremail@example.com" - ``` - - 3. Generate SSH Key: - - ```bash - ssh-keygen -t ed25519 -C "youremail@example.com" - # Press Enter for all prompts (no passphrase) - ``` - 4. Add Key to SSH Agent: - - ```bash - eval "$(ssh-agent -s)" - ssh-add ~/.ssh/id_ed25519 - ``` - - 5. Link to GitHub: - - - Run `cat ~/.ssh/id_ed25519.pub` and copy the output. Go to GitHub Settings > SSH and GPG keys > New SSH Key, and paste your public key there. - - 6. Test Connection: - - ```bash - ssh -T git@github.com - ``` - - - -## Clone and Initialize Genkit -Navigate to the Genkit repository and run the automated setup script for Python samples. - -```bash -git clone git@github.com:firebase/genkit.git -cd genkit/py/samples -./setup.sh -``` - - During setup: - -- Enter your Gemini API Key when prompted. - -- Enter your GCP Project ID. - -- You can skip other optional prompts. - -## Run the Sample -Follow these steps to launch the Google GenAI hello provider sample: - -```bash -# 1. Load your environment variables -source ~/.environment - -# 2. Navigate to the specific sample -cd provider-google-genai-hello - -# 3. Start the sample -./run.sh -``` - -## Access the Developer UI -To view the Genkit Developer UI from your local browser, use SSH port forwarding. This is more secure than opening a firewall port. - -1. Disconnect from your current SSH session if you are still connected. -2. From your **local machine's terminal**, run the following command to connect to your VM and forward the port: - - ```bash - gcloud compute ssh genkit --zone us-central1-a -- -L 4000:localhost:4000 - ``` - -3. Keep this terminal window open. The SSH connection must remain active for port forwarding to work. -4. Now, open your local web browser and navigate to: - - `http://localhost:4000` - -You should see the Genkit Developer UI. diff --git a/py/engdoc/contributing/git_workflow.md b/py/engdoc/contributing/git_workflow.md deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/engdoc/contributing/index.md b/py/engdoc/contributing/index.md deleted file mode 100644 index e289e625e9..0000000000 --- a/py/engdoc/contributing/index.md +++ /dev/null @@ -1,189 +0,0 @@ -# Getting Started - -!!! note - - If you're a user of Genkit and landed here, - this is engineering documentation that someone contributing - to Genkit would use, not necessarily only use it. - - For more information about how to get started with using - Genkit, please see: [User Guide](.) - -## Preparing your account - -### Create a GitHub account - -1. [Sign up](https://github.com/signup) for an account. - -2. Please [enable 2 factor authentication - (2FA)](https://docs.github.com/en/authentication/securing-your-account-with-two-factor-authentication-2fa) - after having created a GitHub account. - - 1. Use the [Google Authenticator - app](https://support.google.com/accounts/answer/1066447?hl=en&co=GENIE.Platform%3DAndroid) - for your device. - - === "Android" - - [Google Authenticator app for Android](https://play.google.com/store/apps/details?id=com.google.android.apps.authenticator2&hl=en_US&pli=1) - - === "iOS" - - [Google Authenticator app for iOS](https://apps.apple.com/us/app/google-authenticator/id388497605) - - 2. Use a physical security key such as the [Google - Titan](https://store.google.com/product/titan_security_key?hl=en-US). - -4. [Generate an SSH - key](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/generating-a-new-ssh-key-and-adding-it-to-the-ssh-agent) - for your workstation to associate it with your GitHub account. - - !!! tip - - You can consider using a script such as the following to easily identify - your key. - - ```bash - #!/usr/bin/env bash - - UNAME_MRS=$(uname -mrs || echo 'unknown') - UNAME_OS="$(uname -s | tr '[:upper:]' '[:lower:]' || echo 'unknown')" - OS_TYPE="$(echo $OSTYPE | tr '[:upper:]' '[:lower:]' || echo 'unknown')" - CPU_ARCH="$(uname -m || echo 'unknown')" - OS_DISTRO="$(cat /etc/os-release 2>/dev/null | grep -E '^ID=.*' | sed -e 's/^ID=\(.*\)/\1/g')" - - ssh-keygen -vvvv -t ed25519 -C "ssh://${USER}@$(hostname)/?arch=${CPU_ARCH}&os=${OS_TYPE}&distro=${OS_DISTRO}×tamp=$(date +'%Y-%m-%dT%H:%M:%S')" - cat ~/.ssh/id_ed25519.pub >> "${HOME}/.ssh/authorized_keys" - cat ~/.ssh/id_ed25519.pub - ``` - - Save this to a file called `gen_ssh_key.sh` and run it as follows after - changing to its container directory: - - ```bash - chmod a+x gen_ssh_key.sh - ./gen_ssh_key.sh - ``` - -5. [Associate your SSH key with your GitHub account](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/adding-a-new-ssh-key-to-your-github-account). - -6. If you're a Googler, also [associate your `@google.com` email - address](https://docs.github.com/en/account-and-profile/setting-up-and-managing-your-personal-account-on-github/managing-email-preferences/adding-an-email-address-to-your-github-account) - with your GitHub account and follow any other requirements to complete this - process. - -### GitHub Organization Membership - -Please talk to the Genkit Dec team (on [discord](https://discord.gg/qXt5zzQKpc)) -to get yourself added to the appropriate groups for your GitHub organization -membership. - -### CLA - -* Ensure you have [signed the -CLA](https://github.com/firebase/genkit/blob/main/CONTRIBUTING.md#sign-our-contributor-license-agreement). - -* For corporate CLAs, please ask the appropriate channels for help. - -## Preparing your workstation - -!!! note - - === "macOS" - - Install [Homebrew](https://brew.sh/) before proceeding. - - === "Debian" - - Your system should have the required software ready to go. - - === "Fedora" - - Your system should have the required software ready to go. - -### Install the GitHub CLI tool - -=== "macOS" - - We're assuming you have [Homebrew](https://brew.sh/) installed - already. - - - ```bash - brew install gh - ``` - -=== "Debian/Ubuntu" - - ```bash - sudo apt install gh - ``` - -=== "Fedora" - - ```bash - sudo dnf install gh - ``` - -### Authenticate the CLI with GitHub - -```bash -gh auth login -``` - -### Check out project-related repositories - -Consider setting up your workspace to mimic the paths found on GitHub for easier -disambiguation of forks: - -```bash -mkdir -p $HOME/code/github.com/firebase/ -cd $HOME/code/github.com/firebase - -gh repo clone https://github.com/firebase/genkit.git -``` - -This should allow you to produce a directory structure similar to the following: - -```bash -zsh❯ tree -L 3 code - code - └── github.com - β”œβ”€β”€ firebase - β”‚Β Β  └── genkit - β”œβ”€β”€ google - β”‚Β Β  └── dotprompt - └── yesudeep - β”œβ”€β”€ dotprompt - └── genkit -``` - -### Configure Git with your legal name and email address. - -!!! note inline end - - Googlers should use their `@google.com` email address - to make commits. - -```bash -git config user.email "{username}@domain.com" -git config user.Name "Your Legal Name." -``` - - -### Engineering workstations - -Run the following command from the `code/github.com/firebase/genkit` repository -working tree and it should install the required tools for you. - -```bash -py/bin/setup -``` - -### CI/CD - -The following is the equivalent used for CI/CD systems. - -```bash -py/bin/setup -a ci -``` diff --git a/py/engdoc/contributing/installation.md b/py/engdoc/contributing/installation.md deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/engdoc/contributing/troubleshooting.md b/py/engdoc/contributing/troubleshooting.md deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/engdoc/extending/api.md b/py/engdoc/extending/api.md deleted file mode 100644 index 5fad3dd83c..0000000000 --- a/py/engdoc/extending/api.md +++ /dev/null @@ -1,277 +0,0 @@ -# API Design - -Genkit is a framework for building AI-powered applications using generative -models. It provides a streamlined way to work with AI models, tools, prompts, -embeddings, and other AI-related functionality. - -The API is structured to make it easy to: - -* Define prompts that can be reused across your application. -* Create tools that AI models can call. -* Work with different AI models through a consistent interface. -* Build complex AI workflows through "flows". -* Store and retrieve data through embeddings and vector search. - -## Design principles - -Genkit is designed with several principles in mind: - -* **Async-first**: Most communication among us and future interactive agents - appear to be largely naturally asynchronous. -* **Type Safety**: Uses build-time and runtime-type information for strong - typing. -* **Modularity**: Components can be mixed and matched. -* **Extensibility**: Plugin system allows adding new features. -* **Developer Experience**: Development tools like Reflection Server help debug - applications. - -## Veneer - -The veneer refers to the user-facing API and exludes the internals of the -library. - -### The `Genkit` Class - -The `Genkit` class is the central part of the framework that: - -* Manages a registry of AI-related components (models, tools, flows, etc.). -* Provides an API for working with AI models and flows. -* Handles configuration and initialization. -* Sets up development tools like the reflection server. - -#### Key features - -| Feature | Description | -|-------------------------|-------------------------------------------------------------------------------| -| **Registry Management** | It maintains a registry to keep track of all components in a Genkit instance. | -| **Plugin System** | Supports loading plugins to extend functionality. | -| **Prompt Management** | Allows defining and using prompts both programmatically and from files. | -| **Model Integration** | Provides methods to work with generative AI models. | - -#### Core Functionality - -`Genkit` defines methods for the following: - -| Category | Function | Description | -|------------------------|---------------------|-----------------------------------------------| -| **Text Generation** | `generate()` | Generates text using AI models | -| | `generate_stream()` | Streaming version for real-time results | -| **Embedding** | `embed()` | Creates vector embeddings of content | -| | `embed_many()` | Batch embedding generation | -| **Retrieval & Search** | `retrieve()` | Fetches documents based on queries | -| | `index()` | Indexes documents for fast retrieval | -| | `rerank()` | Re-orders retrieved documents by relevance | -| **Tools & Functions** | `define_tool()` | Creates tools that models can use | -| | `define_flow()` | Creates workflows that combine multiple steps | -| **Evaluation** | `evaluate()` | Evaluates AI model outputs | - -#### Helper Functions - -The veneer Genkit module may also include: - -* `genkit()`: A factory function to create new Genkit instances -* `shutdown()`: Handles clean shutdown of Genkit servers -* Event handlers for process termination signals - -## Endpoints - -### Telemetry Server - -| Endpoint | HTTP Method | Purpose | Request Body | Response | Content Type | -|------------------------|-------------|---------------------------|--------------------------------------------|----------------------------------------|--------------------| -| `/api/__health` | GET | Health check | - | "OK" (200) | `text/plain` | -| `/api/traces/:traceId` | GET | Retrieve a specific trace | - | Trace data JSON | `application/json` | -| `/api/traces` | POST | Save a new trace | `TraceData` object | "OK" (200) | `text/plain` | -| `/api/traces` | GET | List traces | Query params: `limit`, `continuationToken` | List of traces with continuation token | `application/json` | - -### Flow Server - -| Endpoint | HTTP Method | Purpose | Request Body | Response | Content Type | -|---------------------------------------|-------------|-------------------------------|---------------------|--------------------------------------------------------------------------------|------------------------| -| `/` | POST | Execute a flow | `{ data: }` | `{ result: }` (200) or error (4xx/5xx) | `application/json` | -| `/?stream=true` | POST | Execute a flow with streaming | `{ data: }` | `data: {"message": }` (stream) and `data: {"result": }` (final) | `text/plain` (chunked) | - -### Reflection Server - -TODO: Ideally, these should behave the same, but we're making a note of -differences here for now. - -=== "TypeScript" - - | Endpoint | HTTP Method | Purpose | Request Body | Response | Content Type | - |------------------------------|-------------|-----------------------------|----------------------------------------------------|--------------------------------------|------------------------| - | `/api/__health` | GET | Health check | - | "OK" (200) | `text/plain` | - | `/api/__quitquitquit` | GET | Terminate server | - | "OK" (200) and server stops | `text/plain` | - | `/api/actions` | GET | List registered actions | - | Action metadata with schemas | `application/json` | - | `/api/runAction` | POST | Run an action | `{ key, input, context, telemetryLabels }` | `{ result, telemetry: { traceId } }` | `application/json` | - | `/api/runAction?stream=true` | POST | Run action with streaming | `{ key, input, context, telemetryLabels }` | Stream of chunks and final result | `text/plain` (chunked) | - | `/api/envs` | GET | Get configured environments | - | List of environment names | `application/json` | - | `/api/notify` | POST | Notify of telemetry server | `{ telemetryServerUrl, reflectionApiSpecVersion }` | "OK" (200) | `text/plain` | - -=== "Go" - - | Endpoint | HTTP Method | Purpose | Request Body | Response | Content Type | - |------------------|-------------|----------------------------|----------------------------------------------------|--------------------------------------|--------------------| - | `/api/__health` | GET | Health check | - | 200 OK status | - | - | `/api/actions` | GET | List registered actions | - | Action metadata with schemas | `application/json` | - | `/api/runAction` | POST | Run an action | `{ key, input, context }` | `{ result, telemetry: { traceId } }` | `application/json` | - | `/api/notify` | POST | Notify of telemetry server | `{ telemetryServerUrl, reflectionApiSpecVersion }` | OK response | `application/json` | - -=== "Python" - - | Endpoint | HTTP Method | Purpose | Request Body | Response | Content Type | - |------------------|-------------|-------------------------|--------------|------------------------------|--------------------| - | `/api/__health` | GET | Health check | - | 200 OK status | - | - | `/api/actions` | GET | List registered actions | - | Action metadata with schemas | `application/json` | - | `/api/runAction` | POST | Run an action | Action input | Action output with traceId | `application/json` | - -## Common Patterns - -* **Health check endpoints** (`/api/__health`): All servers implement a simple - health check endpoint. -* **Action/flow execution**: All servers provide endpoints to execute - actions/flows. -* **Streaming support**: JavaScript-based servers support streaming responses. -* **Telemetry integration**: All execution endpoints include telemetry data - (trace IDs) in responses. -* **Error handling**: Standardized error formats with status codes and stack - traces. -* **Content negotiation**: Different response formats based on accept headers or - query parameters. - -# Async-First Design - -Genkit is a library that allows application developers to create AI flows for -their applications using an API that abstracts over various components such as -indexers, retrievers, models, embedders, etc. - -The API is **async-first** because this single-threaded model of dealing with -concurrency is the direction that Python frameworks are taking and Genkit -naturally lives in an async world. Genkit is majorly I/O-bound, not as much -computationally-bound, since its primary purpose is composing various AI -foundational components and setting up typed communication patterns between them. - -### Class Hierarchy - -The implementation uses a three-level class hierarchy: - -```ascii -+---------------------+ -| GenkitRegistry | (in _registry.py) -|---------------------| -| + flow() | Decorator to register flows -| + tool() | Decorator to register tools -| + define_model() | Register model actions -| + define_embedder() | Register embedder actions -| + registry (prop) | -+--------^------------+ - | -+--------|-----------+ -| GenkitBase | (in _base_async.py) -|--------------------| -| + __init__( | -| plugins, | -| model, | -| reflection_ | -| server_spec) | -+--------^-----------+ - | -+--------|-----------+ -| Genkit | (in _aio.py) -|--------------------| -| + generate() | async β€” text generation -| + generate_stream()| streaming generation -| + embed() | async β€” create embeddings -| + retrieve() | async β€” fetch documents -| + rerank() | async β€” reorder documents -| + evaluate() | async β€” evaluate outputs -| + chat() | session-based chat -+--------------------+ -``` - -```mermaid -classDiagram - class GenkitRegistry { - <<_registry.py>> - +flow(name, description) Callable - +tool(name, description) Callable - +define_model(config, fn) Action - +define_embedder(config, fn) Action - +registry() Registry - } - - class GenkitBase { - <<_base_async.py>> - +__init__(plugins, model, reflection_server_spec) - } - - class Genkit { - <<_aio.py>> - +generate(model, prompt, system, ...) GenerateResponseWrapper - +generate_stream(model, prompt, ...) tuple - +embed(embedder, content) EmbedResponse - +retrieve(retriever, query) list - +rerank(reranker, query, documents) list - +evaluate(evaluator, dataset) EvalResponse - } - - GenkitBase --|> GenkitRegistry : inherits - Genkit --|> GenkitBase : inherits -``` - -All methods on the `Genkit` class are `async`. Synchronously-defined flows and -tools are executed using a thread-pool executor internally. - -### Usage - -```python -from genkit.ai import Genkit -from genkit.plugins.google_genai import GoogleAI - -ai = Genkit( - plugins=[GoogleAI()], - model='googleai/gemini-2.0-flash', -) - -@ai.flow() -async def my_flow(query: str) -> str: - response = await ai.generate(prompt=f"Answer this: {query}") - return response.text -``` - -## Implementation - -The `Genkit` class starts a reflection server when the `GENKIT_ENV` environment -variable has been set to `'dev'`. - -Running the following command: - -```bash -genkit start -- uv run sample.py -``` - -sets `GENKIT_ENV='dev'` within a running instance of `sample.py`. - -`genkit start` exposes a developer UI (usually called dev UI for short) that is -used for debugging and that talks to a reflection API server implemented by the -`Genkit` class instance. The reflection API server provides a way for the dev UI -to allow users to debug their custom flows, test features such as models and -plugins, and also observe traces emitted by these components. - -### Concurrency handling - -The implementation avoids using threads for server infrastructure since asyncio -is primarily a single-threaded design. The reflection server runs as a coroutine -on the same event loop. - -#### Scenarios - -- For simple short-lived applications without dev mode, the program exits - normally after completing all flows. - -- For simple short-lived applications with dev mode (`GENKIT_ENV=dev`), the - reflection server starts and prevents the main thread from exiting to enable - debugging. - -- For long-lived servers, the reflection server attaches to the server manager - alongside any application servers written by the end user. diff --git a/py/engdoc/extending/glossary.md b/py/engdoc/extending/glossary.md deleted file mode 100644 index da27167021..0000000000 --- a/py/engdoc/extending/glossary.md +++ /dev/null @@ -1,14 +0,0 @@ -# Glossary - -| **Term** | Definition | -|------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| **Bi-encoder** | A model that compresses the meaning of a document or query into a single vector. Used in the first stage of retrieval. | -| **Context Stuffing** | Overloading the context window with too much information, which can degrade LLM performance. | -| **Context Window** | The maximum amount of text that an LLM can process at once. | -| **LLM Recall** | The ability of an LLM to find specific information within its context window. | -| **Recall** | A metric that measures how many relevant documents are retrieved in a search. | -| **Reranker (Cross-encoder)** | A model that takes a query and a document as input and outputs a similarity score. This score is used to reorder documents by relevance. | -| **Retrieval Augmented Generation (RAG)** | A technique that combines the power of Large Language Models (LLMs) with external knowledge sources to generate more accurate and comprehensive responses. | -| **Semantic Search** | Searching for information based on the meaning of words and phrases, rather than just matching keywords. | -| **Two-Stage Retrieval** | A system that first retrieves a large set of potentially relevant documents using a fast retriever (like a vector search) and then reranks them using a more accurate similarity score generated by a slower reranker before presenting them to an LLM. This approach combines the speed of the first-stage retrieval with the accuracy of the second-stage reranking, resulting in a more efficient and effective RAG pipeline. | -| **Vector Search** | A technique used to perform semantic search by converting text into numerical vectors and comparing their proximity in a vector space. | diff --git a/py/engdoc/extending/index.md b/py/engdoc/extending/index.md deleted file mode 100644 index c51fe449d4..0000000000 --- a/py/engdoc/extending/index.md +++ /dev/null @@ -1,279 +0,0 @@ -# Framework Architecture - -!!! note - - If you're a user of Genkit and landed here, - this is engineering documentation that someone contributing - to Genkit would use, not necessarily only use it. - - For more information about how to get started with using - Genkit, please see: [User Guide](.) - -Genkit models a generative AI framework allowing application developers -to work with abstractions to allow the use of pluggable implementations of the -various elements of generative AI. It has SDKs for JavaScript, Go, and Python. - -## Design - -![Architecture Layers](../img/onion.svg) - -```d2 -genkit: { - ai: { - style: {fill: "#E0F7FA"} - Veneer API - } - blocks: { - style: {fill: "#FFF3E0"} - AI Components: { - prompt: Prompt - model: Model - embedder: Embedder - retriever: Retriever - } - } - core: { - style: {fill: "#E8F5E9"} - Core Foundations: { - flow: Flow - actions: Actions - registry: Registry - reflection_server: Reflection Server - } - } - plugins: { - style: {fill: "#FCE4EC"} - google_genai - google_cloud - vertex_ai - firebase - ollama - anthropic - amazon_bedrock - cloudflare_workers_ai - cohere - compat_oai - deepseek - huggingface - microsoft_foundry - mistral - xai - observability - checks - evaluators - mcp - fastapi - flask - dev_local_vectorstore - } -} - -lib: { - style: {fill: "#EDE7F6"} - handlebars - dotprompt - pydantic - starlette - asgiref - uvicorn - opentelemetry -} - -genkit.blocks -> genkit.core -genkit.blocks -> lib.dotprompt -genkit.core -> lib.asgiref -genkit.core -> lib.opentelemetry -genkit.core -> lib.pydantic -genkit.core -> lib.starlette -genkit.plugins.chroma -> genkit.ai -genkit.plugins.firebase -> genkit.ai -genkit.plugins.google_cloud -> genkit.ai -genkit.plugins.google_genai -> genkit.ai -genkit.plugins.ollama -> genkit.ai -genkit.plugins.pinecone -> genkit.ai -genkit.ai -> genkit.blocks -genkit.ai -> genkit.core -genkit.ai -> lib.uvicorn -lib.dotprompt -> lib.handlebars -``` - -The framework has several layers of abstraction. Think about it as peeling an -onion, starting from the outermost layer: - -1. **User-facing Veneer** is the topmost layer consisting of the APIs exposed to - the application developers. - -2. **1st and 3rd Party Plugins** that extend the framework with additional - functionality. - -3. **AI Abstractions** are higher level generative AI components (e.g. tools, - agents, rerankers, embedders, vector stores, etc.) built atop the core - foundations. - -4. **Core Foundations** (actions and flows) are the primitive building blocks - upon which everything else has been built. Think of these as Lego bricks. - -5. **OpenTelemetry** is used to enable tracing. - -## User-friendly Veneer - -A **flow** is a remotely callable function that wraps user-defined -functionality. A **plugin** extends the framework by adding additional models, -parsers, retrievers, and other components generally used in AI applications. - -## AI Abstractions - -A **prompt** is an instruction provided to a model. Prompt engineering tweaks -these prompts to attempt to coax the model to do what you want. A **prompt -action** is used to render a prompt template producing a request that can be -passed to a **model**. Prompt actions are defined as either code or as -configuration files bearing the `.prompt` (read "dotprompt") extension. -**Dotprompt files** contain configuration in the form of YAML frontmatter -delineated by `---` sequences followed by variable-interpolated UTF-8 encoded -text templates (e.g. using a templating language such as Handlebars): - -```dotprompt ---- -model: vertexai/gemini-2.5-flash -config: - temperature: 0.9 -input: - schema: - properties: - location: {type: string} - style: {type: string} - name: {type: string} - required: [location] - default: - location: a restaurant ---- - -You are the most welcoming AI assistant and are currently working at {{location}}. - -Greet a guest{{#if name}} named {{name}}{{/if}}{{#if style}} in the style of {{style}}{{/if}}. -``` - -A **tool** is an action that can be used by a flow to complete tasks. An -**agent** can use tools (including other agents) to help automate complex tasks -and workflows. - -## Core Foundations - -An **action** is a locally or remotely (JSON based RPC) callable function. It is -strongly-typed, named, observable, uninterrupted operation that can be in -streaming or non-streaming mode. It wraps a function of type **Fn** that takes -an input of type **In**, and returns an output of type **Out**, optionally -streaming values of type **Stream** incrementally by invoking a -**StreamingCallback**. - -An action is a typed JSON-based RPC-over-HTTP function that supports metadata, -streaming, reflection and discovery. A flow is a user-defined action. An action -can depend on other actions. - -!!! note - - Bidirectional streaming is currently not supported. - -It can have **metadata** like description and specifies the structure for its -input and output using language-native typed schema validation (e.g, Zod for -TypeScript and Pydantic for Python) to generate backing JSON schemas usable by -underlying paraphernalia such as a **model** or **plugin**. Every action is -registered in a central in-memory **registry** exposed via a **RESTful -reflection API** that can be used to look up an action by name, test for -membership, and loaded for execution. - -!!! info "Obsolete definition" - - A **flow** used to be defined as an action that can be paused and resumed later, - even on a different machine. While such a feature is useful for long-running - tasks or situations where resources might become temporarily unavailable, this - role is now handled by **Sessions**. **FlowState** used to store the current - state of the flow, including its input, execution history, and cached steps. - **FlowExecution** used to represent a single run of the Flow's function. A - FlowState contains a history of FlowExecutions. - -Most components of Genkit such as tools, agents, prompts, models, retrievers, -embedders, evaluators, rerankers, and indexers are *framework-defined* actions. -A **flow** is a *user-defined action*. - -**OpenTelemetry** integration is baked into each action. - -### Storage, Search and Indexing Actions - -| Term | Definition | -|------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------| -| **Document** | A piece of **content** consisting of multiple **parts** and **metadata** that can be embedded. | -| **Embedder** | An action that embeds a document. | -| **Embedding** | A way to represent a discrete variable, like a word or a sentence, as a continuous vector. | -| **Evaluator** | An action that \_\_\_\_\_. | -| **Indexer** | An action that indexes content for future matching. | -| **Part** | A component of content that may have a content type (MIME type) and kind (text, media, data, tool request, tool response, etc.) associated with it. | -| **Reranker** (Cross-encoder) | A type of model that outputs a similarity score given a query and document pair. The score is used to reorder the documents by relevance to our query. | -| **Retriever** | An action that fetches documents from storage. | -| **Vector Store** | Stores vector embeddings and can be used to perform **similarity search** using algorithms such as cosine similarity, Euclidean distance, dot product, etc. | - -## Communication - -TODO - -## System Architecture Diagram - -```d2 -vars: { - d2-config: { - layout-engine: elk - theme-id: 300 - sketch: true - } -} -runtime: "Runtime (Go, Python, Node.js, ...)" { - app: User Application { - shape: package - } - - library: "Library" { - plugins: "Plugins\n(pinecone,\ngoogleai,\nvertexai...)" - veneers: "Veneers\n(ai.generateFlow,\nai.defineFlow, ...)" - registry: { - shape: stored_data - style.multiple: true - } - otel: "OpenTelemetry" - actions: "Actions" - reflection_api: "Reflection API" - - plugins -> registry: Define models, etc. - plugins -> actions: Define - veneers -> registry - veneers -> actions - reflection_api -> registry: Lookup - reflection_api -> actions: Run - actions -> otel: Implement - } - - app -> library.plugins: Uses - app -> library.veneers: Uses -} -tooling: "Tooling" { - runtime_manager - dev_console: Developer Console UI { - shape: document - } - telemetry_server: Telemetry Server - trace_store: Trace Store { - shape: cylinder - } - eval_and_dataset_store: Evaluation & Dataset Store { - shape: cylinder - } - - dev_console -> runtime_manager: Uses - dev_console -> telemetry_server: Reports to - dev_console -> eval_and_dataset_store: Reads from and writes to - telemetry_server -> trace_store: Writes to -} - -runtime.library.otel -> tooling.telemetry_server: Reports to -tooling.runtime_manager -> runtime.library.reflection_api: Uses - -``` diff --git a/py/engdoc/extending/servers.md b/py/engdoc/extending/servers.md deleted file mode 100644 index 2153ad4489..0000000000 --- a/py/engdoc/extending/servers.md +++ /dev/null @@ -1,106 +0,0 @@ -# Servers - -Starting a Genkit application or developer CLI spawns several server daemons as -either independent processes, threads, goroutines, or coroutines, depending upon -the runtime. The initialization process deals with: - -* **Environment-based startup**: In development mode (`GENKIT_ENV=dev`), the - Reflection server starts automatically. -* **Port selection**: All servers attempt to find available ports, with - configurable defaults. -* **Registration**: Each server registers itself for cleanup on application - exit. -* **Runtime files**: Reflection servers write metadata files to enable tool - discovery (`/.genkit/runtimes/.json`). -* **Traces**: Traces metadata (`/.genkit/traces`). - -## Types - -| Server Type | Purpose | Implementation | Notes | -|---------------------------|-----------------------------------------------------------------------|-------------------------------------------|---------------------------------------------------------| -| Reflection | Development-time API for inspecting and interacting with Genkit | Both Go and JavaScript | Only starts in development mode (`GENKIT_ENV=dev`) | -| Flow | Exposes registered flows as HTTP endpoints | Go (HTTP server) and JavaScript (Express) | Main server for production environment | -| Dev UI | Web interface for monitoring and interacting with Genkit applications | JavaScript only | Provides dashboard, monitoring, and debugging tools | -| Telemetry | Collects and stores traces of Genkit operations | JavaScript only | Can use local file system or Firestore as backing store | -| Engineering documentation | `mkdocs` instances showing this information | | Engineering documentation | - -## Networking - -| Server | Host | Port | Deployment Environment | -|---------------------------|-------------|------------------------------------------|------------------------| -| Flows | `localhost` | 3400 (override `PORT`) | `'dev'`, `'prod'` | -| Dev UI/Tools API | `localhost` | 4000-4099 | `'dev'`, `'prod'` | -| Reflection API | `localhost` | 3100 (override `GENKIT_REFLECTION_PORT`) | `'dev'` | -| Telemetry | `localhost` | 4033 (specified programmatically) | `'dev'`, `'prod'` | -| Engineering documentation | `localhost` | 8000 | `'dev'` | - -## Implementations - -| Server | Sources | -|------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| Flows | [JS](https://github.com/firebase/genkit/blob/main/js/plugins/express/src/index.ts), [Go](TODO), [Python](https://github.com/firebase/genkit/blob/main/py/packages/genkit/src/genkit/core/flows.py) | -| Telemetry | [JS](https://github.com/firebase/genkit/blob/main/genkit-tools/telemetry-server/src/index.ts) | -| Dev UI/Tools API | [JS](https://github.com/firebase/genkit/blob/main/genkit-tools/common/src/server/server.ts) | -| Reflection | [JS](https://github.com/firebase/genkit/blob/main/js/core/src/reflection.ts), [Go](https://github.com/firebase/genkit/blob/main/go/genkit/reflection.go), [Python](https://github.com/firebase/genkit/blob/main/py/packages/genkit/src/genkit/core/reflection.py) | - -## Environment Variables - -| Environment Variable | Server/Component | Default Value | Description | -|--------------------------------|---------------------------------|---------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `DEBUG` | Logger | `false` | Sets the logging level to `'debug'` instead of `'info'`. | -| `FIREBASE_CONFIG` | Firebase plugins | - | JSON configuration for Firebase services. | -| `FIREBASE_DEBUG_FEATURES` | Firebase plugins | - | List of Firebase features to enable debugging for. | -| `FIREBASE_DEBUG_MODE` | Firebase plugins | `false` | Enables debug mode for Firebase plugins. | -| `FIREBASE_PROJECT_ID` | Firebase plugins | - | Project ID for Firebase services. | -| `FIRESTORE_COLLECTION` | Firebase Firestore | - | Name of Firestore collection to use. | -| `GCLOUD_LOCATION` | Vertex AI | - | Location for Vertex AI services. | -| `GCLOUD_PROJECT` | Google Cloud plugins | - | Project ID for Google Cloud services. | -| `GCLOUD_SERVICE_ACCOUNT_CREDS` | Google Cloud plugins | - | JSON credentials for authenticating with Google Cloud. | -| `GENKIT_ENV` | All servers | `'prod'` | Controls the environment mode. Values: `'dev'` (development) or `'prod'` (production). In `'dev'` mode, additional servers are started such as the Reflection server. | -| `GENKIT_GA_DEBUG` | Analytics | `false` | Enables debug mode for Google Analytics in the dev tools. | -| `GENKIT_GA_VALIDATE` | Analytics | `false` | Enables validation mode for Google Analytics in the dev tools. | -| `GENKIT_REFLECTION_PORT` | Reflection Server | 3100 | The port on which the Reflection API server listens in development mode. | -| `GENKIT_RUNTIME_ID` | Reflection Server | Process ID | Custom identifier for the runtime (see `.genkit/runtimes/.json`). | -| `GENKIT_TELEMETRY_SERVER` | Telemetry Client (Go/JS/Python) | - | URL of the telemetry server to send trace data to. | -| `GOOGLE_API_KEY` | Google APIs | - | General API key for Google services, used as fallback. | -| `GOOGLE_CLOUD_PROJECT` | Google Cloud plugins | - | Alternative name for Project ID for Google Cloud services. | -| `GOOGLE_GENAI_API_KEY` | Google Generative AI | - | API key for Google's generative AI services. | -| `PINECONE_API_KEY` | Pinecone plugin | - | API key for authenticating with Pinecone vector database. | -| `PORT` | Flow Server | 3400 | The port on which the HTTP Flow server listens. | -| `WEAVIATE_API_KEY` | Weaviate plugin | - | API key for authenticating with Weaviate. | -| `WEAVIATE_URL` | Weaviate plugin | - | URL for the Weaviate vector database. | - -## Signal Handling - -Many of these servers handle signals to handle graceful termination and clean up. - -| Signal | Handling | Handlers | -|-----------|-----------------------------------|----------------------| -| `SIGTERM` | Graceful termination | All servers | -| `SIGINT` | Graceful termination and clean up | main Genkit instance | - -!!! note annotate "Common Signals" - - | Signal | Number | Description | Default Action | Notes | - |----------------------|----------|---------------------------------------------------------------------------------------------------------|-----------------------|-----------------------------------------------------------------| - | `SIGHUP` | 1 | Hangup signal. Sent when the controlling terminal closes or a process is terminated. | Terminate | Often used to tell daemons to reload their configuration files. | - | `SIGINT` | 2 | Interrupt signal. Sent when the user presses `Ctrl+C`. | Terminate | Typically used to interrupt a running program. | - | `SIGQUIT` | 3 | Quit signal. Sent when the user presses `Ctrl+\\`. | Terminate (core dump) | Similar to SIGINT, but also generates a core dump. | - | `SIGILL` | 4 | Illegal instruction. Sent when a process attempts to execute an invalid instruction. | Terminate (core dump) | Indicates a programming error. | - | `SIGTRAP` | 5 | Trace/breakpoint trap. Sent when a breakpoint is hit during debugging. | Terminate (core dump) | Used by debuggers. | - | `SIGABRT` (`SIGIOT`) | 6 | Abort signal. Sent when a process calls the `abort()` function. | Terminate (core dump) | Indicates an abnormal termination. | - | `SIGBUS` | 7 or 10 | Bus error. Sent when a process attempts to access memory that is not properly aligned. | Terminate (core dump) | Indicates a hardware or memory error. | - | `SIGFPE` | 8 | Floating-point exception. Sent when a process performs an invalid arithmetic operation. | Terminate (core dump) | Indicates an arithmetic error. | - | `SIGKILL` | 9 | Kill signal. Forces a process to terminate immediately. | Terminate | Cannot be caught or ignored. | - | `SIGUSR1` | 10 or 30 | User-defined signal 1. | Terminate | Can be used for custom signal handling. | - | `SIGSEGV` | 11 | Segmentation violation. Sent when a process attempts to access memory that it is not allowed to access. | Terminate (core dump) | Indicates a memory access error. | - | `SIGUSR2` | 12 or 31 | User-defined signal 2. | Terminate | Can be used for custom signal handling. | - | `SIGPIPE` | 13 | Pipe broken. Sent when a process attempts to write to a pipe that has no readers. | Terminate | Indicates a communication error. | - | `SIGALRM` | 14 | Alarm clock. Sent when a timer expires. | Terminate | Used for timeouts. | - | `SIGTERM` | 15 | Termination signal. Sent by the `kill` command by default. | Terminate | Allows a process to perform cleanup before exiting. | - | `SIGCHLD` | 17 or 20 | Child process status changed. Sent to a parent process when a child process terminates or stops. | Ignore | Used for process management. | - | `SIGCONT` | 18 or 19 | Continue signal. Sent to a stopped process to resume execution. | Continue | Used for job control. | - | `SIGSTOP` | 19 or 17 | Stop signal. Forces a process to stop execution. | Stop | Cannot be caught or ignored. | - | `SIGTSTP` | 20 or 18 | Terminal stop signal. Sent when the user presses `Ctrl+Z`. | Stop | Used for job control. | - | `SIGTTIN` | 21 | Terminal input. Sent to a background process that attempts to read from the terminal. | Stop | Used for job control. | - | `SIGTTOU` | 22 | Terminal output. Sent to a background process that attempts to write to the terminal. | Stop | Used for job control. | diff --git a/py/engdoc/extending/tooling/index.md b/py/engdoc/extending/tooling/index.md deleted file mode 100644 index 07dd0c5c77..0000000000 --- a/py/engdoc/extending/tooling/index.md +++ /dev/null @@ -1 +0,0 @@ -# Overview diff --git a/py/engdoc/img/asgi.svg b/py/engdoc/img/asgi.svg deleted file mode 100644 index c7ed5de144..0000000000 --- a/py/engdoc/img/asgi.svg +++ /dev/null @@ -1,4 +0,0 @@ - - - - \ No newline at end of file diff --git a/py/engdoc/img/onion.svg b/py/engdoc/img/onion.svg deleted file mode 100644 index 509f9679e3..0000000000 --- a/py/engdoc/img/onion.svg +++ /dev/null @@ -1,4 +0,0 @@ - - - - \ No newline at end of file diff --git a/py/engdoc/index.md b/py/engdoc/index.md deleted file mode 100644 index d62a62376a..0000000000 --- a/py/engdoc/index.md +++ /dev/null @@ -1,529 +0,0 @@ -# Genkit - -!!! note - - If you're a user of Genkit and landed here, - this is engineering documentation that someone contributing - to Genkit would use, not necessarily only use it. - - For more information about how to get started with using - Genkit, please see: [User Guide](.) - -## What is Genkit? - -Genkit is a framework designed to help you build AI-powered applications and -features. It provides open source libraries and plus developer -tools for testing and debugging. The following language runtimes are supported: - -| Language Runtime | Version | Support Tier | -|------------------|---------|--------------| -| Node.js | 22.0+ | 1 | -| Go | 1.22+ | 1 | -| Python | 3.10+ | 1 | - -It is designed to work with any generative AI model API or vector database. -While we offer integrations for Firebase and Google Cloud, you can use Genkit -independently of any Google services. - -The framework provides an abstraction of components by wrapping them with -building blocks called actions, each of which is maintained in a registry. An -action can expose a component over HTTP as a cloud function or server endpoint -and is inspectable and discoverable via a reflection API. Flows are actions -defined by the user and plugins can be created by third parties to extend the -set of available actions. - -## Key capabilities - -| Feature | Description | -|-----------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| **Unified API for AI generation** | Use one API to generate or stream content from various AI models. Works with multimodal input/output and custom model settings. | -| **Structured output** | Generate or stream structured objects (like JSON) with built-in validation. Simplify integration with your app and convert unstructured data into a usable format. | -| **Tool calling** | Let AI models call your functions and APIs as tools to complete tasks. The model decides when and which tools to use. | -| **Chat** | Genkit offers a chat-specific API that facilitates multi-turn conversations with AI models, which can be stateful and persistent. | -| **Agents** | Create intelligent agents that use tools (including other agents) to help automate complex tasks and workflows. | -| **Data retrieval** | Improve the accuracy and relevance of generated output by integrating your data. Simple APIs help you embed, index, and retrieve information from various sources. | -| **Prompt templating** | Create effective prompts that include rich text templating, model settings, multimodal support, and tool integration - all within a compact, runnable [prompt file](/docs/genkit/dotprompt). | - -See the following code samples for a concrete idea of how to use these -capabilities in code: - -### Feature Parity - -| Feature | Python | JavaScript | Go | -|-------------------|--------|------------|----| -| Agents | ❌ | βœ… | βœ… | -| Chat | βœ… | βœ… | βœ… | -| Data retrieval | βœ… | βœ… | βœ… | -| Generation | βœ… | βœ… | βœ… | -| Prompt templating | βœ… | βœ… | βœ… | -| Structured output | βœ… | βœ… | βœ… | -| Tool calling | βœ… | βœ… | βœ… | - -### Plugin Parity - -| Plugins | Python | JavaScript | Go | -|------------------------|--------|------------|----| -| Amazon Bedrock | βœ… | β€” | β€” | -| Anthropic | βœ… | β€” | β€” | -| Checks | βœ… | βœ… | β€” | -| Cloudflare Workers AI | βœ… | β€” | β€” | -| Cohere | βœ… | β€” | β€” | -| Compat-OAI | βœ… | β€” | β€” | -| DeepSeek | βœ… | β€” | β€” | -| Dev Local Vectorstore | βœ… | βœ… | β€” | -| Dotprompt | βœ… | βœ… | βœ… | -| Evaluators | βœ… | βœ… | β€” | -| FastAPI | βœ… | β€” | β€” | -| Firebase | βœ… | βœ… | βœ… | -| Flask | βœ… | β€” | β€” | -| Google Cloud | βœ… | βœ… | βœ… | -| Google GenAI | βœ… | βœ… | βœ… | -| Hugging Face | βœ… | β€” | β€” | -| MCP | βœ… | β€” | β€” | -| Microsoft Foundry | βœ… | β€” | β€” | -| Mistral | βœ… | β€” | β€” | -| Observability | βœ… | β€” | β€” | -| Ollama | βœ… | βœ… | β€” | -| Vertex AI | βœ… | βœ… | βœ… | -| xAI | βœ… | β€” | β€” | - -## Examples - -### Basic generation - -=== "Python" - - ```python linenums="1" - import asyncio - - from genkit.ai import Genkit - from genkit.plugins.google_genai import GoogleAI - - ai = Genkit( - plugins=[GoogleAI()], - model='googleai/gemini-2.0-flash', - ) - - - async def main() -> None: - response = await ai.generate(prompt='Why is AI awesome?') - print(response.text) - - stream, _ = ai.generate_stream(prompt='Tell me a story') - async for chunk in stream: - print(chunk.text, end='') - - - if __name__ == '__main__': - asyncio.run(main()) - ``` - -=== "JavaScript" - - ```javascript - import { genkit } from 'genkit'; - import { googleAI, gemini15Flash } from '@genkit-ai/googleai'; - - const ai = genkit({ - plugins: [googleAI()], - model: gemini15Flash, // Set default model - }); - - // Simple generation - const { text } = await ai.generate('Why is AI awesome?'); - console.log(text); - - // Streamed generation - const { stream } = await ai.generateStream('Tell me a story'); - for await (const chunk of stream) { - console.log(chunk.text); - } - ``` - -=== "Go" - - ```go - import "fmt" - - func main() { - fmt.Println("Hello") - } - ``` - -### Structured output - -=== "Python" - - ```python - import asyncio - from enum import Enum - - from pydantic import BaseModel - - from genkit.ai import Genkit, Output - from genkit.plugins.google_genai import GoogleAI - - ai = Genkit( - plugins=[GoogleAI()], - model='googleai/gemini-2.0-flash', - ) - - - class Role(str, Enum): - KNIGHT = "knight" - MAGE = "mage" - ARCHER = "archer" - - - class CharacterProfile(BaseModel): - name: str - role: Role - backstory: str - - - async def main() -> None: - response = await ai.generate( - prompt="Create a brief profile for a character in a fantasy video game.", - output=Output(schema=CharacterProfile), - ) - print(response.output) - - - if __name__ == "__main__": - asyncio.run(main()) - ``` - -=== "JavaScript" - - ```javascript - import { genkit, z } from 'genkit'; - import { googleAI, gemini15Flash } from '@genkit-ai/googleai'; - - const ai = genkit({ - plugins: [googleAI()], - model: gemini15Flash, - }); - - const { output } = await ai.generate({ - prompt: 'Create a brief profile for a character in a fantasy video game.', - // Specify output structure using Zod schema - output: { - format: 'json', - schema: z.object({ - name: z.string(), - role: z.enum(['knight', 'mage', 'archer']), - backstory: z.string(), - }), - }, - - }); - - console.log(output); - ``` - -### Function calling - -=== "Python" - - ```python - import asyncio - - from pydantic import BaseModel, Field - - from genkit.ai import Genkit - from genkit.plugins.google_genai import GoogleAI - - ai = Genkit( - plugins=[GoogleAI()], - model='googleai/gemini-2.0-flash', - ) - - - @ai.tool() - async def get_weather(location: str = Field(description="The location to get the current weather for")) -> str: - """Gets the current weather in a given location.""" - return f"The current weather in {location} is 63Β°F and sunny." - - - async def main() -> None: - response = await ai.generate( - prompt="What is the weather like in New York?", - tools=['get_weather'], - ) - print(response.text) - - - if __name__ == "__main__": - asyncio.run(main()) - ``` - -=== "JavaScript" - - ```javascript - import { genkit, z } from 'genkit'; - import { googleAI, gemini15Flash } from '@genkit-ai/googleai'; - - const ai = genkit({ - plugins: [googleAI()], - model: gemini15Flash, - }); - - // Define tool to get current weather for a given location - const getWeather = ai.defineTool( - { - name: "getWeather", - description: "Gets the current weather in a given location", - inputSchema: z.object({ - location: z.string().describe('The location to get the current weather for') - }), - outputSchema: z.string(), - }, - async (input) => { - // Here, we would typically make an API call or database query. For this - // example, we just return a fixed value. - return `The current weather in ${input.location} is 63Β°F and sunny.`; - } - ); - - const { text } = await ai.generate({ - tools: [getWeather], // Give the model a list of tools it can call - prompt: 'What is the weather like in New York? ', - }); - - console.log(text); - ``` - -### Chat - -=== "Python" - - ```python - import asyncio - - from genkit.ai import Genkit - from genkit.plugins.google_genai import GoogleAI - - ai = Genkit( - plugins=[GoogleAI()], - model='googleai/gemini-2.0-flash', - ) - - - async def main() -> None: - response = await ai.generate( - prompt='Hi, my name is Pavel', - system='Talk like a pirate', - ) - print(response.text) - - response = await ai.generate( - prompt='What is my name?', - system='Talk like a pirate', - messages=response.messages, - ) - print(response.text) - # Ahoy there! Your name is Pavel, you scurvy dog! - - - if __name__ == "__main__": - asyncio.run(main()) - ``` - -=== "JavaScript" - - ```javascript - import { genkit, z } from 'genkit'; - import { googleAI, gemini15Flash } from '@genkit-ai/googleai'; - - const ai = genkit({ - plugins: [googleAI()], - model: gemini15Flash, - }); - - const chat = ai.chat({ system: 'Talk like a pirate' }); - - let response = await chat.send('Hi, my name is Pavel'); - - response = await chat.send('What is my name?'); - console.log(response.text); - // Ahoy there! Your name is Pavel, you scurvy dog - ``` - -### Agents - -=== "Python" - - ```python - # Not yet implemented in Python. - # See: https://github.com/firebase/genkit/pull/4212 - ``` - -=== "JavaScript" - - ```javascript - import { genkit, z } from 'genkit'; - import { googleAI, gemini15Flash } from '@genkit-ai/googleai'; - - const ai = genkit({ - plugins: [googleAI()], - model: gemini15Flash, - }); - - // Define prompts that represent specialist agents - const reservationAgent = ai.definePrompt( - { - name: 'reservationAgent', - description: 'Reservation Agent can help manage guest reservations', - tools: [reservationTool, reservationCancelationTool, reservationListTool], - - }, - `{% verbatim %}{{role "system"}}{% endverbatim %} Help guests make and manage reservations` - ); - - const menuInfoAgent = ... - const complaintAgent = ... - - // Define a triage agent that routes to the proper specialist agent - const triageAgent = ai.definePrompt( - { - name: 'triageAgent', - description: 'Triage Agent', - tools: [reservationAgent, menuInfoAgent, complaintAgent], - }, - `{% verbatim %}{{role "system"}}{% endverbatim %} You are an AI customer service agent for Pavel's Cafe. - Greet the user and ask them how you can help. If appropriate, transfer to an - agent that can better handle the request. If you cannot help the customer - with the available tools, politely explain so.` - ); - - // Create a chat to enable multi-turn agent interactions - const chat = ai.chat(triageAgent); - - chat.send('I want a reservation at Pavel\'s Cafe for noon on Tuesday.' ); - ``` - -### Data retrieval - -=== "Python" - - ```python - import asyncio - - from genkit.ai import Genkit - from genkit.plugins.google_genai import GoogleAI - from genkit.plugins.dev_local_vectorstore import DevLocalVectorstore - - ai = Genkit( - plugins=[ - GoogleAI(), - DevLocalVectorstore( - indexes=[{ - 'index_name': 'BobFacts', - 'embedder': 'googleai/text-embedding-004', - }], - ), - ], - model='googleai/gemini-2.0-flash', - ) - - - async def main() -> None: - query = "How old is Bob?" - - docs = await ai.retrieve( - retriever='devLocalVectorstore/BobFacts', - query=query, - ) - - response = await ai.generate( - prompt=f"Use the provided context to answer: {query}", - docs=docs, - ) - print(response.text) - - - if __name__ == "__main__": - asyncio.run(main()) - ``` - -=== "JavaScript" - - ```javascript - import { genkit } from 'genkit'; - import { googleAI, gemini15Flash, textEmbedding004 } from '@genkit-ai/googleai'; - import { devLocalRetrieverRef } from '@genkit-ai/dev-local-vectorstore'; - - const ai = genkit({ - plugins: [ - googleAI() - devLocalVectorstore([ - { - indexName: 'BobFacts', - embedder: textEmbedding004, - }, - ]), - ], - model: gemini15Flash, - }); - - // Reference to a local vector database storing Genkit documentation - const retriever = devLocalRetrieverRef('BobFacts'); - - // Consistent API to retrieve most relevant documents based on semantic similarity to query - const docs = await ai.retrieve( - retriever: retriever, - query: 'How old is bob?', - ); - - const result = await ai.generate({ - prompt: `Use the provided context from the Genkit documentation to answer this query: ${query}`, - docs // Pass retrieved documents to the model - }); - ``` - -### Prompt template - -=== "YAML" - - ```yaml - --- - model: vertexai/gemini-2.5-flash - config: - temperature: 0.9 - input: - schema: - properties: - location: {type: string} - style: {type: string} - name: {type: string} - required: [location] - default: - location: a restaurant - --- - - You are the most welcoming AI assistant and are currently working at {% - verbatim %}{{location}}{% endverbatim %}. - - Greet a guest{% verbatim %}{{#if name}}{% endverbatim %} named {% verbatim %}{{name}}{% endverbatim %}{% verbatim %}{{/if}}{% endverbatim %}{% verbatim %}{{#if style}}{% endverbatim %} in the style of {% verbatim %}{{style}}{% endverbatim %}{% verbatim %}{{/if}}{% endverbatim %}. - ``` - -## Development tools - -Genkit provides a command-line interface (CLI) and a local Developer UI to make -building AI applications easier. These tools help you: - -* **Experiment:** Test and refine your AI functions, prompts, and queries. -* **Debug:** Find and fix issues with detailed execution traces. -* **Evaluate:** Assess generated results across multiple test cases. - -## Connect with us - -* **Join the community:** Stay updated, ask questions, - and share your work on our [Discord server](https://discord.gg/qXt5zzQKpc). -* **Provide feedback:** Report issues or suggest new features - using our GitHub [issue tracker](https://github.com/firebase/genkit/issues). - -## Next steps - -Learn how to build your first AI application with Genkit in our [Get -started](/docs/get_started.md) guide. diff --git a/py/engdoc/user_guide/developer_tools.md b/py/engdoc/user_guide/developer_tools.md deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/engdoc/user_guide/go/index.md b/py/engdoc/user_guide/go/index.md deleted file mode 100644 index 07dd0c5c77..0000000000 --- a/py/engdoc/user_guide/go/index.md +++ /dev/null @@ -1 +0,0 @@ -# Overview diff --git a/py/engdoc/user_guide/introduction.md b/py/engdoc/user_guide/introduction.md deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/engdoc/user_guide/python/getting_started.md b/py/engdoc/user_guide/python/getting_started.md deleted file mode 100644 index da1e322ef0..0000000000 --- a/py/engdoc/user_guide/python/getting_started.md +++ /dev/null @@ -1,3 +0,0 @@ -# Getting Started - -::: genkit.core.action.types.ActionKind diff --git a/py/engdoc/user_guide/python/index.md b/py/engdoc/user_guide/python/index.md deleted file mode 100644 index d9ee6d6d58..0000000000 --- a/py/engdoc/user_guide/python/index.md +++ /dev/null @@ -1,207 +0,0 @@ -# Overview - -The Genkit Python AI SDK exposes components as remotely callable -functions called **actions** or **flows** via a reflection API using the -Asynchronous Server Gateway Interface -([ASGI](https://asgi.readthedocs.io/en/latest/specs/main.html)) over HTTP. - -An action is a typed JSON-based RPC-over-HTTP function that supports metadata, -streaming, reflection and discovery. A flow is a user-defined action. An action -can depend on other actions. - -!!! note "Caveat" - - While a Python ASGI application implements bidirectional streaming, - the Genkit SDK defines only unidirectional streaming - callables for its actions. - -## Dealing with Concurrency - -Concurrency refers to the ability of our programs to deal with multiple, -potentially different, tasks all at once. This might mean interleaving tasks or -even running them in parallel. Parallelism requires multiple resources for -execution, whereas concurrent code isn't necessarily parallel and can apply even -to a single resource. There are several different ways of dealing with -concurrency among which are the following: - -| Model | Description | -|--------------------------|-------------------------------------------------------------------------------------------| -| Multiprocessing | Running multiple independent processes in a single program | -| Multithreading | Using multiple threads in a single process | -| Asynchronous programming | Handling I/O operations without blocking the main program flow | -| Coroutines | Lightweight resumable cooperatively-scheduled concurrent units of work | -| Goroutines (Go) and CSP | Lightweight, managed concurrent execution units that use messaging over concurrent queues | -| Actors | Independent units of computation that communicate via mailboxes | - -The Genkit Python SDK makes use of coroutines, event loops, and asynchronous I/O -to manage concurrency. For synchronous situations, it makes use of [thread pool -executors](https://docs.python.org/3/library/concurrent.futures.html#threadpoolexecutor). - -!!! info "Runnables, callables, and coroutines" - - A **runnable** or **callable** is a unit of executable work—for example, - threads, processes, green threads, goroutines, and coroutines are runnables. - - **Asynchronous** runnables can wait for multiple steps to complete out of order - as opposed to **synchronous** runnables that require a specific order and cannot - proceed to execute subsequent steps without having completed waiting for any - previous blocking steps. The Python standard library offers an - [asyncio](https://docs.python.org/3/library/asyncio.html) module that works with - cooperatively-scheduled multitasking coroutines to manage I/O asynchronously. - - A **coroutine** is a resumable runnable that can pause its execution and then at - a later point in time resume from where it left off. Coroutines are: - - * **Non-blocking** because they do not block the thread of execution when paused - and yield control to the event loop to switch execution to another. - - * **Stateful** because they store their state when paused and resume from where - they left off and therefore do not require restarting. - - * **Cooperative** because they require the programmer to suspend cooperatively - at well-defined suspend points unlike threads that may be pre-emptively - scheduled. - - * **Cancelable** because Python supports terminating their execution by - bubbling up exceptions. - -## Handling Web Requests: WSGI vs ASGI - -WSGI and ASGI are both standards for Python Web applications, but they differ -significantly in how they handle requests and their overall capabilities. - -### WSGI (Web Server Gateway Interface) - -[WSGI](https://wsgi.readthedocs.io/en/latest/what.html), in essence, is embodied -by: - -!!! quote inline end - - "Call into this function with an HTTP request, and wait until an HTTP response - is returned." - -```python -def app(env, start_response): - start_response('200 OK', [('Content-Type','text/plain')]) - return [b"Hello World"] - - -``` - -WSGI is the yesteryear standard interface for Python Web applications. It is: - -* **Synchronous:** Processes one request at a time. Each request must be fully - completed before the server can move on to the next one. - -* **Blocking:** When a request is being processed, it "blocks" the server from - doing anything else until it's finished. - -* **Limited to HTTP:** WSGI only supports the HTTP protocol. - -### ASGI (Asynchronous Server Gateway Interface) - -ASGI, in essence, is embodied by: - -!!! quote inline end - - "Here’s some information about the connection and the protocol, and a pair of - channels between which the server and application can communicate via event - streams." - -```python -async def app(scope, receive, send): - await logger.adebug('Debugging', scope=pformat(value)) - - await send({ - 'type': 'http.response.start', - 'status_code': 200, - 'headers': [], - }) - - # ... - - buffer = bytearray() - more = True - while more: - msg = await receive() - buffer.extend(await msg['body']) - more = msg.get('more_body', False) - - await send({ - 'type': 'http.response.body', - 'body': buffer, - }) - - # ... -``` - -[ο»ΏASGI](https://asgi.readthedocs.io/en/latest/) (Asynchronous Server Gateway -Interface) is a spiritual successor to WSGI, intended to provide a standard -interface between async-capable Python Web servers, frameworks, and -applications. ASGI decomposes protocols into a series of *events* that an -application must *receive* and react to, and *events* the application might -*send* in response. - -#### Applications and Deployment - -An ASGI application is: - -* **Asynchronous:** Can handle concurrent requests without waiting for each one - to finish and in no particular order. - -* **Non-blocking:** Allows the server to work on other tasks while waiting for a - request to complete (e.g., waiting for data from a database). - -* **Supports HTTP and WebSockets:** ASGI supports both HTTP for traditional Web - requests and WebSockets for real-time, two-way communication. - -* **Protocol-independent application** The ASGI protocol server handles the - nitty-gritties of managing connections, parsing input and rendering output in - the correct protocol format. It communicates with the ASGI application in the - form of event streams. - -An ASGI deployment artifact has 2 components: - -* **Protocol server**, which terminates sockets and translates them into - connections and per-connection event messages. An example of a protocol server - is [Uvicorn](https://www.uvicorn.org/). - -* **Application**, which lives inside a protocol server, is called once per - connection, and handles event messages as they happen, emitting its own event - messages back when necessary. - -![ASGI Deployment Artifact](../../img/asgi.svg) - -While one can write applications using the low-level primitives provided by -the ASGI specification, libraries such as -[FastAPI](https://fastapi.tiangolo.com/), -[Starlette](https://www.starlette.io/), and [Litestar](https://litestar.dev/) -provide more convenient interfaces for application development. - -### Key Differences Summarized - -| Feature | WSGI | ASGI | -|-----------------------|--------------------------------------------------|----------------------------------------------------------------------------| -| **Programming Model** | Synchronous, Blocking | Asynchronous, Non-blocking | -| **Concurrency** | Limited | Excellent | -| **Protocol Support** | HTTP Only | HTTP, WebSocket, HTTP/2 | -| **Use Cases** | Traditional Web applications (blogs, e-commerce) | Real-time applications (chat, online games), high-concurrency applications | - -#### When to Choose WSGI - -* For simpler Web applications that primarily serve content and don't require real-time features. - -* When using frameworks like Flask or Django (although Django now has ASGI support). - -#### When to Choose ASGI - -* For applications that need real-time, bidirectional communication (e.g., chat applications, online games). - -* When dealing with high traffic and concurrency. - -* When using modern asynchronous frameworks like Starlette or FastAPI. - -In essence, ASGI is the more modern and versatile standard, offering improved -performance and capabilities for demanding Web applications. However, WSGI -remains relevant for simpler projects where its limitations aren't a significant -concern. diff --git a/py/engdoc/user_guide/typescript/index.md b/py/engdoc/user_guide/typescript/index.md deleted file mode 100644 index 07dd0c5c77..0000000000 --- a/py/engdoc/user_guide/typescript/index.md +++ /dev/null @@ -1 +0,0 @@ -# Overview diff --git a/py/packages/genkit/README.md b/py/packages/genkit/README.md index 1c2560624b..e1126c9f36 100644 --- a/py/packages/genkit/README.md +++ b/py/packages/genkit/README.md @@ -18,7 +18,7 @@ pip install genkit-plugin-google-genai ```python from pydantic import BaseModel, Field -from genkit.ai import Genkit, Output +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI ai = Genkit( @@ -37,7 +37,7 @@ class RpgCharacter(BaseModel): async def generate_character(name: str) -> RpgCharacter: result = await ai.generate( prompt=f'generate an RPG character named {name}', - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, ) return result.output diff --git a/py/packages/genkit/pyproject.toml b/py/packages/genkit/pyproject.toml index 595930f36b..e9daad3ab4 100644 --- a/py/packages/genkit/pyproject.toml +++ b/py/packages/genkit/pyproject.toml @@ -78,7 +78,6 @@ requires-python = ">=3.10" version = "0.5.1" [project.optional-dependencies] -dev-local-vectorstore = ["genkit-plugin-dev-local-vectorstore"] flask = ["genkit-plugin-flask"] google-cloud = ["genkit-plugin-google-cloud"] google-genai = ["genkit-plugin-google-genai"] @@ -106,7 +105,6 @@ testpaths = ["tests"] [tool.uv.sources] genkit-plugin-compat-oai = { workspace = true } -genkit-plugin-dev-local-vectorstore = { workspace = true } genkit-plugin-flask = { workspace = true } genkit-plugin-google-genai = { workspace = true } genkit-plugin-ollama = { workspace = true } diff --git a/py/packages/genkit/src/genkit/__init__.py b/py/packages/genkit/src/genkit/__init__.py index cec6b784f3..2af08aa319 100644 --- a/py/packages/genkit/src/genkit/__init__.py +++ b/py/packages/genkit/src/genkit/__init__.py @@ -14,103 +14,122 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Genkit - Build AI-powered applications with ease. +"""Genkit β€” Build AI-powered applications.""" -Genkit is an open-source Python toolkit designed to help you build -AI-powered features in web and mobile apps. - -Basic usage: - from genkit import Genkit - from genkit.plugins.google_genai import GoogleAI - - ai = Genkit(plugins=[GoogleAI()]) - - @ai.flow() - async def hello(name: str) -> str: - response = await ai.generate(model="gemini-2.0-flash", prompt=f"Hello {name}") - return response.text -""" - -# Main class -# Re-export everything from genkit.ai for backwards compatibility -from genkit.ai import ( - GENKIT_CLIENT_HEADER, - GENKIT_VERSION, - ActionKind, - ActionRunContext, +from genkit._ai._aio import ActionKind, ActionRunContext, Genkit +from genkit._ai._prompt import ( ExecutablePrompt, - FlowWrapper, - GenerateStreamResponse, - GenkitRegistry, - OutputOptions, + ModelStreamResponse, PromptGenerateOptions, ResumeOptions, - SimpleRetrieverOptions, - ToolRunContext, - tool_response, ) -from genkit.ai._aio import Genkit, Output - -# Core types for convenience (also available from genkit.types) -from genkit.blocks.document import Document -from genkit.blocks.interfaces import Input - -# Response types -from genkit.blocks.model import GenerateResponseWrapper - -# Setup plugin discovery (must be done before any plugin imports) -from genkit.core._plugins import extend_plugin_namespace - -# Errors (user-facing) -from genkit.core.error import GenkitError, UserFacingError - -# Plugin interface -from genkit.core.plugin import Plugin -from genkit.core.typing import ( +from genkit._ai._tools import ToolInterruptError, ToolRunContext, tool_response +from genkit._core._action import Action, StreamResponse +from genkit._core._error import GenkitError, PublicError +from genkit._core._model import Document +from genkit._core._plugin import Plugin +from genkit._core._plugins import extend_plugin_namespace +from genkit._core._typing import ( + CustomPart, + DocumentPart, Media, MediaPart, - Message, + Metadata, Part, + ReasoningPart, Role, TextPart, + ToolChoice, + ToolRequest, + ToolRequestPart, + ToolResponse, + ToolResponsePart, +) + +# Import embedder-related types from the embedder namespace +from genkit.embedder import ( + EmbedderOptions, + EmbedderRef, + Embedding, + EmbedRequest, + EmbedResponse, ) +# Import model-related types from the model namespace. +from genkit.model import ( + Constrained, + FinishReason, + Message, + ModelConfig, + ModelInfo, + ModelRequest, + ModelResponse, + ModelResponseChunk, + ModelUsage, + Stage, + Supports, + ToolDefinition, +) + +# Flow is an alias for Action (used in samples for flow type hints) +Flow = Action + extend_plugin_namespace() __all__ = [ # Main class 'Genkit', - 'Input', - 'Output', + 'Flow', # Response types - 'GenerateResponseWrapper', - 'GenerateStreamResponse', + 'Action', + 'StreamResponse', + 'EmbedRequest', + 'EmbedResponse', + 'EmbedderOptions', + 'EmbedderRef', + 'ModelConfig', + 'ModelInfo', + 'ModelStreamResponse', # Errors 'GenkitError', - 'UserFacingError', - # Core types (convenience) - 'Document', + 'PublicError', + 'ToolInterruptError', + # Content types + 'Constrained', + 'CustomPart', + 'Embedding', + 'Metadata', + 'ReasoningPart', + 'FinishReason', + 'ModelUsage', 'Media', 'MediaPart', 'Message', 'Part', 'Role', + 'Stage', + 'Supports', 'TextPart', + 'ToolChoice', + 'ToolDefinition', + 'ToolRequest', + 'ToolRequestPart', + 'ToolResponse', + 'ToolResponsePart', + # Domain types + 'Document', + 'DocumentPart', # Plugin interface 'Plugin', - # From genkit.ai + # AI runtime 'ActionKind', 'ActionRunContext', 'ExecutablePrompt', - 'FlowWrapper', - 'GenkitRegistry', - 'OutputOptions', 'PromptGenerateOptions', 'ResumeOptions', - 'SimpleRetrieverOptions', 'ToolRunContext', 'tool_response', - # Version info - 'GENKIT_CLIENT_HEADER', - 'GENKIT_VERSION', + 'ModelRequest', + 'ModelResponse', + 'ModelResponseChunk', ] diff --git a/py/packages/genkit/src/genkit/_ai/_aio.py b/py/packages/genkit/src/genkit/_ai/_aio.py new file mode 100644 index 0000000000..a1902bb150 --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_aio.py @@ -0,0 +1,1198 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""User-facing asyncio API for Genkit.""" + +from __future__ import annotations + +import asyncio +import inspect +import json +import signal +import socket +import threading +import uuid +from collections.abc import Awaitable, Callable, Coroutine +from pathlib import Path +from typing import Any, ParamSpec, TypeVar, cast, overload + +import anyio +import uvicorn +from opentelemetry import trace as trace_api +from opentelemetry.sdk.trace import TracerProvider +from pydantic import BaseModel + +from genkit._ai._embedding import EmbedderFn, EmbedderOptions, EmbedderRef, define_embedder +from genkit._ai._evaluator import ( + BatchEvaluatorFn, + EvaluatorFn, + EvaluatorRef, + define_batch_evaluator, + define_evaluator, +) +from genkit._ai._formats import built_in_formats +from genkit._ai._formats._types import FormatDef +from genkit._ai._generate import define_generate_action, generate_action +from genkit._ai._model import ( + Message, + ModelConfig, + ModelFn, + ModelMiddleware, + ModelResponse, + ModelResponseChunk, + define_model, +) +from genkit._ai._prompt import ( + ExecutablePrompt, + ModelStreamResponse, + PromptConfig, + define_helper, + define_partial, + define_schema, + load_prompt_folder, + register_prompt_actions, + to_generate_action_options, +) +from genkit._ai._resource import ( + ResourceFn, + ResourceOptions, + define_resource, +) +from genkit._ai._tools import define_tool +from genkit._core._action import Action, ActionKind, ActionRunContext +from genkit._core._background import ( + BackgroundAction, + CancelModelOpFn, + CheckModelOpFn, + StartModelOpFn, + check_operation, + define_background_model, + lookup_background_action, +) +from genkit._core._channel import Channel, run_loop +from genkit._core._dap import ( + DapFn, + DynamicActionProvider, + define_dynamic_action_provider as define_dap_block, +) +from genkit._core._environment import is_dev_environment +from genkit._core._error import GenkitError +from genkit._core._logger import get_logger +from genkit._core._model import Document +from genkit._core._plugin import Plugin +from genkit._core._reflection import ReflectionServer, ServerSpec, create_reflection_asgi_app +from genkit._core._registry import Registry +from genkit._core._tracing import run_in_new_span +from genkit._core._typing import ( + BaseDataPoint, + Embedding, + EmbedRequest, + EvalRequest, + EvalResponse, + ModelInfo, + Operation, + Part, + SpanMetadata, + ToolChoice, +) + +from ._decorators import _FlowDecorator, _FlowDecoratorWithChunk +from ._runtime import RuntimeManager + +logger = get_logger(__name__) + +# TypeVars for generic input/output typing +InputT = TypeVar('InputT') +OutputT = TypeVar('OutputT') +ChunkT = TypeVar('ChunkT') +P = ParamSpec('P') +R = TypeVar('R') +T = TypeVar('T') + + +def _model_supports_long_running(model_action: Action) -> bool: + """Check if a model action supports long-running operations.""" + model_info = model_action.metadata.get('model') if model_action.metadata else None + if not model_info: + return False + # Handle ModelInfo object + if hasattr(model_info, 'supports'): + supports = getattr(model_info, 'supports', None) + return bool(getattr(supports, 'long_running', False)) if supports else False + # Handle dict (cast needed because isinstance narrows too much for type checkers) + if isinstance(model_info, dict): + model_dict = cast(dict[str, Any], model_info) + supports = model_dict.get('supports') + return bool(supports.get('longRunning', False)) if isinstance(supports, dict) else False + return False + + +class Genkit: + """Genkit asyncio user-facing API.""" + + def __init__( + self, + plugins: list[Plugin] | None = None, + model: str | None = None, + prompt_dir: str | Path | None = None, + reflection_server_spec: ServerSpec | None = None, + ) -> None: + self.registry: Registry = Registry() + self._reflection_server_spec: ServerSpec | None = reflection_server_spec + self._reflection_ready = threading.Event() + self._initialize_registry(model, plugins) + # Ensure the default generate action is registered for async usage. + define_generate_action(self.registry) + # In dev mode, start the reflection server immediately in a background + # daemon thread so it's available regardless of which web framework (or + # none) the user chooses. + if is_dev_environment(): + self._start_reflection_background() + + # Load prompts + load_path = prompt_dir + if load_path is None: + default_prompts_path = Path('./prompts') + if default_prompts_path.is_dir(): + load_path = default_prompts_path + + if load_path: + load_prompt_folder(self.registry, dir_path=load_path) + + # ------------------------------------------------------------------------- + # Registry methods + # ------------------------------------------------------------------------- + + @overload + def flow( + self, + name: str | None = None, + *, + description: str | None = None, + chunk_type: None = None, + ) -> _FlowDecorator: ... + + @overload + def flow( + self, + name: str | None = None, + *, + description: str | None = None, + chunk_type: type[ChunkT], + ) -> _FlowDecoratorWithChunk[ChunkT]: ... + + def flow( + self, + name: str | None = None, + *, + description: str | None = None, + chunk_type: type[Any] | None = None, + ) -> _FlowDecorator | _FlowDecoratorWithChunk[Any]: + """Decorator to register an async function as a flow. + + Args: + name: Optional name for the flow. Defaults to the function name. + description: Optional description for the flow. + chunk_type: Optional type for streaming chunks. When provided, + the returned Action will be typed as Action[InputT, OutputT, ChunkT]. + + Example: + @ai.flow() + async def my_flow(x: str) -> int: ... # Action[str, int] + + @ai.flow(chunk_type=str) + async def streaming_flow(x: int, ctx: ActionRunContext) -> str: + ctx.send_chunk("progress") + return "done" + # Action[int, str, str] + """ + if chunk_type is not None: + return _FlowDecoratorWithChunk(self.registry, name, description, chunk_type) + return _FlowDecorator(self.registry, name, description) + + def define_helper(self, name: str, fn: Callable[..., Any]) -> None: + """Register a Handlebars helper function.""" + define_helper(self.registry, name, fn) + + def define_partial(self, name: str, source: str) -> None: + """Register a Handlebars partial template.""" + define_partial(self.registry, name, source) + + def define_schema(self, name: str, schema: type[BaseModel]) -> type[BaseModel]: + """Register a Pydantic schema for use in prompts.""" + define_schema(self.registry, name, schema) + return schema + + def define_json_schema(self, name: str, json_schema: dict[str, object]) -> dict[str, object]: + """Register a JSON schema for use in prompts.""" + self.registry.register_schema(name, json_schema) + return json_schema + + def define_dynamic_action_provider( + self, + name: str, + fn: DapFn, + *, + description: str | None = None, + cache_ttl_millis: int | None = None, + metadata: dict[str, Any] | None = None, + ) -> DynamicActionProvider: + """Register a Dynamic Action Provider (DAP).""" + return define_dap_block( + self.registry, + name, + fn, + description=description, + cache_ttl_millis=cache_ttl_millis, + metadata=metadata, + ) + + def tool( + self, name: str | None = None, description: str | None = None + ) -> Callable[[Callable[P, T]], Callable[P, T]]: + """Decorator to register a function as a tool.""" + + def wrapper(func: Callable[P, T]) -> Callable[P, T]: + return define_tool(self.registry, func, name, description) + + return wrapper + + def define_evaluator( + self, + *, + name: str, + display_name: str, + definition: str, + fn: EvaluatorFn[Any], + is_billed: bool = False, + config_schema: type[BaseModel] | dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + description: str | None = None, + ) -> Action: + """Register an evaluator action.""" + return define_evaluator( + self.registry, + name=name, + display_name=display_name, + definition=definition, + fn=fn, + is_billed=is_billed, + config_schema=config_schema, + metadata=metadata, + description=description, + ) + + def define_batch_evaluator( + self, + *, + name: str, + display_name: str, + definition: str, + fn: BatchEvaluatorFn[Any], + is_billed: bool = False, + config_schema: type[BaseModel] | dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + description: str | None = None, + ) -> Action: + """Register a batch evaluator action.""" + return define_batch_evaluator( + self.registry, + name=name, + display_name=display_name, + definition=definition, + fn=fn, + is_billed=is_billed, + config_schema=config_schema, + metadata=metadata, + description=description, + ) + + def define_model( + self, + name: str, + fn: ModelFn, + config_schema: type[BaseModel] | dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + info: ModelInfo | None = None, + description: str | None = None, + ) -> Action: + """Register a custom model action.""" + return define_model(self.registry, name, fn, config_schema, metadata, info, description) + + def define_background_model( + self, + name: str, + start: StartModelOpFn, + check: CheckModelOpFn, + cancel: CancelModelOpFn | None = None, + label: str | None = None, + info: ModelInfo | None = None, + config_schema: type[BaseModel] | dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + description: str | None = None, + ) -> BackgroundAction: + """Register a background model for long-running AI operations.""" + return define_background_model( + registry=self.registry, + name=name, + start=start, + check=check, + cancel=cancel, + label=label, + info=info, + config_schema=config_schema, + metadata=metadata, + description=description, + ) + + def define_embedder( + self, + name: str, + fn: EmbedderFn, + options: EmbedderOptions | None = None, + metadata: dict[str, object] | None = None, + description: str | None = None, + ) -> Action: + """Register a custom embedder action.""" + return define_embedder(self.registry, name, fn, options, metadata, description) + + def define_format(self, format: FormatDef) -> None: + """Register a custom output format.""" + self.registry.register_value('format', format.name, format) + + # Overload 1: Both input_schema and output_schema typed -> ExecutablePrompt[InputT, OutputT] + @overload + def define_prompt( + self, + name: str | None = None, + *, + variant: str | None = None, + model: str | None = None, + config: dict[str, object] | ModelConfig | None = None, + description: str | None = None, + system: str | list[Part] | None = None, + prompt: str | list[Part] | None = None, + messages: str | list[Message] | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + max_turns: int | None = None, + return_tool_requests: bool | None = None, + metadata: dict[str, object] | None = None, + tools: list[str] | None = None, + tool_choice: ToolChoice | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + input_schema: type[InputT], + output_schema: type[OutputT], + ) -> ExecutablePrompt[InputT, OutputT]: ... + + # Overload 2: Only input_schema typed -> ExecutablePrompt[InputT, Any] + @overload + def define_prompt( + self, + name: str | None = None, + *, + variant: str | None = None, + model: str | None = None, + config: dict[str, object] | ModelConfig | None = None, + description: str | None = None, + system: str | list[Part] | None = None, + prompt: str | list[Part] | None = None, + messages: str | list[Message] | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + max_turns: int | None = None, + return_tool_requests: bool | None = None, + metadata: dict[str, object] | None = None, + tools: list[str] | None = None, + tool_choice: ToolChoice | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + input_schema: type[InputT], + output_schema: dict[str, object] | str | None = None, + ) -> ExecutablePrompt[InputT, Any]: ... + + # Overload 3: Only output_schema typed -> ExecutablePrompt[Any, OutputT] + @overload + def define_prompt( + self, + name: str | None = None, + *, + variant: str | None = None, + model: str | None = None, + config: dict[str, object] | ModelConfig | None = None, + description: str | None = None, + system: str | list[Part] | None = None, + prompt: str | list[Part] | None = None, + messages: str | list[Message] | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + max_turns: int | None = None, + return_tool_requests: bool | None = None, + metadata: dict[str, object] | None = None, + tools: list[str] | None = None, + tool_choice: ToolChoice | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + input_schema: dict[str, object] | str | None = None, + output_schema: type[OutputT], + ) -> ExecutablePrompt[Any, OutputT]: ... + + # Overload 4: Neither typed -> ExecutablePrompt[Any, Any] + @overload + def define_prompt( + self, + name: str | None = None, + *, + variant: str | None = None, + model: str | None = None, + config: dict[str, object] | ModelConfig | None = None, + description: str | None = None, + system: str | list[Part] | None = None, + prompt: str | list[Part] | None = None, + messages: str | list[Message] | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + max_turns: int | None = None, + return_tool_requests: bool | None = None, + metadata: dict[str, object] | None = None, + tools: list[str] | None = None, + tool_choice: ToolChoice | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + input_schema: dict[str, object] | str | None = None, + output_schema: dict[str, object] | str | None = None, + ) -> ExecutablePrompt[Any, Any]: ... + + def define_prompt( + self, + name: str | None = None, + *, + variant: str | None = None, + model: str | None = None, + config: dict[str, object] | ModelConfig | None = None, + description: str | None = None, + system: str | list[Part] | None = None, + prompt: str | list[Part] | None = None, + messages: str | list[Message] | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + max_turns: int | None = None, + return_tool_requests: bool | None = None, + metadata: dict[str, object] | None = None, + tools: list[str] | None = None, + tool_choice: ToolChoice | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + input_schema: type | dict[str, object] | str | None = None, + output_schema: type | dict[str, object] | str | None = None, + ) -> ExecutablePrompt[Any, Any]: + """Register a prompt template.""" + executable_prompt = ExecutablePrompt( + self.registry, + variant=variant, + model=model, + config=config, + description=description, + input_schema=input_schema, + system=system, + prompt=prompt, + messages=messages, + output_format=output_format, + output_content_type=output_content_type, + output_instructions=output_instructions, + output_schema=output_schema, + output_constrained=output_constrained, + max_turns=max_turns, + return_tool_requests=return_tool_requests, + metadata=metadata, + tools=tools, + tool_choice=tool_choice, + use=use, + docs=docs, + name=name, + ) + if name: + register_prompt_actions(self.registry, executable_prompt, name, variant) + return executable_prompt + + # Overload 1: Neither typed -> ExecutablePrompt[Any, Any] + @overload + def prompt( + self, + name: str, + *, + variant: str | None = None, + input_schema: None = None, + output_schema: None = None, + ) -> ExecutablePrompt[Any, Any]: ... + + # Overload 2: Only input_schema typed + @overload + def prompt( + self, + name: str, + *, + variant: str | None = None, + input_schema: type[InputT], + output_schema: None = None, + ) -> ExecutablePrompt[InputT, Any]: ... + + # Overload 3: Only output_schema typed + @overload + def prompt( + self, + name: str, + *, + variant: str | None = None, + input_schema: None = None, + output_schema: type[OutputT], + ) -> ExecutablePrompt[Any, OutputT]: ... + + # Overload 4: Both input_schema and output_schema typed + @overload + def prompt( + self, + name: str, + *, + variant: str | None = None, + input_schema: type[InputT], + output_schema: type[OutputT], + ) -> ExecutablePrompt[InputT, OutputT]: ... + + def prompt( + self, + name: str, + *, + variant: str | None = None, + input_schema: type[InputT] | None = None, + output_schema: type[OutputT] | None = None, + ) -> ExecutablePrompt[InputT, OutputT] | ExecutablePrompt[Any, Any]: + """Look up a prompt by name and optional variant.""" + return ExecutablePrompt( + registry=self.registry, + name=name, + variant=variant, + input_schema=input_schema, + output_schema=output_schema, + ) + + def define_resource( + self, + *, + fn: ResourceFn, + name: str | None = None, + uri: str | None = None, + template: str | None = None, + description: str | None = None, + metadata: dict[str, object] | None = None, + ) -> Action: + """Register a resource action.""" + opts: ResourceOptions = {} + if name: + opts['name'] = name + if uri: + opts['uri'] = uri + if template: + opts['template'] = template + if description: + opts['description'] = description + if metadata: + opts['metadata'] = metadata + + return define_resource(self.registry, opts, fn) + + # ------------------------------------------------------------------------- + # Server infrastructure methods + # ------------------------------------------------------------------------- + + def _start_reflection_background(self) -> None: + """Start the Dev UI reflection server in a background daemon thread.""" + + async def _run_server() -> None: + sockets: list[socket.socket] | None = None + spec = self._reflection_server_spec + if spec is None: + # Bind to port 0 to let OS choose available port, pass socket to uvicorn + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(('127.0.0.1', 0)) + sock.listen(2048) + host, port = sock.getsockname() + spec = ServerSpec(scheme='http', host=host, port=port) + self._reflection_server_spec = spec + sockets = [sock] + + app = create_reflection_asgi_app(registry=self.registry) + config = uvicorn.Config(app, host=spec.host, port=spec.port, loop='asyncio') + server = ReflectionServer(config, ready=self._reflection_ready) + async with RuntimeManager(spec, lazy_write=True) as runtime_manager: + server_task = asyncio.create_task(server.serve(sockets=sockets)) + await asyncio.to_thread(self._reflection_ready.wait) + + if server.should_exit: + logger.warning(f'Reflection server at {spec.url} failed to start.') + return + + runtime_manager.write_runtime_file() + await logger.ainfo(f'Genkit Dev UI reflection server running at {spec.url}') + await server_task + + threading.Thread( + target=lambda: asyncio.run(_run_server()), + daemon=True, + name='genkit-reflection-server', + ).start() + + def _initialize_registry(self, model: str | None, plugins: list[Plugin] | None) -> None: + """Initialize the registry with default model and plugins.""" + self.registry.default_model = model + for fmt in built_in_formats: + self.define_format(fmt) + + if not plugins: + logger.warning('No plugins provided to Genkit') + else: + for plugin in plugins: + if isinstance(plugin, Plugin): # pyright: ignore[reportUnnecessaryIsInstance] + self.registry.register_plugin(plugin) + else: + raise ValueError(f'Invalid {plugin=} provided to Genkit: must be of type `genkit.ai.Plugin`') + + def run_main(self, coro: Coroutine[Any, Any, T]) -> T | None: + """Run the user's main coroutine, blocking in dev mode for the reflection server.""" + if not is_dev_environment(): + logger.info('Running in production mode.') + return run_loop(coro) + + logger.info('Running in development mode.') + + async def dev_runner() -> T | None: + user_result: T | None = None + try: + user_result = await coro + logger.debug('User coroutine completed successfully.') + except Exception: + logger.exception('User coroutine failed') + + # Block until Ctrl+C (SIGINT handled by anyio) or SIGTERM, keeping + # the daemon reflection thread alive. + logger.info('Script done β€” Dev UI running. Press Ctrl+C to stop.') + try: + async with anyio.create_task_group() as tg: + + async def _handle_sigterm(tg_: anyio.abc.TaskGroup) -> None: # type: ignore[name-defined] + with anyio.open_signal_receiver(signal.SIGTERM) as sigs: + async for _ in sigs: + tg_.cancel_scope.cancel() + return + + tg.start_soon(_handle_sigterm, tg) + await anyio.sleep_forever() + except anyio.get_cancelled_exc_class(): + pass + + logger.info('Dev UI server stopped.') + return user_result + + return anyio.run(dev_runner) + + # ------------------------------------------------------------------------- + # Genkit-specific methods (generation, embedding, retrieval, etc.) + # ------------------------------------------------------------------------- + + def _resolve_embedder_name(self, embedder: str | EmbedderRef | None) -> str: + """Resolve embedder name from string or EmbedderRef.""" + if isinstance(embedder, EmbedderRef): + return embedder.name + elif isinstance(embedder, str): + return embedder + else: + raise ValueError('Embedder must be specified as a string name or an EmbedderRef.') + + # Overload: output_schema=type[T] -> ModelResponse[T] + @overload + async def generate( + self, + *, + model: str | None = None, + prompt: str | list[Part] | None = None, + system: str | list[Part] | None = None, + messages: list[Message] | None = None, + tools: list[str] | None = None, + return_tool_requests: bool | None = None, + tool_choice: ToolChoice | None = None, + tool_responses: list[Part] | None = None, + config: dict[str, object] | ModelConfig | None = None, + max_turns: int | None = None, + context: dict[str, object] | None = None, + output_schema: type[OutputT], + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + ) -> ModelResponse[OutputT]: ... + + # Overload: no output_schema, dict, or union -> ModelResponse[Any] + @overload + async def generate( + self, + *, + model: str | None = None, + prompt: str | list[Part] | None = None, + system: str | list[Part] | None = None, + messages: list[Message] | None = None, + tools: list[str] | None = None, + return_tool_requests: bool | None = None, + tool_choice: ToolChoice | None = None, + tool_responses: list[Part] | None = None, + config: dict[str, object] | ModelConfig | None = None, + max_turns: int | None = None, + context: dict[str, object] | None = None, + output_schema: type | dict | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + ) -> ModelResponse[Any]: ... + + async def generate( + self, + *, + model: str | None = None, + prompt: str | list[Part] | None = None, + system: str | list[Part] | None = None, + messages: list[Message] | None = None, + tools: list[str] | None = None, + return_tool_requests: bool | None = None, + tool_choice: ToolChoice | None = None, + tool_responses: list[Part] | None = None, + config: dict[str, object] | ModelConfig | None = None, + max_turns: int | None = None, + context: dict[str, object] | None = None, + output_schema: type | dict | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + ) -> ModelResponse[Any]: + """Generate text or structured data using a language model.""" + return await generate_action( + self.registry, + await to_generate_action_options( + self.registry, + PromptConfig( + model=model, + prompt=prompt, + system=system, + messages=messages, + tools=tools, + return_tool_requests=return_tool_requests, + tool_choice=tool_choice, + tool_responses=tool_responses, + config=config, + max_turns=max_turns, + output_format=output_format, + output_content_type=output_content_type, + output_instructions=output_instructions, + output_schema=output_schema, + output_constrained=output_constrained, + docs=docs, + ), + ), + middleware=use, + context=context if context else ActionRunContext._current_context(), # pyright: ignore[reportPrivateUsage] + ) + + # Overload: output_schema=type[T] -> ModelStreamResponse[T] + @overload + def generate_stream( + self, + *, + model: str | None = None, + prompt: str | list[Part] | None = None, + system: str | list[Part] | None = None, + messages: list[Message] | None = None, + tools: list[str] | None = None, + return_tool_requests: bool | None = None, + tool_choice: ToolChoice | None = None, + config: dict[str, object] | ModelConfig | None = None, + max_turns: int | None = None, + context: dict[str, object] | None = None, + output_schema: type[OutputT], + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + timeout: float | None = None, + ) -> ModelStreamResponse[OutputT]: ... + + # Overload: no output_schema, dict, or union -> ModelStreamResponse[Any] + @overload + def generate_stream( + self, + *, + model: str | None = None, + prompt: str | list[Part] | None = None, + system: str | list[Part] | None = None, + messages: list[Message] | None = None, + tools: list[str] | None = None, + return_tool_requests: bool | None = None, + tool_choice: ToolChoice | None = None, + config: dict[str, object] | ModelConfig | None = None, + max_turns: int | None = None, + context: dict[str, object] | None = None, + output_schema: type | dict | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + timeout: float | None = None, + ) -> ModelStreamResponse[Any]: ... + + def generate_stream( + self, + *, + model: str | None = None, + prompt: str | list[Part] | None = None, + system: str | list[Part] | None = None, + messages: list[Message] | None = None, + tools: list[str] | None = None, + return_tool_requests: bool | None = None, + tool_choice: ToolChoice | None = None, + config: dict[str, object] | ModelConfig | None = None, + max_turns: int | None = None, + context: dict[str, object] | None = None, + output_schema: type | dict | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + timeout: float | None = None, + ) -> ModelStreamResponse[Any]: + """Stream generated text, returning a ModelStreamResponse with .stream and .response.""" + channel: Channel[ModelResponseChunk, ModelResponse[Any]] = Channel(timeout=timeout) + + async def _run_generate() -> ModelResponse[Any]: + return await generate_action( + self.registry, + await to_generate_action_options( + self.registry, + PromptConfig( + model=model, + prompt=prompt, + system=system, + messages=messages, + tools=tools, + return_tool_requests=return_tool_requests, + tool_choice=tool_choice, + config=config, + max_turns=max_turns, + output_format=output_format, + output_content_type=output_content_type, + output_instructions=output_instructions, + output_schema=output_schema, + output_constrained=output_constrained, + docs=docs, + ), + ), + on_chunk=lambda c: channel.send(c), + middleware=use, + context=context if context else ActionRunContext._current_context(), # pyright: ignore[reportPrivateUsage] + ) + + response_future: asyncio.Future[ModelResponse[Any]] = asyncio.create_task(_run_generate()) + channel.set_close_future(response_future) + + return ModelStreamResponse[Any](channel=channel, response_future=response_future) + + async def embed( + self, + *, + embedder: str | EmbedderRef | None = None, + content: str | Document | None = None, + metadata: dict[str, object] | None = None, + options: dict[str, object] | None = None, + ) -> list[Embedding]: + """Generate vector embeddings for a single document or string.""" + embedder_name = self._resolve_embedder_name(embedder) + embedder_config: dict[str, object] = {} + + # Extract config and version from EmbedderRef (not done for embed_many per JS behavior) + if isinstance(embedder, EmbedderRef): + embedder_config = embedder.config or {} + if embedder.version: + embedder_config['version'] = embedder.version # Handle version from ref + + # Merge options passed to embed() with config from EmbedderRef + final_options = {**(embedder_config or {}), **(options or {})} + + embed_action = await self.registry.resolve_embedder(embedder_name) + if embed_action is None: + raise ValueError(f'Embedder "{embedder_name}" not found') + + if content is None: + raise ValueError('Content must be specified for embedding.') + + documents = [Document.from_text(content, metadata)] if isinstance(content, str) else [content] + + response = ( + await embed_action.run( + EmbedRequest( + input=documents, # pyright: ignore[reportArgumentType] + options=final_options, + ) + ) + ).response + return response.embeddings + + async def embed_many( + self, + *, + embedder: str | EmbedderRef | None = None, + content: list[str] | list[Document] | None = None, + metadata: dict[str, object] | None = None, + options: dict[str, object] | None = None, + ) -> list[Embedding]: + """Generate vector embeddings for multiple documents in a single batch call.""" + if content is None: + raise ValueError('Content must be specified for embedding.') + + # Convert strings to Documents if needed + documents: list[Document] = [ + Document.from_text(item, metadata) if isinstance(item, str) else item for item in content + ] + + # Resolve embedder name (JS embedMany does not extract config/version from ref) + embedder_name = self._resolve_embedder_name(embedder) + + embed_action = await self.registry.resolve_embedder(embedder_name) + if embed_action is None: + raise ValueError(f'Embedder "{embedder_name}" not found') + + response = (await embed_action.run(EmbedRequest(input=documents, options=options))).response # type: ignore[arg-type] + return response.embeddings + + async def evaluate( + self, + evaluator: str | EvaluatorRef | None = None, + dataset: list[BaseDataPoint] | None = None, + options: dict[str, object] | None = None, + eval_run_id: str | None = None, + ) -> EvalResponse: + """Evaluate a dataset using the specified evaluator.""" + evaluator_name: str = '' + evaluator_config: dict[str, object] = {} + + if isinstance(evaluator, EvaluatorRef): + evaluator_name = evaluator.name + evaluator_config = evaluator.config_schema or {} + elif isinstance(evaluator, str): + evaluator_name = evaluator + else: + raise ValueError('Evaluator must be specified as a string name or an EvaluatorRef.') + + final_options = {**(evaluator_config or {}), **(options or {})} + + eval_action = await self.registry.resolve_evaluator(evaluator_name) + if eval_action is None: + raise ValueError(f'Evaluator "{evaluator_name}" not found') + + if not eval_run_id: + eval_run_id = str(uuid.uuid4()) + + if dataset is None: + raise ValueError('Dataset must be specified for evaluation.') + + return ( + await eval_action.run( + EvalRequest( + dataset=dataset, + options=final_options, + eval_run_id=eval_run_id, + ) + ) + ).response + + @staticmethod + def current_context() -> dict[str, Any] | None: + """Get the current execution context, or None if not in an action.""" + return ActionRunContext._current_context() # pyright: ignore[reportPrivateUsage] + + def dynamic_tool( + self, + *, + name: str, + fn: Callable[..., object], + description: str | None = None, + metadata: dict[str, object] | None = None, + ) -> Action: + """Create an unregistered tool action for passing directly to generate().""" + tool_meta: dict[str, object] = metadata.copy() if metadata else {} + tool_meta['type'] = 'tool' + tool_meta['dynamic'] = True + return Action( + kind=ActionKind.TOOL, + name=name, + fn=fn, # type: ignore[arg-type] # dynamic tools may be sync + description=description, + metadata=tool_meta, + ) + + async def flush_tracing(self) -> None: + """Flush all pending trace spans to exporters.""" + provider = trace_api.get_tracer_provider() + if isinstance(provider, TracerProvider): + await asyncio.to_thread(provider.force_flush) + + async def run( + self, + *, + name: str, + fn: Callable[[], Awaitable[T]], + metadata: dict[str, Any] | None = None, + ) -> T: + """Run a function as a discrete traced step within a flow.""" + if not inspect.iscoroutinefunction(fn): + raise TypeError('fn must be a coroutine function') + + span_metadata = SpanMetadata(name=name, metadata=metadata) + with run_in_new_span(span_metadata, labels={'genkit:type': 'flowStep'}) as span: + try: + result = await fn() + output = ( + result.model_dump_json(by_alias=True, exclude_none=True) + if isinstance(result, BaseModel) + else json.dumps(result) + ) + span.set_attribute('genkit:output', output) + return result + except Exception: + # We catch all exceptions here to ensure they are captured by + # the trace span context manager before being re-raised. + # The run_in_new_span context manager handles recording + # the exception details. + raise + + async def check_operation(self, operation: Operation) -> Operation: + """Check the status of a long-running background operation.""" + return await check_operation(self.registry, operation) + + async def cancel_operation(self, operation: Operation) -> Operation: + """Cancel a long-running background operation.""" + if not operation.action: + raise ValueError('Provided operation is missing original request information') + + background_action = await lookup_background_action(self.registry, operation.action) + if background_action is None: + raise ValueError(f'Failed to resolve background action from original request: {operation.action}') + + return await background_action.cancel(operation) + + async def generate_operation( + self, + *, + model: str | None = None, + prompt: str | list[Part] | None = None, + system: str | list[Part] | None = None, + messages: list[Message] | None = None, + tools: list[str] | None = None, + return_tool_requests: bool | None = None, + tool_choice: ToolChoice | None = None, + config: dict[str, object] | ModelConfig | None = None, + max_turns: int | None = None, + context: dict[str, object] | None = None, + output_schema: type | dict | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_constrained: bool | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + ) -> Operation: + """Generate content using a long-running model, returning an Operation to poll.""" + # Resolve the model and check for long_running support + resolved_model = model or self.registry.default_model + if not resolved_model: + raise GenkitError( + status='INVALID_ARGUMENT', + message='No model specified for generate_operation.', + ) + + model_action = await self.registry.resolve_action(ActionKind.MODEL, resolved_model) + if not model_action: + raise GenkitError( + status='NOT_FOUND', + message=f"Model '{resolved_model}' not found.", + ) + + # Check if model supports long-running operations + if not _model_supports_long_running(model_action): + raise GenkitError( + status='INVALID_ARGUMENT', + message=f"Model '{model_action.name}' does not support long running operations.", + ) + + # Call generate + response = await self.generate( + model=model, + prompt=prompt, + system=system, + messages=messages, + tools=tools, + return_tool_requests=return_tool_requests, + tool_choice=tool_choice, + config=config, + max_turns=max_turns, + context=context, + output_schema=output_schema, + output_format=output_format, + output_content_type=output_content_type, + output_instructions=output_instructions, + output_constrained=output_constrained, + use=use, + docs=docs, + ) + + # Extract operation from response + if not hasattr(response, 'operation') or not response.operation: + raise GenkitError( + status='FAILED_PRECONDITION', + message=f"Model '{model_action.name}' did not return an operation.", + ) + + return response.operation diff --git a/py/packages/genkit/src/genkit/_ai/_decorators.py b/py/packages/genkit/src/genkit/_ai/_decorators.py new file mode 100644 index 0000000000..52a69261df --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_decorators.py @@ -0,0 +1,77 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Flow decorator classes for type-safe flow registration.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any, Generic, TypeVar, cast, overload + +from genkit._core._action import Action, ActionRunContext +from genkit._core._flow import define_flow +from genkit._core._registry import Registry + +# TypeVars for generic input/output typing +InputT = TypeVar('InputT') +OutputT = TypeVar('OutputT') +ChunkT = TypeVar('ChunkT') + + +class _FlowDecorator: + """Decorator class for flow registration with proper type inference.""" + + def __init__(self, registry: Registry, name: str | None, description: str | None) -> None: + self._registry = registry + self._name = name + self._description = description + + @overload + def __call__(self, func: Callable[[], Awaitable[OutputT]]) -> Action[None, OutputT]: ... + + @overload + def __call__(self, func: Callable[[InputT], Awaitable[OutputT]]) -> Action[InputT, OutputT]: ... + + @overload + def __call__(self, func: Callable[[InputT, ActionRunContext], Awaitable[OutputT]]) -> Action[InputT, OutputT]: ... + + def __call__(self, func: Callable[..., Awaitable[Any]]) -> Action[Any, Any]: + return define_flow(self._registry, func, self._name, self._description) + + +class _FlowDecoratorWithChunk(Generic[ChunkT]): + """Decorator class for streaming flow registration with chunk type inference.""" + + def __init__(self, registry: Registry, name: str | None, description: str | None, chunk_type: type[ChunkT]) -> None: + self._registry = registry + self._name = name + self._description = description + self._chunk_type = chunk_type + + @overload + def __call__(self, func: Callable[[], Awaitable[OutputT]]) -> Action[None, OutputT, ChunkT]: ... + + @overload + def __call__(self, func: Callable[[InputT], Awaitable[OutputT]]) -> Action[InputT, OutputT, ChunkT]: ... + + @overload + def __call__( + self, func: Callable[[InputT, ActionRunContext], Awaitable[OutputT]] + ) -> Action[InputT, OutputT, ChunkT]: ... + + def __call__(self, func: Callable[..., Awaitable[Any]]) -> Action[Any, Any, ChunkT]: + # Cast is safe: chunk_type is purely for static typing, runtime behavior is identical + return cast(Action[Any, Any, ChunkT], define_flow(self._registry, func, self._name, self._description)) diff --git a/py/packages/genkit/src/genkit/_ai/_embedding.py b/py/packages/genkit/src/genkit/_ai/_embedding.py new file mode 100644 index 0000000000..880b71484d --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_embedding.py @@ -0,0 +1,154 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Embedding types and utilities for Genkit.""" + +from collections.abc import Awaitable, Callable +from typing import Any, ClassVar, cast + +from pydantic import BaseModel, ConfigDict +from pydantic.alias_generators import to_camel +from typing_extensions import Never + +from genkit._core._action import Action, ActionKind, ActionMetadata, get_func_description +from genkit._core._model import Document +from genkit._core._registry import Registry +from genkit._core._schema import to_json_schema +from genkit._core._typing import EmbedRequest, EmbedResponse + + +class EmbedderSupports(BaseModel): + """Embedder capability support.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) + + input: list[str] | None = None + multilingual: bool | None = None + + +class EmbedderOptions(BaseModel): + """Configuration options for an embedder.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, alias_generator=to_camel) + + config_schema: dict[str, Any] | None = None + label: str | None = None + supports: EmbedderSupports | None = None + dimensions: int | None = None + + +class EmbedderRef(BaseModel): + """Reference to an embedder with configuration.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) + + name: str + config: Any | None = None + version: str | None = None + + +class Embedder: + """Runtime embedder wrapper around an embedder Action.""" + + def __init__(self, name: str, action: Action[EmbedRequest, EmbedResponse, Never]) -> None: + """Initialize with embedder name and backing action.""" + self.name: str = name + self._action: Action[EmbedRequest, EmbedResponse, Never] = action + + async def embed( + self, + documents: list[Document], + options: dict[str, Any] | None = None, + ) -> EmbedResponse: + """Generate embeddings for a list of documents.""" + # Document veneer is compatible with DocumentData at runtime + return ( + await self._action.run(EmbedRequest(input=documents, options=options)) # type: ignore[arg-type] + ).response + + +EmbedderFn = Callable[[EmbedRequest], Awaitable[EmbedResponse]] + + +def embedder_action_metadata( + name: str, + options: EmbedderOptions | None = None, +) -> ActionMetadata: + """Create ActionMetadata for an embedder action.""" + options = options if options is not None else EmbedderOptions() + embedder_metadata_dict: dict[str, object] = {'embedder': {}} + embedder_info = cast(dict[str, object], embedder_metadata_dict['embedder']) + + if options.label: + embedder_info['label'] = options.label + + embedder_info['dimensions'] = options.dimensions + + if options.supports: + embedder_info['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True) + + embedder_info['customOptions'] = options.config_schema if options.config_schema else None + + return ActionMetadata( + kind=ActionKind.EMBEDDER, + name=name, + input_json_schema=to_json_schema(EmbedRequest), + output_json_schema=to_json_schema(EmbedResponse), + metadata=embedder_metadata_dict, + ) + + +def create_embedder_ref(name: str, config: dict[str, Any] | None = None, version: str | None = None) -> EmbedderRef: + """Creates an EmbedderRef instance.""" + return EmbedderRef(name=name, config=config, version=version) + + +def define_embedder( + registry: Registry, + name: str, + fn: EmbedderFn, + options: EmbedderOptions | None = None, + metadata: dict[str, object] | None = None, + description: str | None = None, +) -> Action: + """Register a custom embedder action.""" + embedder_meta: dict[str, object] = dict(metadata) if metadata else {} + embedder_info: dict[str, object] + existing_embedder = embedder_meta.get('embedder') + if isinstance(existing_embedder, dict): + embedder_info = {str(key): value for key, value in existing_embedder.items()} + else: + embedder_info = {} + embedder_meta['embedder'] = embedder_info + + if options: + if options.label: + embedder_info['label'] = options.label + if options.dimensions: + embedder_info['dimensions'] = options.dimensions + if options.supports: + embedder_info['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True) + if options.config_schema: + embedder_info['customOptions'] = to_json_schema(options.config_schema) + + embedder_description = get_func_description(fn, description) + return registry.register_action( + name=name, + kind=ActionKind.EMBEDDER, + fn=fn, + metadata=embedder_meta, + description=embedder_description, + ) diff --git a/py/packages/genkit/src/genkit/_ai/_evaluator.py b/py/packages/genkit/src/genkit/_ai/_evaluator.py new file mode 100644 index 0000000000..e9585ab579 --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_evaluator.py @@ -0,0 +1,244 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Evaluator type definitions for the Genkit framework.""" + +import json +import traceback +import uuid +from collections.abc import Callable, Coroutine +from typing import Any, ClassVar, TypeVar, cast + +from pydantic import BaseModel, ConfigDict +from pydantic.alias_generators import to_camel + +from genkit._core._action import Action, ActionKind, ActionMetadata +from genkit._core._logger import get_logger +from genkit._core._registry import Registry +from genkit._core._schema import to_json_schema +from genkit._core._tracing import run_in_new_span +from genkit._core._typing import ( + BaseDataPoint, + EvalFnResponse, + EvalRequest, + EvalResponse, + EvalStatusEnum, + Score, + SpanMetadata, +) + +logger = get_logger(__name__) + +EVALUATOR_METADATA_KEY_DISPLAY_NAME = 'evaluatorDisplayName' +EVALUATOR_METADATA_KEY_DEFINITION = 'evaluatorDefinition' +EVALUATOR_METADATA_KEY_IS_BILLED = 'evaluatorIsBilled' + +T = TypeVar('T') + +# User-provided evaluator function that evaluates a single datapoint. +# Must be async (coroutine function). +EvaluatorFn = Callable[[BaseDataPoint, T], Coroutine[Any, Any, EvalFnResponse]] + +# User-provided batch evaluator function that evaluates an EvaluationRequest +BatchEvaluatorFn = Callable[[EvalRequest, T], Coroutine[Any, Any, list[EvalFnResponse]]] + + +class EvaluatorRef(BaseModel): + """Reference to an evaluator.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, alias_generator=to_camel) + + name: str + config_schema: dict[str, object] | None = None + + +def evaluator_ref(name: str, config_schema: dict[str, object] | None = None) -> EvaluatorRef: + """Create an EvaluatorRef.""" + return EvaluatorRef(name=name, config_schema=config_schema) + + +def evaluator_action_metadata( + name: str, + config_schema: type | dict[str, Any] | None = None, +) -> ActionMetadata: + """Create ActionMetadata for an evaluator action.""" + return ActionMetadata( + kind=ActionKind.EVALUATOR, + name=name, + input_json_schema=to_json_schema(EvalRequest), + output_json_schema=to_json_schema(list[EvalFnResponse]), + metadata={'evaluator': {'customOptions': to_json_schema(config_schema) if config_schema else None}}, + ) + + +def _get_func_description(func: Callable[..., Any], description: str | None = None) -> str: + """Return description if provided, otherwise use the function's docstring.""" + if description is not None: + return description + if func.__doc__ is not None: + return func.__doc__ + return '' + + +def define_evaluator( + registry: Registry, + name: str, + display_name: str, + definition: str, + fn: EvaluatorFn[Any], + is_billed: bool = False, + config_schema: type[BaseModel] | dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + description: str | None = None, +) -> Action: + """Register an evaluator that runs the callback on each dataset sample.""" + evaluator_meta: dict[str, object] = dict(metadata) if metadata else {} + evaluator_info: dict[str, object] + existing_evaluator = evaluator_meta.get('evaluator') + if isinstance(existing_evaluator, dict): + evaluator_info = {str(key): value for key, value in existing_evaluator.items()} + else: + evaluator_info = {} + evaluator_meta['evaluator'] = evaluator_info + evaluator_info[EVALUATOR_METADATA_KEY_DEFINITION] = definition + evaluator_info[EVALUATOR_METADATA_KEY_DISPLAY_NAME] = display_name + evaluator_info[EVALUATOR_METADATA_KEY_IS_BILLED] = is_billed + label_value = evaluator_info.get('label') + if not isinstance(label_value, str) or not label_value: + evaluator_info['label'] = name + if config_schema: + evaluator_info['customOptions'] = to_json_schema(config_schema) + + evaluator_description = _get_func_description(fn, description) + + async def eval_stepper_fn(req: EvalRequest) -> EvalResponse: + eval_responses: list[EvalFnResponse] = [] + for index in range(len(req.dataset)): + datapoint = req.dataset[index] + if datapoint.test_case_id is None: + datapoint.test_case_id = str(uuid.uuid4()) + span_metadata = SpanMetadata( + name=f'Test Case {datapoint.test_case_id}', + metadata={'evaluator:evalRunId': req.eval_run_id}, + ) + try: + # Try to run with tracing, but fallback if tracing infrastructure fails + # (e.g., in environments with NonRecordingSpans like pre-commit) + try: + with run_in_new_span(span_metadata, labels={'genkit:type': 'evaluator'}) as span: + span_id = format(span.get_span_context().span_id, '016x') + trace_id = format(span.get_span_context().trace_id, '032x') + try: + input_json = ( + datapoint.model_dump_json(by_alias=True, exclude_none=True) + if isinstance(datapoint, BaseModel) + else json.dumps(datapoint) + ) + span.set_attribute('genkit:input', input_json) + test_case_output = await fn(datapoint, req.options) + test_case_output.span_id = span_id + test_case_output.trace_id = trace_id + output_json = ( + test_case_output.model_dump_json(by_alias=True, exclude_none=True) + if isinstance(test_case_output, BaseModel) + else json.dumps(test_case_output) + ) + span.set_attribute('genkit:output', output_json) + eval_responses.append(test_case_output) + except Exception as e: + logger.debug(f'eval_stepper_fn error: {e!s}') + logger.debug(traceback.format_exc()) + evaluation = Score( + error=f'Evaluation of test case {datapoint.test_case_id} failed: \n{e!s}', + status=EvalStatusEnum.FAIL, + ) + eval_responses.append( + # The ty type checker only recognizes aliases, so we use them + # to pass both ty check and runtime validation. + EvalFnResponse( + span_id=span_id, + trace_id=trace_id, + test_case_id=datapoint.test_case_id, + evaluation=evaluation, + ) + ) + # Raise to mark span as failed + raise e + except (AttributeError, UnboundLocalError): + # Fallback: run without span + try: + test_case_output = await fn(datapoint, req.options) + eval_responses.append(test_case_output) + except Exception as e: + logger.debug(f'eval_stepper_fn error: {e!s}') + logger.debug(traceback.format_exc()) + evaluation = Score( + error=f'Evaluation of test case {datapoint.test_case_id} failed: \n{e!s}', + status=EvalStatusEnum.FAIL, + ) + eval_responses.append( + EvalFnResponse( + test_case_id=datapoint.test_case_id, + evaluation=evaluation, + ) + ) + except Exception: # noqa: S112 - intentionally continue processing other datapoints + # Continue to process other points + continue + return EvalResponse(eval_responses) + + return registry.register_action( + name=name, + kind=ActionKind.EVALUATOR, + fn=eval_stepper_fn, + metadata=evaluator_meta, + description=evaluator_description, + ) + + +def define_batch_evaluator( + registry: Registry, + name: str, + display_name: str, + definition: str, + fn: BatchEvaluatorFn[Any], + is_billed: bool = False, + config_schema: type[BaseModel] | dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + description: str | None = None, +) -> Action: + """Register a batch evaluator that runs the callback on the entire dataset.""" + evaluator_meta: dict[str, object] = metadata.copy() if metadata else {} + if 'evaluator' not in evaluator_meta: + evaluator_meta['evaluator'] = {} + # Cast to dict for nested operations - pyrefly doesn't narrow nested dict types + evaluator_dict = cast(dict[str, object], evaluator_meta['evaluator']) + evaluator_dict[EVALUATOR_METADATA_KEY_DEFINITION] = definition + evaluator_dict[EVALUATOR_METADATA_KEY_DISPLAY_NAME] = display_name + evaluator_dict[EVALUATOR_METADATA_KEY_IS_BILLED] = is_billed + if 'label' not in evaluator_dict or not evaluator_dict['label']: + evaluator_dict['label'] = name + if config_schema: + evaluator_dict['customOptions'] = to_json_schema(config_schema) + + evaluator_description = _get_func_description(fn, description) + return registry.register_action( + name=name, + kind=ActionKind.EVALUATOR, + fn=fn, + metadata=evaluator_meta, + description=evaluator_description, + ) diff --git a/py/packages/genkit/src/genkit/_ai/_formats/__init__.py b/py/packages/genkit/src/genkit/_ai/_formats/__init__.py new file mode 100644 index 0000000000..3cf457dfe7 --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_formats/__init__.py @@ -0,0 +1,39 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + + +"""Genkit format package. Provides implementation for various formats like json, jsonl, etc.""" + +from genkit._ai._formats._array import ArrayFormat +from genkit._ai._formats._enum import EnumFormat +from genkit._ai._formats._json import JsonFormat +from genkit._ai._formats._jsonl import JsonlFormat +from genkit._ai._formats._text import TextFormat +from genkit._ai._formats._types import FormatDef, Formatter + + +def package_name() -> str: + """Get the fully qualified package name.""" + return 'genkit._ai._formats' + + +built_in_formats = [ + ArrayFormat(), + EnumFormat(), + JsonFormat(), + JsonlFormat(), + TextFormat(), +] diff --git a/py/packages/genkit/src/genkit/blocks/formats/array.py b/py/packages/genkit/src/genkit/_ai/_formats/_array.py similarity index 81% rename from py/packages/genkit/src/genkit/blocks/formats/array.py rename to py/packages/genkit/src/genkit/_ai/_formats/_array.py index d4a6194f2f..a48fa0f93f 100644 --- a/py/packages/genkit/src/genkit/blocks/formats/array.py +++ b/py/packages/genkit/src/genkit/_ai/_formats/_array.py @@ -16,17 +16,17 @@ """Implementation of Array output format.""" +import json from typing import Any -from genkit.blocks.formats.types import FormatDef, Formatter, FormatterConfig -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - MessageWrapper, +from genkit._ai._formats._types import FormatDef, Formatter, FormatterConfig +from genkit._ai._model import ( + Message, + ModelResponseChunk, ) -from genkit.codec import dump_json -from genkit.core._compat import override -from genkit.core.error import GenkitError -from genkit.core.extract import extract_items +from genkit._core._compat import override +from genkit._core._error import GenkitError +from genkit._core._extract_json import extract_json_array_from_text class ArrayFormat(FormatDef): @@ -38,7 +38,7 @@ class ArrayFormat(FormatDef): The formatter automatically handles: 1. Validating that the schema is of type `array`. 2. Injecting instructions into the prompt to output a JSON array. - 3. Parsing the response (both full messages and streaming chunks) using `extract_items` + 3. Parsing the response (both full messages and streaming chunks) using `extract_json_array_from_text` to recover valid JSON objects from potentially incomplete or noisy output. Usage: @@ -91,12 +91,12 @@ def handle(self, schema: dict[str, object] | None) -> Formatter[Any, Any]: message="Must supply an 'array' schema type when using the 'items' parser format.", ) - def message_parser(msg: MessageWrapper) -> list[object]: + def message_parser(msg: Message) -> list[object]: """Parses a complete message into a list of items.""" - result = extract_items(msg.text, 0) + result = extract_json_array_from_text(msg.text, 0) return result.items - def chunk_parser(chunk: GenerateResponseChunkWrapper) -> list[object]: + def chunk_parser(chunk: ModelResponseChunk) -> list[object]: """Parses a streaming chunk into a list of items.""" # Calculate the length of text from previous chunks previous_text_len = len(chunk.accumulated_text) - len(chunk.text) @@ -104,9 +104,9 @@ def chunk_parser(chunk: GenerateResponseChunkWrapper) -> list[object]: # Find cursor position from previous text cursor = 0 if previous_text_len > 0: - cursor = extract_items(chunk.accumulated_text[:previous_text_len]).cursor + cursor = extract_json_array_from_text(chunk.accumulated_text[:previous_text_len]).cursor - result = extract_items(chunk.accumulated_text, cursor) + result = extract_json_array_from_text(chunk.accumulated_text, cursor) return result.items instructions = None @@ -114,7 +114,7 @@ def chunk_parser(chunk: GenerateResponseChunkWrapper) -> list[object]: instructions = f"""Output should be a JSON array conforming to the following schema: ``` -{dump_json(schema, indent=2)} +{json.dumps(schema, indent=2)} ``` """ return Formatter( diff --git a/py/packages/genkit/src/genkit/blocks/formats/enum.py b/py/packages/genkit/src/genkit/_ai/_formats/_enum.py similarity index 90% rename from py/packages/genkit/src/genkit/blocks/formats/enum.py rename to py/packages/genkit/src/genkit/_ai/_formats/_enum.py index bbba93d04e..0c2b0d0d6f 100644 --- a/py/packages/genkit/src/genkit/blocks/formats/enum.py +++ b/py/packages/genkit/src/genkit/_ai/_formats/_enum.py @@ -19,13 +19,13 @@ import re from typing import Any -from genkit.blocks.formats.types import FormatDef, Formatter, FormatterConfig -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - MessageWrapper, +from genkit._ai._formats._types import FormatDef, Formatter, FormatterConfig +from genkit._ai._model import ( + Message, + ModelResponseChunk, ) -from genkit.core._compat import override -from genkit.core.error import GenkitError +from genkit._core._compat import override +from genkit._core._error import GenkitError class EnumFormat(FormatDef): @@ -86,11 +86,11 @@ def handle(self, schema: dict[str, object] | None) -> Formatter[Any, Any]: message="Must supply a schema of type 'string' with an 'enum' property when using the enum format.", ) - def message_parser(msg: MessageWrapper) -> str: + def message_parser(msg: Message) -> str: """Parses a complete message, removing quotes.""" return re.sub(r'[\'"]', '', msg.text).strip() - def chunk_parser(chunk: GenerateResponseChunkWrapper) -> str: + def chunk_parser(chunk: ModelResponseChunk) -> str: """Parses a chunk, removing quotes from accumulated text.""" return re.sub(r'[\'"]', '', chunk.accumulated_text).strip() diff --git a/py/packages/genkit/src/genkit/blocks/formats/json.py b/py/packages/genkit/src/genkit/_ai/_formats/_json.py similarity index 85% rename from py/packages/genkit/src/genkit/blocks/formats/json.py rename to py/packages/genkit/src/genkit/_ai/_formats/_json.py index 1d5f46231b..ee9b85120b 100644 --- a/py/packages/genkit/src/genkit/blocks/formats/json.py +++ b/py/packages/genkit/src/genkit/_ai/_formats/_json.py @@ -16,16 +16,16 @@ """Implementation of JSON output format.""" +import json from typing import Any -from genkit.blocks.formats.types import FormatDef, Formatter, FormatterConfig -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - MessageWrapper, +from genkit._ai._formats._types import FormatDef, Formatter, FormatterConfig +from genkit._ai._model import ( + Message, + ModelResponseChunk, ) -from genkit.codec import dump_json -from genkit.core._compat import override -from genkit.core.extract import extract_json +from genkit._core._compat import override +from genkit._core._extract_json import extract_json class JsonFormat(FormatDef): @@ -78,7 +78,7 @@ def handle(self, schema: dict[str, object] | None) -> Formatter[Any, Any]: the provided schema. """ - def message_parser(msg: MessageWrapper) -> object: + def message_parser(msg: Message) -> object: """Extracts JSON from a Message object. Concatenates the text content of all parts in the message and @@ -92,15 +92,15 @@ def message_parser(msg: MessageWrapper) -> object: """ return extract_json(msg.text) - def chunk_parser(chunk: GenerateResponseChunkWrapper) -> object: - """Extracts JSON from a GenerateResponseChunkWrapper object. + def chunk_parser(chunk: ModelResponseChunk) -> object: + """Extracts JSON from a ModelResponseChunk object. Extracts a JSON object from the accumulated text in the given chunk. Returns None if no valid JSON is found yet (common during streaming when receiving preamble text). Args: - chunk: The GenerateResponseChunkWrapper object to parse. + chunk: The ModelResponseChunk object to parse. Returns: A JSON object extracted from the chunk's accumulated text, @@ -115,7 +115,7 @@ def chunk_parser(chunk: GenerateResponseChunkWrapper) -> object: Output should be in JSON format and conform to the following schema: ``` -{dump_json(schema, indent=2)} +{json.dumps(schema, indent=2)} ``` """ diff --git a/py/packages/genkit/src/genkit/blocks/formats/jsonl.py b/py/packages/genkit/src/genkit/_ai/_formats/_jsonl.py similarity index 90% rename from py/packages/genkit/src/genkit/blocks/formats/jsonl.py rename to py/packages/genkit/src/genkit/_ai/_formats/_jsonl.py index b95387cc8f..ef23db4fea 100644 --- a/py/packages/genkit/src/genkit/blocks/formats/jsonl.py +++ b/py/packages/genkit/src/genkit/_ai/_formats/_jsonl.py @@ -16,19 +16,19 @@ """Implementation of JSONL output format.""" +import json from typing import Any, cast import json5 -from genkit.blocks.formats.types import FormatDef, Formatter, FormatterConfig -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - MessageWrapper, +from genkit._ai._formats._types import FormatDef, Formatter, FormatterConfig +from genkit._ai._model import ( + Message, + ModelResponseChunk, ) -from genkit.codec import dump_json -from genkit.core._compat import override -from genkit.core.error import GenkitError -from genkit.core.extract import extract_json +from genkit._core._compat import override +from genkit._core._error import GenkitError +from genkit._core._extract_json import extract_json class JsonlFormat(FormatDef): @@ -98,7 +98,7 @@ def handle(self, schema: dict[str, object] | None) -> Formatter[Any, Any]: ), ) - def message_parser(msg: MessageWrapper) -> list[object]: + def message_parser(msg: Message) -> list[object]: """Parses a complete message into a list of objects.""" lines = [line.strip() for line in msg.text.split('\n') if line.strip().startswith('{')] items = [] @@ -108,7 +108,7 @@ def message_parser(msg: MessageWrapper) -> list[object]: items.append(extracted) return items - def chunk_parser(chunk: GenerateResponseChunkWrapper) -> list[object]: + def chunk_parser(chunk: ModelResponseChunk) -> list[object]: """Parses a streaming chunk into a list of objects found in that chunk.""" # Calculate the length of text from previous chunks previous_text_len = len(chunk.accumulated_text) - len(chunk.text) @@ -143,7 +143,7 @@ def chunk_parser(chunk: GenerateResponseChunkWrapper) -> list[object]: 'Output should be JSONL format, a sequence of JSON objects (one per line) ' 'separated by a newline `\\n` character. Each line should be a JSON object ' 'conforming to the following schema:\n\n' - f'```\n{dump_json(schema["items"], indent=2)}\n```\n' + f'```\n{json.dumps(schema["items"], indent=2)}\n```\n' ) return Formatter( chunk_parser=chunk_parser, diff --git a/py/packages/genkit/src/genkit/blocks/formats/text.py b/py/packages/genkit/src/genkit/_ai/_formats/_text.py similarity index 82% rename from py/packages/genkit/src/genkit/blocks/formats/text.py rename to py/packages/genkit/src/genkit/_ai/_formats/_text.py index 983cd6d367..792c231884 100644 --- a/py/packages/genkit/src/genkit/blocks/formats/text.py +++ b/py/packages/genkit/src/genkit/_ai/_formats/_text.py @@ -18,12 +18,12 @@ from typing import Any -from genkit.blocks.formats.types import FormatDef, Formatter, FormatterConfig -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - MessageWrapper, +from genkit._ai._formats._types import FormatDef, Formatter, FormatterConfig +from genkit._ai._model import ( + Message, + ModelResponseChunk, ) -from genkit.core._compat import override +from genkit._core._compat import override class TextFormat(FormatDef): @@ -63,7 +63,7 @@ def handle(self, schema: dict[str, object] | None) -> Formatter[Any, Any]: A Formatter instance configured for text handling. """ - def message_parser(msg: MessageWrapper) -> str: + def message_parser(msg: Message) -> str: """Extracts text from a Message object. Args: @@ -74,11 +74,11 @@ def message_parser(msg: MessageWrapper) -> str: """ return msg.text - def chunk_parser(chunk: GenerateResponseChunkWrapper) -> str: - """Extracts text from a GenerateResponseChunkWrapper object. + def chunk_parser(chunk: ModelResponseChunk) -> str: + """Extracts text from a ModelResponseChunk object. Args: - chunk: The GenerateResponseChunkWrapper object. + chunk: The ModelResponseChunk object. Returns: The text content from the current chunk only. diff --git a/py/packages/genkit/src/genkit/blocks/formats/types.py b/py/packages/genkit/src/genkit/_ai/_formats/_types.py similarity index 87% rename from py/packages/genkit/src/genkit/blocks/formats/types.py rename to py/packages/genkit/src/genkit/_ai/_formats/_types.py index 929761fde0..f959913221 100644 --- a/py/packages/genkit/src/genkit/blocks/formats/types.py +++ b/py/packages/genkit/src/genkit/_ai/_formats/_types.py @@ -20,11 +20,11 @@ from collections.abc import Callable from typing import Any, Generic, TypeVar -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - MessageWrapper, +from genkit._ai._model import ( + Message, + ModelResponseChunk, ) -from genkit.core.typing import ( +from genkit._core._typing import ( OutputConfig, ) @@ -53,22 +53,22 @@ class Formatter(Generic[OutputT, ChunkT]): def __init__( self, - message_parser: Callable[[MessageWrapper], OutputT], - chunk_parser: Callable[[GenerateResponseChunkWrapper], ChunkT], + message_parser: Callable[[Message], OutputT], + chunk_parser: Callable[[ModelResponseChunk], ChunkT], instructions: str | None, ) -> None: """Initializes a Formatter. Args: message_parser: A callable that parses a Message into type OutputT. - chunk_parser: A callable that parses a GenerateResponseChunkWrapper into type ChunkT. + chunk_parser: A callable that parses a ModelResponseChunk into type ChunkT. instructions: Optional instructions for the formatter. """ self.instructions: str | None = instructions self.__message_parser = message_parser self.__chunk_parser = chunk_parser - def parse_message(self, message: MessageWrapper) -> OutputT: + def parse_message(self, message: Message) -> OutputT: """Parses a message. Args: @@ -79,7 +79,7 @@ def parse_message(self, message: MessageWrapper) -> OutputT: """ return self.__message_parser(message) - def parse_chunk(self, chunk: GenerateResponseChunkWrapper) -> ChunkT: + def parse_chunk(self, chunk: ModelResponseChunk) -> ChunkT: """Parses a chunk. Args: diff --git a/py/packages/genkit/src/genkit/blocks/generate.py b/py/packages/genkit/src/genkit/_ai/_generate.py similarity index 61% rename from py/packages/genkit/src/genkit/blocks/generate.py rename to py/packages/genkit/src/genkit/_ai/_generate.py index 357923e0fa..b6094a9b25 100644 --- a/py/packages/genkit/src/genkit/blocks/generate.py +++ b/py/packages/genkit/src/genkit/_ai/_generate.py @@ -17,36 +17,33 @@ """Generate action.""" import copy +import inspect import re from collections.abc import Callable from typing import Any, cast -from genkit.blocks.formats.types import FormatDef, Formatter -from genkit.blocks.messages import inject_instructions -from genkit.blocks.middleware import augment_with_context -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - GenerateResponseWrapper, - MessageWrapper, +from pydantic import BaseModel + +from genkit._ai._formats._types import FormatDef, Formatter +from genkit._ai._messages import inject_instructions +from genkit._ai._middleware import augment_with_context +from genkit._ai._model import ( + Message, ModelMiddleware, + ModelRequest, + ModelResponse, + ModelResponseChunk, ) -from genkit.blocks.resource import ResourceArgument, ResourceInput, find_matching_resource, resolve_resources -from genkit.blocks.tools import ToolInterruptError -from genkit.codec import dump_dict -from genkit.core.action import Action, ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.error import GenkitError -from genkit.core.logging import get_logger -from genkit.core.registry import Registry -from genkit.core.typing import ( +from genkit._ai._resource import ResourceArgument, ResourceInput, find_matching_resource, resolve_resources +from genkit._ai._tools import ToolInterruptError +from genkit._core._action import Action, ActionKind, ActionRunContext +from genkit._core._error import GenkitError +from genkit._core._logger import get_logger +from genkit._core._registry import Registry +from genkit._core._typing import ( FinishReason, GenerateActionOptions, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - Message, Metadata, - OutputConfig, Part, Role, ToolDefinition, @@ -56,8 +53,6 @@ ToolResponsePart, ) -StreamingCallback = Callable[[GenerateResponseChunkWrapper], None] - DEFAULT_MAX_TURNS = 5 logger = get_logger(__name__) @@ -90,16 +85,20 @@ def _redact_data_uris(obj: Any) -> Any: # noqa: ANN401 def define_generate_action(registry: Registry) -> None: """Registers generate action in the provided registry.""" - async def generate_action_fn(input: GenerateActionOptions, ctx: ActionRunContext) -> GenerateResponse: + async def generate_action_fn( + input: GenerateActionOptions, + ctx: ActionRunContext, + ) -> ModelResponse: + on_chunk = cast(Callable[[ModelResponseChunk], None], ctx.streaming_callback) if ctx.is_streaming else None return await generate_action( registry=registry, raw_request=input, - on_chunk=ctx.send_chunk if ctx.is_streaming else None, + on_chunk=on_chunk, context=ctx.context, ) _ = registry.register_action( - kind=cast(ActionKind, ActionKind.UTIL), + kind=ActionKind.UTIL, name='generate', fn=generate_action_fn, ) @@ -108,26 +107,13 @@ async def generate_action_fn(input: GenerateActionOptions, ctx: ActionRunContext async def generate_action( registry: Registry, raw_request: GenerateActionOptions, - on_chunk: StreamingCallback | None = None, + on_chunk: Callable[[ModelResponseChunk], None] | None = None, message_index: int = 0, current_turn: int = 0, middleware: list[ModelMiddleware] | None = None, context: dict[str, Any] | None = None, -) -> GenerateResponseWrapper: - """Generate action. - - Args: - registry: The registry to use for the action. - raw_request: The raw request to generate. - on_chunk: The callback to use for the action. - message_index: The index of the message to use for the action. - current_turn: The current turn of the action. - middleware: The middleware to use for the action. - context: The context to use for the action. - - Returns: - The generated response. - """ +) -> ModelResponse: + """Execute a generation request with tool calling and middleware support.""" model, tools, format_def = await resolve_parameters(registry, raw_request) raw_request, formatter = apply_format(raw_request, format_def) @@ -156,27 +142,14 @@ async def generate_action( request = await action_to_generate_request(raw_request, tools, model) - logger.debug('generate request', model=model.name, request=_redact_data_uris(dump_dict(request))) + logger.debug('generate request', model=model.name, request=_redact_data_uris(request.model_dump())) - prev_chunks: list[GenerateResponseChunk] = [] + prev_chunks: list[ModelResponseChunk] = [] - chunk_role: Role = cast(Role, Role.MODEL) + chunk_role: Role = Role.MODEL - def make_chunk(role: Role, chunk: GenerateResponseChunk) -> GenerateResponseChunkWrapper: - """Creates a GenerateResponseChunkWrapper from role and data. - - This convenience method wraps a raw chunk, adds metadata like the - current message index and previous chunks, appends the raw chunk - to the internal `prev_chunks` list, and increments the message index - if the role changes. - - Args: - role: The role (e.g., MODEL, TOOL) associated with this chunk. - chunk: The raw GenerateResponseChunk data. - - Returns: - A GenerateResponseChunkWrapper containing the chunk and metadata. - """ + def make_chunk(role: Role, chunk: ModelResponseChunk) -> ModelResponseChunk: + """Wrap a raw chunk with metadata and track message index changes.""" nonlocal chunk_role, message_index if role != chunk_role and len(prev_chunks) > 0: @@ -187,35 +160,24 @@ def make_chunk(role: Role, chunk: GenerateResponseChunk) -> GenerateResponseChun prev_to_send = copy.copy(prev_chunks) prev_chunks.append(chunk) - def chunk_parser(chunk: GenerateResponseChunkWrapper) -> Any: # noqa: ANN401 - """Parse a chunk using the current formatter.""" + def chunk_parser(chunk: ModelResponseChunk) -> Any: # noqa: ANN401 if formatter is None: return None return formatter.parse_chunk(chunk) - return GenerateResponseChunkWrapper( + return ModelResponseChunk( chunk, index=message_index, previous_chunks=prev_to_send, chunk_parser=chunk_parser if formatter else None, ) - def wrap_chunks(role: Role | None = None) -> Callable[[GenerateResponseChunk], None]: - """Wrap and process a model response chunk. - - This function prepares model response chunks for the stream callback. - - Args: - role: The role to use for the wrapped chunks (default: MODEL). - chunk: The original model response chunk. - - Returns: - The result of passing the processed chunk to the callback. - """ + def wrap_chunks(role: Role | None = None) -> Callable[[ModelResponseChunk], None]: + """Return a callback that wraps chunks with the given role for streaming.""" if role is None: - role = cast(Role, Role.MODEL) + role = Role.MODEL - def wrapper(chunk: GenerateResponseChunk) -> None: + def wrapper(chunk: ModelResponseChunk) -> None: if on_chunk is not None: on_chunk(make_chunk(role, chunk)) @@ -237,46 +199,61 @@ def wrapper(chunk: GenerateResponseChunk) -> None: if raw_request.docs and not supports_context: middleware.append(augment_with_context()) - async def dispatch(index: int, req: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - """Dispatches model request, passing it through middleware if present. - - Args: - index: The index of the middleware to use. - req: The request to dispatch. - ctx: The context to use for the action. - - Returns: - The generated response. - """ + async def dispatch( + index: int, + req: ModelRequest, + ctx: ActionRunContext, + chunk_callback: Callable[[ModelResponseChunk], None] | None, + ) -> ModelResponse: + """Dispatch request through middleware chain to the model.""" if not middleware or index == len(middleware): - # end of the chain, call the original model action + # End of the chain, call the original model action return ( - await model.arun( + await model.run( input=req, context=ctx.context, - on_chunk=ctx.send_chunk if ctx.is_streaming else None, + on_chunk=cast(Callable[[object], None], chunk_callback) if chunk_callback else None, ) ).response current_middleware = middleware[index] + n_params = len(inspect.signature(current_middleware).parameters) + + if n_params == 4: + # Streaming middleware: (req, ctx, on_chunk, next) -> response + async def next_fn_streaming( + modified_req: ModelRequest | None = None, + modified_ctx: ActionRunContext | None = None, + modified_on_chunk: Callable[[ModelResponseChunk], None] | None = None, + ) -> ModelResponse: + return await dispatch( + index + 1, + modified_req if modified_req else req, + modified_ctx if modified_ctx else ctx, + modified_on_chunk if modified_on_chunk is not None else chunk_callback, + ) - async def next_fn( - modified_req: GenerateRequest | None = None, - modified_ctx: ActionRunContext | None = None, - ) -> GenerateResponse: - return await dispatch( - index + 1, - modified_req if modified_req else req, - modified_ctx if modified_ctx else ctx, - ) + return await current_middleware(req, ctx, chunk_callback, next_fn_streaming) + else: + # Simple middleware: (req, ctx, next) -> response + async def next_fn_simple( + modified_req: ModelRequest | None = None, + modified_ctx: ActionRunContext | None = None, + ) -> ModelResponse: + return await dispatch( + index + 1, + modified_req if modified_req else req, + modified_ctx if modified_ctx else ctx, + chunk_callback, + ) - return await current_middleware(req, ctx, next_fn) + return await current_middleware(req, ctx, next_fn_simple) # if resolving the 'resume' option above generated a tool message, stream it. if resumed_tool_message and on_chunk: - wrap_chunks(cast(Role, Role.TOOL))( - GenerateResponseChunk( - role=cast(Role, resumed_tool_message.role), + wrap_chunks(Role.TOOL)( + ModelResponseChunk( + role=resumed_tool_message.role, content=resumed_tool_message.content, ) ) @@ -284,21 +261,11 @@ async def next_fn( model_response = await dispatch( 0, request, - ActionRunContext( - on_chunk=cast(Callable[[object], None], wrap_chunks()) if on_chunk else None, - context=context, - ), + ActionRunContext(context=context), + wrap_chunks() if on_chunk else None, ) - def message_parser(msg: MessageWrapper) -> Any: # noqa: ANN401 - """Parse a message using the current formatter. - - Args: - msg: The message to parse. - - Returns: - The parsed message content. - """ + def message_parser(msg: Message) -> Any: # noqa: ANN401 if formatter is None: return None return formatter.parse_message(msg) @@ -306,14 +273,16 @@ def message_parser(msg: MessageWrapper) -> Any: # noqa: ANN401 # Extract schema_type for runtime Pydantic validation schema_type = raw_request.output.schema_type if raw_request.output else None - response = GenerateResponseWrapper( - model_response, - request, - message_parser=message_parser if formatter else None, - schema_type=schema_type, - ) + # Plugin returns ModelResponse directly. Framework sets request and + # any output format context (message_parser, schema_type) as private attrs. + response = model_response + response.request = request + if formatter: + response._message_parser = message_parser + if schema_type: + response._schema_type = schema_type - logger.debug('generate response', response=_redact_data_uris(dump_dict(response))) + logger.debug('generate response', response=_redact_data_uris(response.model_dump())) response.assert_valid() generated_msg = response.message @@ -343,29 +312,24 @@ def message_parser(msg: MessageWrapper) -> Any: # noqa: ANN401 revised_model_msg, tool_msg, transfer_preamble, - ) = await resolve_tool_requests(registry, raw_request, generated_msg._original_message) # pyright: ignore[reportPrivateUsage] + ) = await resolve_tool_requests(registry, raw_request, generated_msg) # if an interrupt message is returned, stop the tool loop and return a # response. if revised_model_msg: - interrupted_resp = GenerateResponseWrapper( - response, - request, - message_parser=message_parser if formatter else None, - schema_type=schema_type, - ) + interrupted_resp = response.model_copy(deep=False) interrupted_resp.finish_reason = FinishReason.INTERRUPTED interrupted_resp.finish_message = 'One or more tool calls resulted in interrupts.' - interrupted_resp.message = MessageWrapper(revised_model_msg) + interrupted_resp.message = Message(revised_model_msg) return interrupted_resp # If the loop will continue, stream out the tool response message... if on_chunk and tool_msg: on_chunk( make_chunk( - cast(Role, Role.TOOL), - GenerateResponseChunk( - role=cast(Role, tool_msg.role), + Role.TOOL, + ModelResponseChunk( + role=tool_msg.role, content=tool_msg.content, ), ) @@ -394,21 +358,7 @@ def message_parser(msg: MessageWrapper) -> Any: # noqa: ANN401 def apply_format( raw_request: GenerateActionOptions, format_def: FormatDef | None ) -> tuple[GenerateActionOptions, Formatter[Any, Any] | None]: - """Applies formatting instructions and configuration to the request. - - If a format definition is provided, this function deep copies the request, - resolves formatting instructions, applies them to the messages, and updates - output configuration based on the format definition and the request's - original output settings. - - Args: - raw_request: The original generation request options. - format_def: The format definition to apply, or None. - - Returns: - A tuple containing the potentially modified request options and the - resolved Formatter instance (or None if no format was applied). - """ + """Apply format definition to request, injecting instructions and output config.""" if not format_def: return raw_request, None @@ -416,10 +366,11 @@ def apply_format( formatter = format_def(raw_request.output.json_schema if raw_request.output else None) - instructions = resolve_instructions( - formatter, - raw_request.output.instructions if raw_request.output else None, - ) + # Extract instructions - handle bool | str | None type + # Schema allows: str (custom instructions), True (use defaults), False (disable), None (default behavior) + raw_instructions = raw_request.output.instructions if raw_request.output else None + str_instructions = raw_instructions if isinstance(raw_instructions, str) else None + instructions = resolve_instructions(formatter, str_instructions) should_inject = False if raw_request.output and raw_request.output.instructions is not None: @@ -430,7 +381,7 @@ def apply_format( should_inject = True if should_inject and instructions is not None: - out_request.messages = inject_instructions(out_request.messages, instructions) + out_request.messages = inject_instructions(out_request.messages, instructions) # type: ignore[arg-type] # Ensure output is set before modifying its properties if out_request.output is None: @@ -449,23 +400,11 @@ def apply_format( return (out_request, formatter) -def resolve_instructions(formatter: Formatter[Any, Any], instructions_opt: bool | str | None) -> str | None: - """Resolve instructions based on formatter and instruction options. - - Args: - formatter: The formatter to use for resolving instructions. - instructions_opt: The instruction options: True/False, a string, or - None. - - Returns: - The resolved instructions or None if no instructions should be used. - """ - if isinstance(instructions_opt, str): +def resolve_instructions(formatter: Formatter[Any, Any], instructions_opt: str | None) -> str | None: + """Return custom instructions if provided, otherwise use formatter defaults.""" + if instructions_opt is not None: # user provided instructions return instructions_opt - if instructions_opt is False: - # user says no instructions - return None if not formatter: return None # pyright: ignore[reportUnreachable] - defensive check return formatter.instructions @@ -474,36 +413,13 @@ def resolve_instructions(formatter: Formatter[Any, Any], instructions_opt: bool def apply_transfer_preamble( next_request: GenerateActionOptions, _preamble: GenerateActionOptions ) -> GenerateActionOptions: - """Applies relevant properties from a preamble request to the next request. - - This function is intended to copy settings (like model, config, etc.) - from a preamble, which might be generated by a tool during processing, - to the subsequent generation request. - - Note: This function is currently a placeholder (TODO). - - Args: - next_request: The generation request to be modified. - preamble: The preamble request containing properties to transfer. - - Returns: - The potentially updated `next_request`. - """ + """Transfer preamble settings to the next request. (TODO: not yet implemented).""" # TODO(#4338): implement me return next_request def _extract_resource_uri(resource_obj: Any) -> str | None: # noqa: ANN401 - """Extract URI from a resource object. - - Handles various Pydantic wrapper structures (Resource, Resource1, RootModel, dict). - - Args: - resource_obj: The resource object to extract URI from. - - Returns: - The extracted URI string, or None if not found. - """ + """Extract URI from a resource object, unwrapping Pydantic structures as needed.""" # Direct uri attribute (Resource1, ResourceInput, etc.) if hasattr(resource_obj, 'uri'): return resource_obj.uri @@ -524,15 +440,7 @@ def _extract_resource_uri(resource_obj: Any) -> str | None: # noqa: ANN401 async def apply_resources(registry: Registry, raw_request: GenerateActionOptions) -> GenerateActionOptions: - """Applies resources to the request messages by hydrating resource parts. - - Args: - registry: The registry to use for resolving resources. - raw_request: The generation request. - - Returns: - The updated generation request with hydrated resources. - """ + """Resolve and hydrate resource parts in the request messages.""" # Quick check if any message has a resource part has_resource = False for msg in raw_request.messages: @@ -593,7 +501,7 @@ async def apply_resources(registry: Registry, raw_request: GenerateActionOptions ) # Execute the resource - response = await resource_action.arun(resource_input, on_chunk=None, context=None) + response = await resource_action.run(resource_input, on_chunk=None, context=None) # response.response is ResourceOutput which has .content (list of Parts) # It usually returns a dict if coming from dynamic_resource (model_dump called) @@ -615,14 +523,7 @@ async def apply_resources(registry: Registry, raw_request: GenerateActionOptions def assert_valid_tool_names(_raw_request: GenerateActionOptions) -> None: - """Assert that tool names in the request are valid. - - Args: - raw_request: The generation request to validate. - - Raises: - ValueError: If any tool names are invalid. - """ + """Validate tool names in the request. (TODO: not yet implemented).""" # TODO(#4338): implement me pass @@ -630,16 +531,7 @@ def assert_valid_tool_names(_raw_request: GenerateActionOptions) -> None: async def resolve_parameters( registry: Registry, request: GenerateActionOptions ) -> tuple[Action[Any, Any, Any], list[Action[Any, Any, Any]], FormatDef | None]: - """Resolve parameters for the generate action. - - Args: - registry: The registry to use for the action. - request: The generation request to resolve parameters for. - - Returns: - A tuple containing the model action, the list of tool actions, and the - format definition. - """ + """Resolve model, tools, and format from registry for a generation request.""" model = request.model if request.model is not None else registry.default_model if not model: raise Exception('No model configured.') @@ -651,7 +543,7 @@ async def resolve_parameters( tools: list[Action[Any, Any, Any]] = [] if request.tools: for tool_name in request.tools: - tool_action = await registry.resolve_action(cast(ActionKind, ActionKind.TOOL), tool_name) + tool_action = await registry.resolve_action(ActionKind.TOOL, tool_name) if tool_action is None: raise Exception(f'Unable to resolve tool {tool_name}') tools.append(tool_action) @@ -668,50 +560,29 @@ async def resolve_parameters( async def action_to_generate_request( options: GenerateActionOptions, resolved_tools: list[Action[Any, Any, Any]], _model: Action[Any, Any, Any] -) -> GenerateRequest: - """Convert generate action options to a generate request. - - Args: - options: The generation options to convert. - resolved_tools: The resolved tools to use for the action. - model: The model to use for the action. - - Returns: - The generated request. - """ +) -> ModelRequest: + """Convert GenerateActionOptions to a ModelRequest with tool definitions.""" # TODO(#4340): add warning when tools are not supported in ModelInfo # TODO(#4341): add warning when toolChoice is not supported in ModelInfo tool_defs = [to_tool_definition(tool) for tool in resolved_tools] if resolved_tools else [] - return GenerateRequest( - messages=options.messages, - config=options.config if options.config is not None else {}, - docs=options.docs, + output = options.output + return ModelRequest( + # Field validators auto-wrap MessageData -> Message and DocumentData -> Document + messages=options.messages, # type: ignore[arg-type] + config=options.config if options.config is not None else {}, # type: ignore[arg-type] + docs=options.docs if options.docs else None, # type: ignore[arg-type] tools=tool_defs, tool_choice=options.tool_choice, - output=OutputConfig( - content_type=options.output.content_type if options.output else None, - format=options.output.format if options.output else None, - schema=options.output.json_schema if options.output else None, - constrained=options.output.constrained if options.output else None, - ), + output_format=output.format if output else None, + output_schema=output.json_schema if output else None, + output_constrained=output.constrained if output else None, + output_content_type=output.content_type if output else None, ) def to_tool_definition(tool: Action) -> ToolDefinition: - """Convert an action to a tool definition. - - Args: - tool: The action to convert. - - Returns: - The converted tool definition. - - Note: - Tool names may contain '/' characters (e.g., 'plugin/action'). - The full name is preserved here; model plugins are responsible for - escaping names if the provider API doesn't support certain characters. - """ + """Convert an Action to a ToolDefinition for model requests.""" tdef = ToolDefinition( name=tool.name, description=tool.description or '', @@ -724,17 +595,7 @@ def to_tool_definition(tool: Action) -> ToolDefinition: async def resolve_tool_requests( registry: Registry, request: GenerateActionOptions, message: Message ) -> tuple[Message | None, Message | None, GenerateActionOptions | None]: - """Resolve tool requests for the generate action. - - Args: - registry: The registry to use for the action. - request: The generation request to resolve tool requests for. - message: The message to resolve tool requests for. - - Returns: - A tuple containing the revised model message, the tool message, and the - transfer preamble. - """ + """Execute tool requests in a message, returning responses or interrupt info.""" # TODO(#4342): prompt transfer tool_dict: dict[str, Action] = {} if request.tools: @@ -779,20 +640,7 @@ async def resolve_tool_requests( def _to_pending_response(request: ToolRequestPart, response: ToolResponsePart) -> Part: - """Updates a ToolRequestPart to mark it as pending with its response data. - - This is used when a tool call is made, and its response is available, - but the overall generation loop might continue (e.g., for other tool calls). - The response output is stored in the request's metadata under 'pendingOutput'. - - Args: - request: The original ToolRequestPart. - response: The corresponding ToolResponsePart. - - Returns: - A new Part containing the original tool request but with updated metadata - indicating the pending output. - """ + """Mark a tool request as pending with its response stored in metadata.""" metadata = request.metadata.root if request.metadata else {} metadata['pendingOutput'] = response.tool_response.output # Part is a RootModel, so we pass content via 'root' parameter @@ -805,26 +653,9 @@ def _to_pending_response(request: ToolRequestPart, response: ToolResponsePart) - async def _resolve_tool_request(tool: Action, tool_request_part: ToolRequestPart) -> tuple[Part | None, Part | None]: - """Executes a tool action and returns its response or interrupt part. - - Calls the tool's `arun_raw` method with the input from the tool request. - Handles potential ToolInterruptErrors by returning a specific interrupt Part. - - Args: - tool: The resolved tool Action to execute. - tool_request_part: The Part containing the tool request details. - - Returns: - A tuple containing: - - A Part with the ToolResponse if successful, else None. - - A Part marking an interrupt if a ToolInterruptError occurred, else None. - - Raises: - GenkitError: If any error other than ToolInterruptError occurs during - tool execution. - """ + """Execute a tool and return (response_part, interrupt_part).""" try: - tool_response = (await tool.arun_raw(tool_request_part.tool_request.input)).response + tool_response = (await tool.run(tool_request_part.tool_request.input)).response # Part is a RootModel, so we pass content via 'root' parameter return ( Part( @@ -832,7 +663,7 @@ async def _resolve_tool_request(tool: Action, tool_request_part: ToolRequestPart tool_response=ToolResponse( name=tool_request_part.tool_request.name, ref=tool_request_part.tool_request.ref, - output=dump_dict(tool_response), + output=tool_response.model_dump() if isinstance(tool_response, BaseModel) else tool_response, ) ) ), @@ -865,19 +696,8 @@ async def _resolve_tool_request(tool: Action, tool_request_part: ToolRequestPart async def resolve_tool(registry: Registry, tool_name: str) -> Action: - """Resolve a tool by name from the registry. - - Args: - registry: The registry to resolve the tool from. - tool_name: The name of the tool to resolve. - - Returns: - The resolved tool action. - - Raises: - ValueError: If the tool could not be resolved. - """ - tool = await registry.resolve_action(kind=cast(ActionKind, ActionKind.TOOL), name=tool_name) + """Resolve a tool by name from the registry.""" + tool = await registry.resolve_action(kind=ActionKind.TOOL, name=tool_name) if tool is None: raise ValueError(f'Unable to resolve tool {tool_name}') return tool @@ -885,32 +705,8 @@ async def resolve_tool(registry: Registry, tool_name: str) -> Action: async def _resolve_resume_options( _registry: Registry, raw_request: GenerateActionOptions -) -> tuple[GenerateActionOptions, GenerateResponse | None, Message | None]: - """Resolves tool calls from a previous turn when resuming generation. - - Handles the `resume` option in GenerateActionOptions. It processes the - last model message (which should contain tool requests), resolves pending - outputs or applies provided responses/restarts from the `resume` argument, - and constructs a tool message to append to the history. - - Args: - registry: The action registry (unused in current implementation). - raw_request: The incoming generation request potentially with resume options. - - Returns: - A tuple containing: - - revised_request: The request updated with the new tool message in history - and the `resume` field cleared. - - interrupted_response: Currently always None (interrupts during resume - are not fully supported yet). - - tool_message: The constructed tool message with resolved responses, - or None if not resuming. - - Raises: - GenkitError: If the request format is invalid for resuming (e.g., last - message is not a model message with tool requests) or if - interrupted tool requests are not handled by the resume options. - """ +) -> tuple[GenerateActionOptions, ModelResponse | None, Message | None]: + """Handle resume options by resolving pending tool calls from a previous turn.""" if not raw_request.resume: return (raw_request, None, None) @@ -961,27 +757,7 @@ async def _resolve_resume_options( def _resolve_resumed_tool_request(raw_request: GenerateActionOptions, tool_request_part: Part) -> tuple[Part, Part]: - """Resolves a single tool request based on resume options. - - Checks if the tool request part has pending output in its metadata or if a - corresponding response is provided in the `raw_request.resume.respond` list. - Constructs the appropriate request and response parts for the tool history. - - Args: - raw_request: The overall generation request containing resume options. - tool_request_part: The specific tool request Part from the previous turn. - - Returns: - A tuple containing: - - request_part: The ToolRequest Part to be kept/added in the history. - Metadata might be updated (e.g., 'resolvedInterrupt'). - - response_part: The corresponding ToolResponse Part derived from pending - output or provided responses. - - Raises: - GenkitError: If the tool request cannot be resolved (neither pending - output nor a provided response is found). - """ + """Resolve a single tool request from pending output or resume.respond list.""" # Type narrowing: ensure we're working with a ToolRequestPart if not isinstance(tool_request_part.root, ToolRequestPart): raise GenkitError( @@ -1004,7 +780,7 @@ def _resolve_resumed_tool_request(raw_request: GenerateActionOptions, tool_reque tool_response=ToolResponse( name=tool_req_root.tool_request.name, ref=tool_req_root.tool_request.ref, - output=dump_dict(pending_output), + output=pending_output.model_dump() if isinstance(pending_output, BaseModel) else pending_output, ), metadata=Metadata(root=metadata), ) @@ -1046,15 +822,7 @@ def _resolve_resumed_tool_request(raw_request: GenerateActionOptions, tool_reque def _find_corresponding_tool_response(responses: list[ToolResponsePart], request: ToolRequestPart) -> Part | None: - """Finds a response part corresponding to a given request part by name and ref. - - Args: - responses: A list of ToolResponseParts to search within. - request: The ToolRequestPart to find a corresponding response for. - - Returns: - The matching ToolResponsePart as a Part, or None if no match is found. - """ + """Find a response matching the request by name and ref.""" for p in responses: if p.tool_response.name == request.tool_request.name and p.tool_response.ref == request.tool_request.ref: return Part(root=p) @@ -1064,25 +832,18 @@ def _find_corresponding_tool_response(responses: list[ToolResponsePart], request # TODO(#4336): extend GenkitError class GenerationResponseError(Exception): # TODO(#4337): use status enum - """Error for generation responses.""" + """Error raised when a generation request fails.""" def __init__( self, - response: GenerateResponse, + response: ModelResponse, message: str, status: str, details: dict[str, Any], ) -> None: - """Initialize the GenerationResponseError. - - Args: - response: The generation response. - message: The message to display. - status: The status of the generation response. - details: The details of the generation response. - """ + """Initialize with the failed response and error details.""" super().__init__(message) - self.response: GenerateResponse = response + self.response: ModelResponse = response self.message: str = message self.status: str = status self.details: dict[str, Any] = details diff --git a/py/packages/genkit/src/genkit/blocks/messages.py b/py/packages/genkit/src/genkit/_ai/_messages.py similarity index 61% rename from py/packages/genkit/src/genkit/blocks/messages.py rename to py/packages/genkit/src/genkit/_ai/_messages.py index ab58046b7e..8ffc4dec7e 100644 --- a/py/packages/genkit/src/genkit/blocks/messages.py +++ b/py/packages/genkit/src/genkit/_ai/_messages.py @@ -18,8 +18,8 @@ from typing import Any -from genkit.core.typing import ( - Message, +from genkit._ai._model import Message +from genkit._core._typing import ( Metadata, Part, Role, @@ -28,14 +28,7 @@ def _get_metadata_dict(metadata: Metadata | dict[str, Any] | None) -> dict[str, Any]: - """Safely extracts the dict from a Metadata RootModel or returns the dict directly. - - Args: - metadata: The metadata, which can be a Metadata RootModel, a dict, or None. - - Returns: - The underlying dict, or an empty dict if metadata is None. - """ + """Extract dict from Metadata RootModel or return dict directly.""" if metadata is None: return {} if isinstance(metadata, Metadata): @@ -44,16 +37,7 @@ def _get_metadata_dict(metadata: Metadata | dict[str, Any] | None) -> dict[str, def _is_output_part(part: Part, require_pending: bool = False, require_non_pending: bool = False) -> bool: - """Check if a part has output purpose metadata. - - Args: - part: The part to check. - require_pending: If True, only match pending output parts. - require_non_pending: If True, only match non-pending output parts. - - Returns: - True if the part matches the criteria. - """ + """Check if a part has purpose='output' metadata, optionally filtering by pending state.""" metadata_dict = _get_metadata_dict(part.root.metadata) if metadata_dict.get('purpose') != 'output': return False @@ -65,32 +49,7 @@ def _is_output_part(part: Part, require_pending: bool = False, require_non_pendi def inject_instructions(messages: list[Message], instructions: str) -> list[Message]: - """Injects instructions as a new Part into a list of messages. - - This function attempts to add the provided `instructions` string as a new - text Part with `metadata={'purpose': 'output'}` into the message history. - - Injection Logic: - - If `instructions` is empty, the original list is returned. - - If any message already contains a non-pending output Part, the original list is returned. - - Otherwise, it looks for a target message: - 1. The first message containing a Part marked as pending output. - 2. If none, the first system message. - 3. If none, the last user message. - - If a target message is found: - - If the target contains a pending output Part, it's replaced by the new instruction Part. - - Otherwise, the new instruction Part is appended to the target message's content. - - A *new* list containing the potentially modified message is returned. - - Args: - messages: A list of Message objects representing the conversation history. - instructions: The text instructions to inject. If empty, no injection occurs. - - Returns: - A new list of Message objects with the instructions injected into the - appropriate message, or a copy of the original list if no suitable place - for injection was found or if instructions were empty. - """ + """Inject output instructions into the message list (system, pending output, or last user).""" if not instructions: return messages diff --git a/py/packages/genkit/src/genkit/blocks/middleware.py b/py/packages/genkit/src/genkit/_ai/_middleware.py similarity index 55% rename from py/packages/genkit/src/genkit/blocks/middleware.py rename to py/packages/genkit/src/genkit/_ai/_middleware.py index 30aa8cf4cf..479e74d81a 100644 --- a/py/packages/genkit/src/genkit/blocks/middleware.py +++ b/py/packages/genkit/src/genkit/_ai/_middleware.py @@ -16,17 +16,18 @@ """Middleware for the Genkit framework.""" -from genkit.blocks.model import ( +from collections.abc import Awaitable, Callable + +from genkit._ai._model import ( + Message, ModelMiddleware, - ModelMiddlewareNext, + ModelRequest, + ModelResponse, text_from_content, ) -from genkit.core.action import ActionRunContext -from genkit.core.typing import ( - DocumentData, - GenerateRequest, - GenerateResponse, - Message, +from genkit._core._action import ActionRunContext +from genkit._core._model import Document +from genkit._core._typing import ( Metadata, Part, TextPart, @@ -35,20 +36,8 @@ CONTEXT_PREFACE = '\n\nUse the following information to complete your task:\n\n' -def context_item_template(d: DocumentData, index: int) -> str: - """Renders a DocumentData object into a formatted string for context injection. - - Creates a string representation of the document, typically for inclusion in a - prompt. It attempts to use metadata fields ('ref', 'id') or the provided index - as a citation marker. - - Args: - d: The DocumentData object to render. - index: The index of the document in a list, used as a fallback citation. - - Returns: - A formatted string representing the document content with a citation. - """ +def context_item_template(d: Document, index: int) -> str: + """Render a document as a citation line for context injection.""" out = '- ' ref = (d.metadata and (d.metadata.get('ref') or d.metadata.get('id'))) or index out += f'[{ref}]: ' @@ -57,36 +46,13 @@ def context_item_template(d: DocumentData, index: int) -> str: def augment_with_context() -> ModelMiddleware: - """Returns a ModelMiddleware that augments the prompt with document context. - - This middleware checks if the `GenerateRequest` includes documents (`req.docs`). - If documents are present, it finds the last user message and injects the - rendered content of the documents into it as a special context Part. - - Returns: - A ModelMiddleware function. - """ + """Middleware that injects document context into the last user message.""" async def middleware( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next_middleware: ModelMiddlewareNext, - ) -> GenerateResponse: - """The actual middleware logic to inject context. - - Checks for documents in the request. If found, locates the last user message. - It then either replaces an existing pending context Part or appends a new - context Part (rendered from the documents) to the user message before - passing the request to the next middleware or the model. - - Args: - req: The incoming GenerateRequest. - ctx: The ActionRunContext. - next_middleware: The next function in the middleware chain. - - Returns: - The result from the next middleware or the final GenerateResponse. - """ + next_middleware: Callable[..., Awaitable[ModelResponse]], + ) -> ModelResponse: if not req.docs: return await next_middleware(req, ctx) @@ -111,7 +77,7 @@ async def middleware( out = CONTEXT_PREFACE for i, doc_data in enumerate(req.docs): - doc = DocumentData(content=doc_data.content, metadata=doc_data.metadata) + doc = Document(content=doc_data.content, metadata=doc_data.metadata) out += context_item_template(doc, i) out += '\n' @@ -129,14 +95,7 @@ async def middleware( def last_user_message(messages: list[Message]) -> Message | None: - """Finds the last message with the role 'user' in a list of messages. - - Args: - messages: A list of message dictionaries. - - Returns: - The last message with the role 'user', or None if no such message exists. - """ + """Find the last user message in a list.""" for i in range(len(messages) - 1, -1, -1): if messages[i].role == 'user': return messages[i] diff --git a/py/packages/genkit/src/genkit/_ai/_model.py b/py/packages/genkit/src/genkit/_ai/_model.py new file mode 100644 index 0000000000..fba260334c --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_model.py @@ -0,0 +1,165 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Model type definitions for the Genkit framework.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from typing import Any, cast + +from pydantic import BaseModel + +from genkit._core._action import ( + Action, + ActionKind, + ActionMetadata, + ActionRunContext, + get_func_description, +) +from genkit._core._model import ( + Message, + ModelConfig, + ModelMiddleware, + ModelRef, + ModelRequest, + ModelResponse, + ModelResponseChunk, + get_basic_usage_stats, + text_from_content, + text_from_message, +) +from genkit._core._registry import Registry +from genkit._core._schema import to_json_schema +from genkit._core._typing import ModelInfo + +# Type alias for model functions (must be async) +# Use ctx.send_chunk() for streaming +ModelFn = Callable[[ModelRequest, ActionRunContext], Awaitable[ModelResponse[Any]]] + + +def model_action_metadata( + name: str, + info: dict[str, object] | None = None, + config_schema: type | dict[str, Any] | None = None, +) -> ActionMetadata: + """Create ActionMetadata for a model action.""" + info = info if info is not None else {} + return ActionMetadata( + kind=ActionKind.MODEL, + name=name, + input_json_schema=to_json_schema(ModelRequest), + output_json_schema=to_json_schema(ModelResponse), + metadata={'model': {**info, 'customOptions': to_json_schema(config_schema) if config_schema else None}}, + ) + + +def model_ref( + name: str, + namespace: str | None = None, + info: ModelInfo | None = None, + version: str | None = None, + config: dict[str, object] | None = None, +) -> ModelRef: + """Create a ModelRef, optionally prefixing name with namespace.""" + # Logic: if (options.namespace && !name.startsWith(options.namespace + '/')) + final_name = f'{namespace}/{name}' if namespace and not name.startswith(f'{namespace}/') else name + + return ModelRef(name=final_name, info=info, version=version, config=config) + + +def define_model( + registry: Registry, + name: str, + fn: ModelFn, + config_schema: type[BaseModel] | dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + info: ModelInfo | None = None, + description: str | None = None, +) -> Action: + """Register a custom model action.""" + # Build model options dict + model_options: dict[str, object] = {} + + # Start with info if provided + if info: + model_options.update(info.model_dump()) + + # Check if metadata has model info + if metadata and 'model' in metadata: + existing = metadata['model'] + if isinstance(existing, dict): + existing_dict = cast(dict[str, object], existing) + for key, value in existing_dict.items(): + if isinstance(key, str) and key not in model_options: + model_options[key] = value + + # Default label to name if not set + if 'label' not in model_options or not model_options['label']: + model_options['label'] = name + + # Add config schema if provided + if config_schema: + model_options['customOptions'] = to_json_schema(config_schema) + + # Build the final metadata dict + model_meta: dict[str, object] = metadata.copy() if metadata else {} + model_meta['model'] = model_options + + model_description = get_func_description(fn, description) + return registry.register_action( + name=name, + kind=ActionKind.MODEL, + fn=fn, + metadata=model_meta, + description=model_description, + ) + + +# ============================================================================= +# Model config types (from model_types.py) +# ============================================================================= + + +def get_request_api_key(config: Mapping[str, object] | ModelConfig | object | None) -> str | None: + """Extract API key from config (snake_case or camelCase).""" + if config is None: + return None + + if isinstance(config, ModelConfig): + return config.api_key + + if isinstance(config, Mapping): + config_mapping = cast(Mapping[str, object], config) + api_key = config_mapping.get('api_key') + if isinstance(api_key, str) and api_key: + return api_key + else: + # Defensive fallback for plugin-specific config classes that inherit from + # ModelConfig or expose an api_key attribute. + api_key_attr = getattr(config, 'api_key', None) + if isinstance(api_key_attr, str) and api_key_attr: + return api_key_attr + + return None + + +def get_effective_api_key( + config: Mapping[str, object] | ModelConfig | object | None, + plugin_api_key: str | None, +) -> str | None: + """Return request API key if set, otherwise plugin API key.""" + return get_request_api_key(config) or plugin_api_key diff --git a/py/packages/genkit/src/genkit/_ai/_prompt.py b/py/packages/genkit/src/genkit/_ai/_prompt.py new file mode 100644 index 0000000000..2df5c3de55 --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_prompt.py @@ -0,0 +1,1282 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + + +"""Prompt management and templating.""" + +import asyncio +import os +import weakref +from collections.abc import AsyncIterable, Awaitable, Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any, ClassVar, Generic, TypedDict, TypeVar, cast + +from dotpromptz.typing import ( + DataArgument, + PromptFunction, + PromptInputConfig, + PromptMetadata, +) +from pydantic import BaseModel, ConfigDict + +from genkit._ai._generate import ( + generate_action, + to_tool_definition, +) +from genkit._ai._model import ( + Message, + ModelMiddleware, + ModelRequest, + ModelResponse, + ModelResponseChunk, +) +from genkit._core._action import Action, ActionKind, ActionRunContext, StreamingCallback, create_action_key +from genkit._core._channel import Channel +from genkit._core._error import GenkitError +from genkit._core._logger import get_logger +from genkit._core._model import Document, ModelConfig +from genkit._core._registry import Registry +from genkit._core._schema import to_json_schema +from genkit._core._typing import ( + GenerateActionOptions, + GenerateActionOutputConfig, + OutputConfig, + Part, + Resume, + Role, + TextPart, + ToolChoice, + ToolRequestPart, + ToolResponsePart, +) + +ModelStreamingCallback = StreamingCallback + +logger = get_logger(__name__) + +# TypeVars for generic input/output typing +InputT = TypeVar('InputT') +OutputT = TypeVar('OutputT') + + +class OutputOptions(TypedDict, total=False): + """Output format/schema configuration for prompt generation.""" + + format: str | None + content_type: str | None + instructions: str | None + schema: type | dict[str, Any] | str | None + json_schema: dict[str, Any] | None + constrained: bool | None + + +class ResumeOptions(TypedDict, total=False): + """Options for resuming generation after a tool interrupt.""" + + respond: ToolResponsePart | list[ToolResponsePart] | None + restart: ToolRequestPart | list[ToolRequestPart] | None + metadata: dict[str, Any] | None + + +class PromptGenerateOptions(TypedDict, total=False): + """Runtime options for prompt execution (config, tools, messages, etc.).""" + + model: str | None + config: dict[str, Any] | ModelConfig | None + messages: list[Message] | None + docs: list[Document] | None + tools: list[str] | None + resources: list[str] | None + tool_choice: ToolChoice | None + output: OutputOptions | None + resume: ResumeOptions | None + return_tool_requests: bool | None + max_turns: int | None + on_chunk: ModelStreamingCallback | None + use: list[ModelMiddleware] | None + context: dict[str, Any] | None + step_name: str | None + metadata: dict[str, Any] | None + + +class ModelStreamResponse(Generic[OutputT]): + """Response from streaming prompt execution with stream and response properties.""" + + def __init__( + self, + channel: Channel[ModelResponseChunk, ModelResponse[OutputT]], + response_future: asyncio.Future[ModelResponse[OutputT]], + ) -> None: + """Initialize with streaming channel and response future.""" + self._channel: Channel[ModelResponseChunk, ModelResponse[OutputT]] = channel + self._response_future: asyncio.Future[ModelResponse[OutputT]] = response_future + + @property + def stream(self) -> AsyncIterable[ModelResponseChunk]: + """Get the async iterable of response chunks. + + Returns: + An async iterable that yields ModelResponseChunk objects + as they are received from the model. Each chunk contains: + - text: The partial text generated so far + - index: The chunk index + - Additional metadata from the model + """ + return self._channel + + @property + def response(self) -> Awaitable[ModelResponse[OutputT]]: + """Get the awaitable for the complete response. + + Returns: + An awaitable that resolves to a ModelResponse containing: + - text: The complete generated text + - output: The typed output (when using Output[T]) + - messages: The full message history + - usage: Token usage statistics + - finish_reason: Why generation stopped (e.g., 'stop', 'length') + - Any tool calls or interrupts from the response + """ + return self._response_future + + +@dataclass +class PromptCache: + """Model for a prompt cache.""" + + user_prompt: PromptFunction[Any] | None = None + system: PromptFunction[Any] | None = None + messages: PromptFunction[Any] | None = None + + +class PromptConfig(BaseModel): + """Model for a prompt action.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) + + variant: str | None = None + model: str | None = None + config: dict[str, Any] | ModelConfig | None = None + description: str | None = None + input_schema: type | dict[str, Any] | str | None = None + system: str | list[Part] | None = None + prompt: str | list[Part] | None = None + messages: str | list[Message] | None = None + output_format: str | None = None + output_content_type: str | None = None + output_instructions: str | None = None + output_schema: type | dict[str, Any] | str | None = None + output_constrained: bool | None = None + max_turns: int | None = None + return_tool_requests: bool | None = None + metadata: dict[str, Any] | None = None + tools: list[str] | None = None + tool_choice: ToolChoice | None = None + use: list[ModelMiddleware] | None = None + docs: list[Document] | None = None + tool_responses: list[Part] | None = None + resources: list[str] | None = None + + +class ExecutablePrompt(Generic[InputT, OutputT]): + """A callable prompt with typed input/output that generates AI responses.""" + + def __init__( + self, + registry: Registry, + variant: str | None = None, + model: str | None = None, + config: dict[str, Any] | ModelConfig | None = None, + description: str | None = None, + input_schema: type | dict[str, Any] | str | None = None, + system: str | list[Part] | None = None, + prompt: str | list[Part] | None = None, + messages: str | list[Message] | None = None, + output_format: str | None = None, + output_content_type: str | None = None, + output_instructions: str | None = None, + output_schema: type | dict[str, Any] | str | None = None, + output_constrained: bool | None = None, + max_turns: int | None = None, + return_tool_requests: bool | None = None, + metadata: dict[str, Any] | None = None, + tools: list[str] | None = None, + tool_choice: ToolChoice | None = None, + use: list[ModelMiddleware] | None = None, + docs: list[Document] | None = None, + resources: list[str] | None = None, + name: str | None = None, + ns: str | None = None, + ) -> None: + """Initialize prompt with configuration, templates, and schema options.""" + self._registry = registry + self._variant = variant + self._model = model + self._config = config + self._description = description + self._input_schema = input_schema + self._system = system + self._prompt = prompt + self._messages = messages + self._output_format = output_format + self._output_content_type = output_content_type + self._output_instructions = output_instructions + self._output_schema = output_schema + self._output_constrained = output_constrained + self._max_turns = max_turns + self._return_tool_requests = return_tool_requests + self._metadata = metadata + self._tools = tools + self._tool_choice = tool_choice + self._use = use + self._docs = docs + self._resources = resources + self._cache_prompt: PromptCache = PromptCache() + self._name = name + self._ns = ns + self._prompt_action: Action | None = None + + @property + def ref(self) -> dict[str, Any]: + """Reference object with prompt name and metadata.""" + return { + 'name': registry_definition_key(self._name, self._variant, self._ns) if self._name else None, + 'metadata': self._metadata, + } + + async def _ensure_resolved(self) -> None: + if self._prompt_action or not self._name: + return + + # Preserve Pydantic schema type if it was explicitly provided via ai.prompt(..., output=Output(schema=T)) + # The resolved prompt from .prompt file will have a dict schema, but we want to keep the Pydantic type + # for runtime validation to get proper typed output. + original_output_schema = self._output_schema + + resolved = await lookup_prompt(self._registry, self._name, self._variant) + self._model = resolved._model + self._config = resolved._config + self._description = resolved._description + self._input_schema = resolved._input_schema + self._system = resolved._system + self._prompt = resolved._prompt + self._messages = resolved._messages + self._output_format = resolved._output_format + self._output_content_type = resolved._output_content_type + self._output_instructions = resolved._output_instructions + # Keep original Pydantic type if provided, otherwise use resolved (dict) schema + if isinstance(original_output_schema, type) and issubclass(original_output_schema, BaseModel): + self._output_schema = original_output_schema + else: + self._output_schema = resolved._output_schema + self._output_constrained = resolved._output_constrained + self._max_turns = resolved._max_turns + self._return_tool_requests = resolved._return_tool_requests + self._metadata = resolved._metadata + self._tools = resolved._tools + self._tool_choice = resolved._tool_choice + self._use = resolved._use + self._docs = resolved._docs + self._resources = resolved._resources + self._prompt_action = resolved._prompt_action + + async def __call__( + self, + input: InputT | None = None, + opts: PromptGenerateOptions | None = None, + ) -> ModelResponse[OutputT]: + """Execute the prompt and return the response.""" + await self._ensure_resolved() + effective_opts: PromptGenerateOptions = opts if opts else {} + + # Extract streaming callback and middleware from opts + on_chunk = effective_opts.get('on_chunk') + middleware = effective_opts.get('use') or self._use + context = effective_opts.get('context') + + result = await generate_action( + self._registry, + await self.render(input=input, opts=effective_opts), + on_chunk=on_chunk, + middleware=middleware, + context=context if context else ActionRunContext._current_context(), # pyright: ignore[reportPrivateUsage] + ) + # Cast to preserve the generic type parameter + return cast(ModelResponse[OutputT], result) + + def stream( + self, + input: InputT | None = None, + opts: PromptGenerateOptions | None = None, + *, + timeout: float | None = None, + ) -> ModelStreamResponse[OutputT]: + """Stream the prompt execution, returning (stream, response_future).""" + effective_opts: PromptGenerateOptions = opts if opts else {} + channel: Channel[ModelResponseChunk, ModelResponse[OutputT]] = Channel(timeout=timeout) + + # Create a copy of opts with the streaming callback + stream_opts: PromptGenerateOptions = { + **effective_opts, + 'on_chunk': lambda c: channel.send(cast(ModelResponseChunk, c)), + } + + resp = self.__call__(input=input, opts=stream_opts) + response_future: asyncio.Future[ModelResponse[OutputT]] = asyncio.create_task(resp) + channel.set_close_future(response_future) + + return ModelStreamResponse[OutputT](channel=channel, response_future=response_future) + + async def render( + self, + input: InputT | dict[str, Any] | None = None, + opts: PromptGenerateOptions | None = None, + ) -> GenerateActionOptions: + """Render the prompt template without executing, returning GenerateActionOptions.""" + await self._ensure_resolved() + if opts is None: + opts = cast(PromptGenerateOptions, {}) + output_opts = opts.get('output') or {} + context = opts.get('context') + + # Config merge requires special handling (dict merge with Pydantic conversion) + merged_config: dict[str, Any] | ModelConfig | None + if opts.get('config') is not None: + base = ( + self._config.model_dump(exclude_none=True) + if isinstance(self._config, BaseModel) + else (self._config or {}) + ) + opt_config = opts.get('config') + override = ( + opt_config.model_dump(exclude_none=True) if isinstance(opt_config, BaseModel) else (opt_config or {}) + ) + merged_config = {**base, **override} if base or override else None + else: + merged_config = self._config + + # Metadata merge (combine dicts) + merged_metadata = ( + {**(self._metadata or {}), **(opts.get('metadata') or {})} if opts.get('metadata') else self._metadata + ) + + def _or(opt_val: Any, default: Any) -> Any: # noqa: ANN401 + return opt_val if opt_val is not None else default + + prompt_options = PromptConfig( + model=opts.get('model') or self._model, + prompt=self._prompt, + system=self._system, + messages=self._messages, + tools=opts.get('tools') or self._tools, + return_tool_requests=_or(opts.get('return_tool_requests'), self._return_tool_requests), + tool_choice=opts.get('tool_choice') or self._tool_choice, + config=merged_config, + max_turns=_or(opts.get('max_turns'), self._max_turns), + output_format=output_opts.get('format') or self._output_format, + output_content_type=output_opts.get('content_type') or self._output_content_type, + output_instructions=_or(output_opts.get('instructions'), self._output_instructions), + output_schema=output_opts.get('schema') or output_opts.get('json_schema') or self._output_schema, + output_constrained=_or(output_opts.get('constrained'), self._output_constrained), + input_schema=self._input_schema, + metadata=merged_metadata, + docs=self._docs, + resources=opts.get('resources') or self._resources, + ) + + model = prompt_options.model or self._registry.default_model + if model is None: + raise GenkitError(status='INVALID_ARGUMENT', message='No model configured.') + + resolved_msgs: list[Message] = [] + # Convert input to dict for render functions + # If input is a Pydantic model, convert to dict; otherwise use as-is + render_input: dict[str, Any] + if input is None: + render_input = {} + elif isinstance(input, dict): + # Type narrow: input is dict here, assign to dict[str, Any] typed variable + render_input = {str(k): v for k, v in input.items()} + elif isinstance(input, BaseModel): + # Pydantic v2 model + render_input = input.model_dump() + elif hasattr(input, 'dict'): + # Pydantic v1 model + dict_func = getattr(input, 'dict', None) + render_input = cast(Callable[[], dict[str, Any]], dict_func)() + else: + # Fallback: cast to dict (should not happen with proper typing) + render_input = cast(dict[str, Any], input) + # Get opts.messages for history (matching JS behavior) + opts_messages = opts.get('messages') + + # Render system prompt + if prompt_options.system: + result = await render_system_prompt( + self._registry, render_input, prompt_options, self._cache_prompt, context + ) + resolved_msgs.append(result) + + # Handle messages (matching JS behavior): + # - If prompt has messages template: render it (opts.messages passed as history to resolvers) + # - If prompt has no messages: use opts.messages directly + if prompt_options.messages: + # Prompt defines messages - render them (resolvers receive opts_messages as history) + resolved_msgs.extend( + await render_message_prompt( + self._registry, + render_input, + prompt_options, + self._cache_prompt, + context, + history=opts_messages, + ) + ) + elif opts_messages: + # Prompt has no messages template - use opts.messages directly + resolved_msgs.extend(opts_messages) + + # Render user prompt + if prompt_options.prompt: + result = await render_user_prompt(self._registry, render_input, prompt_options, self._cache_prompt, context) + resolved_msgs.append(result) + + # If schema is set but format is not explicitly set, default to 'json' format + if prompt_options.output_schema and not prompt_options.output_format: + output_format = 'json' + else: + output_format = prompt_options.output_format + + # Build output config + output = GenerateActionOutputConfig() + if output_format: + output.format = output_format + if prompt_options.output_content_type: + output.content_type = prompt_options.output_content_type + if prompt_options.output_instructions is not None: + output.instructions = prompt_options.output_instructions + _resolve_output_schema(self._registry, prompt_options.output_schema, output) + if prompt_options.output_constrained is not None: + output.constrained = prompt_options.output_constrained + + # Handle resume options + resume = None + resume_opts = opts.get('resume') + if resume_opts: + respond = resume_opts.get('respond') + if respond: + resume = Resume(respond=respond) if isinstance(respond, list) else Resume(respond=[respond]) + + # Merge docs: opts.docs extends prompt docs + merged_docs = await render_docs(render_input, prompt_options, context) + opts_docs = opts.get('docs') + if opts_docs: + merged_docs = [*merged_docs, *opts_docs] if merged_docs else list(opts_docs) + + return GenerateActionOptions( + model=model, + messages=resolved_msgs, # type: ignore[arg-type] + config=prompt_options.config, + tools=prompt_options.tools, + return_tool_requests=prompt_options.return_tool_requests, + tool_choice=prompt_options.tool_choice, + output=output, + max_turns=prompt_options.max_turns, + docs=merged_docs, # type: ignore[arg-type] + resume=resume, + ) + + async def as_tool(self) -> Action: + """Expose this prompt as a tool. + + Returns the PROMPT action, which can be used as a tool. + """ + await self._ensure_resolved() + # If we have a direct reference to the action, use it + if self._prompt_action is not None: + return self._prompt_action + + # Otherwise, try to look it up using name/variant/ns + if self._name is None: + raise GenkitError( + status='FAILED_PRECONDITION', + message=( + 'Prompt name not available. This prompt was not created via define_prompt_async() or load_prompt().' + ), + ) + + lookup_key = registry_lookup_key(self._name, self._variant, self._ns) + + action = await self._registry.resolve_action_by_key(lookup_key) + + if action is None or action.kind != ActionKind.PROMPT: + raise GenkitError( + status='NOT_FOUND', + message=f'PROMPT action not found for prompt "{self._name}"', + ) + + return action + + +def register_prompt_actions( + registry: Registry, + executable_prompt: ExecutablePrompt[Any, Any], + name: str, + variant: str | None = None, +) -> None: + """Register PROMPT and EXECUTABLE_PROMPT actions for a prompt. + + This links the executable prompt to actions in the registry, enabling + lookup and DevUI integration. + """ + action_metadata: dict[str, object] = { + 'type': 'prompt', + 'source': 'programmatic', + 'prompt': { + 'name': name, + 'variant': variant or '', + }, + } + + async def prompt_action_fn(input: Any = None) -> ModelRequest: # noqa: ANN401 + options = await executable_prompt.render(input=input) + return await to_generate_request(registry, options) + + async def executable_prompt_action_fn(input: Any = None) -> GenerateActionOptions: # noqa: ANN401 + return await executable_prompt.render(input=input) + + action_name = registry_definition_key(name, variant) + prompt_action = registry.register_action( + kind=ActionKind.PROMPT, + name=action_name, + fn=prompt_action_fn, + metadata=action_metadata, + ) + + executable_prompt_action = registry.register_action( + kind=ActionKind.EXECUTABLE_PROMPT, + name=action_name, + fn=executable_prompt_action_fn, + metadata=action_metadata, + ) + + # Link them + executable_prompt._prompt_action = prompt_action # pyright: ignore[reportPrivateUsage] + setattr(prompt_action, '_executable_prompt', weakref.ref(executable_prompt)) # noqa: B010 + setattr(executable_prompt_action, '_executable_prompt', weakref.ref(executable_prompt)) # noqa: B010 + + +def _resolve_output_schema( + registry: Registry, + output_schema: type | dict[str, Any] | str | None, + output: GenerateActionOutputConfig, +) -> None: + """Resolve output schema and populate the output config. + + Handles three types of output_schema: + - str: Schema name - look up JSON schema and type from registry + - Pydantic type: Store both JSON schema and type for runtime validation + - dict: Raw JSON schema - convert directly + + Args: + registry: The registry to use for schema lookups. + output_schema: The schema to resolve (string name, Pydantic type, or dict). + output: The output config to populate with json_schema and schema_type. + """ + if output_schema is None: + return + + if isinstance(output_schema, str): + # Schema name - look up from registry + resolved_schema = registry.lookup_schema(output_schema) + if resolved_schema: + output.json_schema = resolved_schema + # Also look up the schema type for runtime validation + schema_type = registry.lookup_schema_type(output_schema) + if schema_type: + output.schema_type = schema_type + elif isinstance(output_schema, type) and issubclass(output_schema, BaseModel): + # Pydantic type - store both JSON schema and type + output.json_schema = to_json_schema(output_schema) + output.schema_type = output_schema + else: + # dict (raw JSON schema) + output.json_schema = to_json_schema(output_schema) + + +async def to_generate_action_options(registry: Registry, options: PromptConfig) -> GenerateActionOptions: + """Convert PromptConfig to GenerateActionOptions.""" + model = options.model or registry.default_model + if model is None: + raise GenkitError(status='INVALID_ARGUMENT', message='No model configured.') + + cache = PromptCache() + resolved_msgs: list[Message] = [] + # Use empty dict instead of None for render functions + render_input: dict[str, Any] = {} + if options.system: + result = await render_system_prompt(registry, render_input, options, cache) + resolved_msgs.append(result) + if options.messages: + resolved_msgs.extend(await render_message_prompt(registry, render_input, options, cache)) + if options.prompt: + result = await render_user_prompt(registry, render_input, options, cache) + resolved_msgs.append(result) + + # If is schema is set but format is not explicitly set, default to + # `json` format. + output_format = 'json' if options.output_schema and not options.output_format else options.output_format + + output = GenerateActionOutputConfig() + if output_format: + output.format = output_format + if options.output_content_type: + output.content_type = options.output_content_type + if options.output_instructions is not None: + output.instructions = options.output_instructions + _resolve_output_schema(registry, options.output_schema, output) + if options.output_constrained is not None: + output.constrained = options.output_constrained + + resume = None + if options.tool_responses: + # Filter for only ToolResponsePart instances + tool_response_parts = [r.root for r in options.tool_responses if isinstance(r.root, ToolResponsePart)] + if tool_response_parts: + resume = Resume(respond=tool_response_parts) + + return GenerateActionOptions( + model=model, + messages=resolved_msgs, # type: ignore[arg-type] + config=options.config, + tools=options.tools, + return_tool_requests=options.return_tool_requests, + tool_choice=options.tool_choice, + output=output, + max_turns=options.max_turns, + docs=await render_docs(render_input, options), # type: ignore[arg-type] + resume=resume, + ) + + +async def to_generate_request(registry: Registry, options: GenerateActionOptions) -> ModelRequest: + """Convert GenerateActionOptions to ModelRequest, resolving tool names.""" + tools: list[Action] = [] + if options.tools: + for tool_name in options.tools: + tool_action = await registry.resolve_action(ActionKind.TOOL, tool_name) + if tool_action is None: + raise GenkitError(status='NOT_FOUND', message=f'Unable to resolve tool {tool_name}') + tools.append(tool_action) + + tool_defs = [to_tool_definition(tool) for tool in tools] if tools else [] + + if not options.messages: + raise GenkitError( + status='INVALID_ARGUMENT', + message='at least one message is required in generate request', + ) + + output_config = OutputConfig( + content_type=options.output.content_type if options.output else None, + format=options.output.format if options.output else None, + schema=options.output.json_schema if options.output else None, + constrained=options.output.constrained if options.output else None, + ) + return ModelRequest( + # Field validators auto-wrap MessageData -> Message and DocumentData -> Document + messages=options.messages, # type: ignore[arg-type] + config=options.config if options.config is not None else {}, # type: ignore[arg-type] + docs=options.docs if options.docs else None, # type: ignore[arg-type] + tools=tool_defs, + tool_choice=options.tool_choice, + output_format=output_config.format, + output_schema=output_config.schema, + output_constrained=output_config.constrained, + output_content_type=output_config.content_type, + ) + + +def _normalize_prompt_arg( + prompt: str | list[Part] | None, +) -> list[Part]: + """Convert string/Part/list to list[Part].""" + if not prompt: + return [] + if isinstance(prompt, str): + # Part is a RootModel, so we pass content via 'root' parameter + return [Part(root=TextPart(text=prompt))] + elif isinstance(prompt, list): + return prompt + elif isinstance(prompt, Part): # pyright: ignore[reportUnnecessaryIsInstance] + return [prompt] + else: + return [] # pyright: ignore[reportUnreachable] - defensive fallback + + +async def _render_template( + registry: Registry, + role: Role, + template: str | list[Part] | None, + input: dict[str, Any], + input_schema: type | dict[str, Any] | str | None, + metadata: dict[str, Any] | None, + compiled_fn: PromptFunction[Any] | None, + context: dict[str, Any] | None, +) -> tuple[Message, PromptFunction[Any] | None]: + """Compile and render a prompt template, returning (message, compiled_fn).""" + if isinstance(template, str): + if compiled_fn is None: + compiled_fn = await registry.dotprompt.compile(template) + + if metadata: + context = {**(context or {}), 'state': metadata.get('state')} + + rendered_parts = cast( + list[Part], + await render_dotprompt_to_parts( + context or {}, + compiled_fn, + input, + PromptMetadata( + input=PromptInputConfig( + schema=to_json_schema(input_schema) if input_schema else None, + ) + ), + ), + ) + return Message(role=role, content=rendered_parts), compiled_fn + + return Message(role=role, content=_normalize_prompt_arg(template)), compiled_fn + + +async def render_system_prompt( + registry: Registry, + input: dict[str, Any], + options: PromptConfig, + prompt_cache: PromptCache, + context: dict[str, Any] | None = None, +) -> Message: + """Render the system prompt.""" + msg, prompt_cache.system = await _render_template( + registry, + Role.SYSTEM, + options.system, + input, + options.input_schema, + options.metadata, + prompt_cache.system, + context, + ) + return msg + + +async def render_dotprompt_to_parts( + context: dict[str, Any], + prompt_function: PromptFunction[Any], + input_: dict[str, Any], + options: PromptMetadata[Any] | None = None, +) -> list[dict[str, Any]]: + """Execute a compiled dotprompt function and return parts as dicts.""" + # Flatten input and context for template resolution + flattened_data = {**(context or {}), **(input_ or {})} + rendered = await prompt_function( + data=DataArgument[dict[str, Any]]( + input=flattened_data, + context=context, + ), + options=options, + ) + + if len(rendered.messages) > 1: + raise Exception('parts template must produce only one message') + + # Convert parts to dicts for Pydantic re-validation when creating new Message + part_rendered: list[dict[str, Any]] = [] + for message in rendered.messages: + for part in message.content: + part_rendered.append(part.model_dump()) + + return part_rendered + + +async def render_message_prompt( + registry: Registry, + input: dict[str, Any], + options: PromptConfig, + prompt_cache: PromptCache, + context: dict[str, Any] | None = None, + history: list[Message] | None = None, +) -> list[Message]: + """Render a messages template (string or list) into Message objects.""" + if isinstance(options.messages, str): + if prompt_cache.messages is None: + prompt_cache.messages = await registry.dotprompt.compile(options.messages) + + if options.metadata: + context = {**(context or {}), 'state': options.metadata.get('state')} + + # Convert history to dict format for template + messages_ = None + if history: + messages_ = [e.model_dump() for e in history] + + # Flatten input and context for template resolution + flattened_data = {**(context or {}), **(input or {})} + rendered = await prompt_cache.messages( + data=DataArgument[dict[str, Any]]( + input=flattened_data, + context=context, + messages=messages_, # type: ignore[arg-type] + ), + options=PromptMetadata( + input=PromptInputConfig( + schema=to_json_schema(options.input_schema) if options.input_schema else None, + ) + ), + ) + return [Message.model_validate(e.model_dump()) for e in rendered.messages] + + elif isinstance(options.messages, list): + return [m if isinstance(m, Message) else Message.model_validate(m) for m in options.messages] + + raise TypeError(f'Unsupported type for messages: {type(options.messages)}') + + +async def render_user_prompt( + registry: Registry, + input: dict[str, Any], + options: PromptConfig, + prompt_cache: PromptCache, + context: dict[str, Any] | None = None, +) -> Message: + """Render the user prompt.""" + msg, prompt_cache.user_prompt = await _render_template( + registry, + Role.USER, + options.prompt, + input, + options.input_schema, + options.metadata, + prompt_cache.user_prompt, + context, + ) + return msg + + +async def render_docs( + input: dict[str, Any], + options: PromptConfig, + context: dict[str, Any] | None = None, +) -> list[Document] | None: + """Return the docs from options (placeholder for future doc rendering).""" + return options.docs + + +def registry_definition_key(name: str, variant: str | None = None, ns: str | None = None) -> str: + """Generate a registry definition key for a prompt. + + Format: "ns/name.variant" where ns and variant are optional. + + Args: + name: The prompt name. + variant: Optional variant name. + ns: Optional namespace. + + Returns: + Registry key string. + """ + parts = [] + if ns: + parts.append(ns) + parts.append(name) + if variant: + parts[-1] = f'{parts[-1]}.{variant}' + return '/'.join(parts) + + +def registry_lookup_key(name: str, variant: str | None = None, ns: str | None = None) -> str: + """Generate a registry lookup key for a prompt. + + Args: + name: The prompt name. + variant: Optional variant name. + ns: Optional namespace. + + Returns: + Registry lookup key string. + """ + return f'/prompt/{registry_definition_key(name, variant, ns)}' + + +def define_partial(registry: Registry, name: str, source: str) -> None: + """Define a partial template in the registry. + + Partials are reusable template fragments that can be included in other prompts. + Files starting with `_` are treated as partials. + + Args: + registry: The registry to register the partial in. + name: The name of the partial. + source: The template source code. + """ + _ = registry.dotprompt.define_partial(name, source) + logger.debug(f'Registered Dotprompt partial "{name}"') + + +def define_helper(registry: Registry, name: str, fn: Callable[..., Any]) -> None: + """Define a Handlebars helper function in the registry. + + Args: + registry: The registry to register the helper in. + name: The name of the helper function. + fn: The helper function to register. + """ + _ = registry.dotprompt.define_helper(name, fn) + logger.debug(f'Registered Dotprompt helper "{name}"') + + +def define_schema(registry: Registry, name: str, schema: type[BaseModel]) -> None: + """Register a Pydantic schema for use in prompts. + + Schemas registered with this function can be referenced by name in + .prompt files using the `output.schema` field. + + Args: + registry: The registry to register the schema in. + name: The name of the schema. + schema: The Pydantic model class to register. + + Example: + ```python + from genkit._ai._prompt import define_schema + + define_schema(registry, 'Recipe', Recipe) + ``` + + Then in a .prompt file: + ```yaml + output: + schema: Recipe + ``` + """ + json_schema = to_json_schema(schema) + registry.register_schema(name, json_schema, schema_type=schema) + logger.debug(f'Registered schema "{name}"') + + +def _transform_prompt_metadata( + raw_metadata: Any, # noqa: ANN401 + variant: str | None, + template: str, + registry_key: str, +) -> dict[str, Any]: + """Transform dotprompt metadata into the format ExecutablePrompt expects.""" + # Convert Pydantic model to dict if needed + if hasattr(raw_metadata, 'model_dump'): + md = raw_metadata.model_dump(by_alias=True) + elif hasattr(raw_metadata, 'dict'): + md = raw_metadata.dict(by_alias=True) # pyright: ignore[reportDeprecated] + else: + md = cast(dict[str, Any], raw_metadata) + + # Preserve raw for accessing maxTurns, toolChoice, etc. + if hasattr(raw_metadata, 'raw'): + md['raw'] = raw_metadata.raw + + if variant: + md['variant'] = variant + + # Clean up null descriptions (matches JS behavior) + output = md.get('output') + if output and isinstance(output, dict): + schema = output.get('schema') + if schema and isinstance(schema, dict) and schema.get('description') is None: + schema.pop('description', None) + + input_cfg = md.get('input') + if input_cfg and isinstance(input_cfg, dict): + schema = input_cfg.get('schema') + if schema and isinstance(schema, dict) and schema.get('description') is None: + schema.pop('description', None) + + # Build metadata structure + metadata = { + **md, + **(md.get('metadata', {}) if isinstance(md.get('metadata'), dict) else {}), + 'type': 'prompt', + 'prompt': {**md, 'template': template}, + } + + raw = md.get('raw') + if raw and isinstance(raw, dict) and raw.get('metadata'): + metadata['metadata'] = {**raw['metadata']} + + return { + 'name': registry_key, + 'model': md.get('model'), + 'config': md.get('config'), + 'tools': md.get('tools'), + 'description': md.get('description'), + 'output': { + 'jsonSchema': output.get('schema') if isinstance(output, dict) else None, + 'format': output.get('format') if isinstance(output, dict) else None, + }, + 'input': { + 'default': input_cfg.get('default') if isinstance(input_cfg, dict) else None, + 'jsonSchema': input_cfg.get('schema') if isinstance(input_cfg, dict) else None, + }, + 'metadata': metadata, + 'maxTurns': raw.get('maxTurns') if isinstance(raw, dict) else None, + 'toolChoice': raw.get('toolChoice') if isinstance(raw, dict) else None, + 'returnToolRequests': raw.get('returnToolRequests') if isinstance(raw, dict) else None, + 'messages': template, + } + + +def load_prompt(registry: Registry, path: Path, filename: str, prefix: str = '', ns: str = '') -> None: + """Load a .prompt file and register it as a lazy-loaded prompt.""" + if not filename.endswith('.prompt'): + raise ValueError(f"Invalid prompt filename: {filename}. Must end with '.prompt'") + + base_name = filename.removesuffix('.prompt') + name = f'{prefix}{base_name}' if prefix else base_name + variant: str | None = None + + if '.' in name: + parts = name.split('.') + name = parts[0] + variant = parts[1] + + file_path = path / (prefix.rstrip('/') + '/' + filename if prefix else filename) + + with Path(file_path).open(encoding='utf-8') as f: + source = f.read() + + parsed_prompt = registry.dotprompt.parse(source) + registry_key = registry_definition_key(name, variant, ns) + + # Memoized prompt instance + _cached_prompt: ExecutablePrompt[Any, Any] | None = None + + async def create_prompt_from_file() -> ExecutablePrompt[Any, Any]: + nonlocal _cached_prompt + if _cached_prompt is not None: + return _cached_prompt + + raw_metadata = await registry.dotprompt.render_metadata(parsed_prompt) + metadata = _transform_prompt_metadata(raw_metadata, variant, parsed_prompt.template, registry_key) + + executable_prompt = ExecutablePrompt( + registry=registry, + variant=metadata.get('variant'), + model=metadata.get('model'), + config=metadata.get('config'), + description=metadata.get('description'), + input_schema=metadata.get('input', {}).get('jsonSchema'), + output_schema=metadata.get('output', {}).get('jsonSchema'), + output_constrained=True if metadata.get('output', {}).get('jsonSchema') else None, + output_format=metadata.get('output', {}).get('format'), + messages=metadata.get('messages'), + max_turns=metadata.get('maxTurns'), + tool_choice=metadata.get('toolChoice'), + return_tool_requests=metadata.get('returnToolRequests'), + metadata=metadata.get('metadata'), + tools=metadata.get('tools'), + name=name, + ns=ns, + ) + + # Wire up action references + definition_key = registry_definition_key(name, variant, ns) + prompt_action = await registry.resolve_action_by_key(create_action_key(ActionKind.PROMPT, definition_key)) + exec_prompt_action = await registry.resolve_action_by_key( + create_action_key(ActionKind.EXECUTABLE_PROMPT, definition_key) + ) + + if prompt_action and prompt_action.kind == ActionKind.PROMPT: + executable_prompt._prompt_action = prompt_action # pyright: ignore[reportPrivateUsage] + setattr(prompt_action, '_executable_prompt', weakref.ref(executable_prompt)) # noqa: B010 + + # Update schemas on actions for Dev UI + for action in [prompt_action, exec_prompt_action]: + if action: + if metadata.get('input', {}).get('jsonSchema'): + action.input_schema = metadata['input']['jsonSchema'] + if metadata.get('output', {}).get('jsonSchema'): + action.output_schema = metadata['output']['jsonSchema'] + + _cached_prompt = executable_prompt + return executable_prompt + + action_metadata: dict[str, object] = { + 'type': 'prompt', + 'lazy': True, + 'source': 'file', + 'prompt': {'name': name, 'variant': variant or ''}, + } + + async def prompt_action_fn(input: Any = None) -> ModelRequest: # noqa: ANN401 + prompt = await create_prompt_from_file() + return await to_generate_request(registry, await prompt.render(input=input)) + + async def executable_prompt_action_fn(input: Any = None) -> GenerateActionOptions: # noqa: ANN401 + prompt = await create_prompt_from_file() + return await prompt.render(input=input) + + action_name = registry_definition_key(name, variant, ns) + prompt_action = registry.register_action( + kind=ActionKind.PROMPT, name=action_name, fn=prompt_action_fn, metadata=action_metadata + ) + executable_prompt_action = registry.register_action( + kind=ActionKind.EXECUTABLE_PROMPT, name=action_name, fn=executable_prompt_action_fn, metadata=action_metadata + ) + + setattr(prompt_action, '_async_factory', create_prompt_from_file) # noqa: B010 + setattr(executable_prompt_action, '_async_factory', create_prompt_from_file) # noqa: B010 + + logger.debug(f'Registered prompt "{registry_key}" from "{file_path}"') + + +def load_prompt_folder_recursively(registry: Registry, dir_path: Path, ns: str, sub_dir: str = '') -> None: + """Recursively load all prompt files from a directory. + + Args: + registry: The registry to register prompts in. + dir_path: Base path to the prompts directory. + ns: Namespace for prompts. + sub_dir: Current subdirectory being processed (for recursion). + """ + full_path = dir_path / sub_dir if sub_dir else dir_path + + if not full_path.exists() or not full_path.is_dir(): + return + + # Iterate through directory entries + try: + for entry in os.scandir(full_path): + if entry.is_file() and entry.name.endswith('.prompt'): + if entry.name.startswith('_'): + # This is a partial + partial_name = entry.name[1:-7] # Remove "_" prefix and ".prompt" suffix + with Path(entry.path).open(encoding='utf-8') as f: + source = f.read() + + # Strip frontmatter if present + if source.startswith('---'): + end_frontmatter = source.find('---', 3) + if end_frontmatter != -1: + source = source[end_frontmatter + 3 :].strip() + + define_partial(registry, partial_name, source) + logger.debug(f'Registered Dotprompt partial "{partial_name}" from "{entry.path}"') + else: + # This is a regular prompt + prefix_with_slash = f'{sub_dir}/' if sub_dir else '' + load_prompt(registry, dir_path, entry.name, prefix_with_slash, ns) + elif entry.is_dir(): + # Recursively process subdirectories + new_sub_dir = os.path.join(sub_dir, entry.name) if sub_dir else entry.name + load_prompt_folder_recursively(registry, dir_path, ns, new_sub_dir) + except PermissionError: + logger.warning(f'Permission denied accessing directory: {full_path}') + except Exception as e: + logger.exception(f'Error loading prompts from {full_path}', exc_info=e) + + +def load_prompt_folder(registry: Registry, dir_path: str | Path = './prompts', ns: str = '') -> None: + """Load all prompt files from a directory. + + This is the main entry point for loading prompts from a directory. + It recursively processes all `.prompt` files and registers them. + + Args: + registry: The registry to register prompts in. + dir_path: Path to the prompts directory. Defaults to './prompts'. + ns: Namespace for prompts. Defaults to 'dotprompt'. + """ + path = Path(dir_path).resolve() + + if not path.exists(): + logger.warning(f'Prompt directory does not exist: {path}') + return + + if not path.is_dir(): + logger.warning(f'Prompt path is not a directory: {path}') + return + + load_prompt_folder_recursively(registry, path, ns, '') + logger.info(f'Loaded prompts from directory: {path}') + + +async def lookup_prompt(registry: Registry, name: str, variant: str | None = None) -> ExecutablePrompt[Any, Any]: + """Look up a prompt by name from the registry.""" + # Try without namespace first (for programmatic prompts) + # Use create_action_key to build the full key: "/prompt/" + definition_key = registry_definition_key(name, variant, None) + lookup_key = create_action_key(ActionKind.PROMPT, definition_key) + action = await registry.resolve_action_by_key(lookup_key) + + # If not found and no namespace was specified, try with default 'dotprompt' namespace + # (for file-based prompts) + if not action: + definition_key = registry_definition_key(name, variant, 'dotprompt') + lookup_key = create_action_key(ActionKind.PROMPT, definition_key) + action = await registry.resolve_action_by_key(lookup_key) + + if action: + # First check if we've stored the ExecutablePrompt directly + prompt_ref = getattr(action, '_executable_prompt', None) + if prompt_ref is not None: + if isinstance(prompt_ref, weakref.ReferenceType): + resolved = prompt_ref() + if resolved is not None: + return resolved + if isinstance(prompt_ref, ExecutablePrompt): + return prompt_ref + # Otherwise, create it from the factory (lazy loading) + async_factory = getattr(action, '_async_factory', None) + if callable(async_factory): + # Cast to async callable - getattr returns object but we've verified it's callable + async_factory_fn = cast(Callable[[], Awaitable[ExecutablePrompt]], async_factory) + executable_prompt = await async_factory_fn() + if getattr(action, '_executable_prompt', None) is None: + setattr(action, '_executable_prompt', executable_prompt) # noqa: B010 + return executable_prompt + # This shouldn't happen if prompts are loaded correctly + raise GenkitError( + status='INTERNAL', + message=f'Prompt action found but no ExecutablePrompt available for {name}', + ) + + variant_str = f' (variant {variant})' if variant else '' + raise GenkitError( + status='NOT_FOUND', + message=f'Prompt {name}{variant_str} not found', + ) + + +async def prompt( + registry: Registry, + name: str, + variant: str | None = None, +) -> ExecutablePrompt[Any, Any]: + """Look up a prompt by name and optional variant.""" + return await lookup_prompt(registry, name, variant) + + +# Renamed β€” use ModelStreamResponse diff --git a/py/packages/genkit/src/genkit/blocks/resource.py b/py/packages/genkit/src/genkit/_ai/_resource.py similarity index 54% rename from py/packages/genkit/src/genkit/blocks/resource.py rename to py/packages/genkit/src/genkit/_ai/_resource.py index 668fdf1d0a..8a936d2ebd 100644 --- a/py/packages/genkit/src/genkit/blocks/resource.py +++ b/py/packages/genkit/src/genkit/_ai/_resource.py @@ -12,15 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Resource module for defining and managing resources. - -Resources in Genkit represent addressable content or data processing units containing -unstructured data (Post, PDF, etc.) that can be retrieved or generated. They are -identified by URIs (e.g. `file://`, `http://`, `gs://`) and can be static (fixed URI) -or dynamic (using URI templates). -This module provides tools to define resource actions that can resolve these URIs -and return content (`ResourceOutput`) containing `Part`s. -""" +"""Resource module for defining and managing resources.""" import inspect import re @@ -29,25 +21,13 @@ from pydantic import BaseModel -from genkit.aio import ensure_async -from genkit.core.action import Action, ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry -from genkit.core.typing import Metadata, Part +from genkit._core._action import Action, ActionKind, ActionRunContext +from genkit._core._registry import Registry +from genkit._core._typing import Metadata, Part class ResourceOptions(TypedDict, total=False): - """Options for defining a resource. - - Attributes: - name: Resource name. If not specified, uri or template will be used as name. - uri: The URI of the resource. Can contain template variables for simple matches, - but `template` is preferred for pattern matching. - template: The URI template (ex. `my://resource/{id}`). See RFC6570 for specification. - Used for matching variable resources. - description: A description of the resource, used for documentation and discovery. - metadata: Arbitrary metadata to attach to the resource action. - """ + """Options for defining a resource (name, uri/template, description, metadata).""" name: str uri: str @@ -57,47 +37,27 @@ class ResourceOptions(TypedDict, total=False): class ResourceInput(BaseModel): - """Input structure for a resource request. - - Attributes: - uri: The full URI being requested/resolved. - """ + """Input for a resource request containing the URI to resolve.""" uri: str class ResourceOutput(BaseModel): - """Output structure from a resource resolution. - - Attributes: - content: A list of `Part` objects representing the resource content. - """ + """Output from a resource resolution containing content parts.""" content: list[Part] ResourcePayload = ResourceOutput | dict[str, Any] -ResourceFn = Callable[..., Awaitable[ResourcePayload] | ResourcePayload] +ResourceFn = Callable[..., Awaitable[ResourcePayload]] ResourceArgument = Action | str async def resolve_resources(registry: Registry, resources: list[ResourceArgument] | None = None) -> list[Action]: - """Resolves a list of resource names or actions into a list of Action objects. - - Args: - registry: The registry to lookup resources in. - resources: A list of resource references, which can be either direct `Action` - objects or strings (names/URIs). - - Returns: - A list of resolved `Action` objects. - - Raises: - ValueError: If a resource reference is invalid or cannot be found. - """ + """Resolve resource names/actions to Action objects.""" if not resources: return [] @@ -113,25 +73,11 @@ async def resolve_resources(registry: Registry, resources: list[ResourceArgument async def lookup_resource_by_name(registry: Registry, name: str) -> Action: - """Looks up a resource action by name in the registry. - - Tries to resolve the name directly, or with common prefixes like `/resource/` - or `/dynamic-action-provider/`. - - Args: - registry: The registry to search. - name: The name or URI of the resource to lookup. - - Returns: - The found `Action`. - - Raises: - ValueError: If the resource cannot be found. - """ + """Look up a resource action by name, trying common prefixes.""" resource = ( - await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), name) - or await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), f'/resource/{name}') - or await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), f'/dynamic-action-provider/{name}') + await registry.resolve_action(ActionKind.RESOURCE, name) + or await registry.resolve_action(ActionKind.RESOURCE, f'/resource/{name}') + or await registry.resolve_action(ActionKind.RESOURCE, f'/dynamic-action-provider/{name}') ) if not resource: raise ValueError(f'Resource {name} not found') @@ -139,22 +85,10 @@ async def lookup_resource_by_name(registry: Registry, name: str) -> Action: def define_resource(registry: Registry, opts: ResourceOptions, fn: ResourceFn) -> Action: - """Defines a resource and registers it with the given registry. - - This creates a resource action that can handle requests for a specific URI - or URI template. - - Args: - registry: The registry to register the resource with. - opts: Options defining the resource (name, uri, template, etc.). - fn: The function that implements resource content retrieval. - - Returns: - The registered `Action` for the resource. - """ + """Register a resource action for a specific URI or template.""" action = dynamic_resource(opts, fn) - action.matches = create_matcher(opts.get('uri'), opts.get('template')) # type: ignore[attr-defined] + action.matches = create_matcher(opts.get('uri'), opts.get('template')) # Mark as not dynamic since it's being registered action.metadata['dynamic'] = False @@ -165,40 +99,15 @@ def define_resource(registry: Registry, opts: ResourceOptions, fn: ResourceFn) - def resource(opts: ResourceOptions, fn: ResourceFn) -> Action: - """Defines a dynamic resource action without immediate registration. - - This is an alias for `dynamic_resource`. Useful for defining resources that - might be registered later or used as standalone actions. - - Args: - opts: Options defining the resource. - fn: The resource implementation function. - - Returns: - The created `Action`. - """ + """Create a dynamic resource action (alias for dynamic_resource).""" return dynamic_resource(opts, fn) def dynamic_resource(opts: ResourceOptions, fn: ResourceFn) -> Action: - """Defines a dynamic resource action. - - Creates an `Action` of kind `RESOURCE` that wraps the provided function. - The wrapper handles: - 1. Input validation and matching against the URI/Template. - 2. Execution of the resource function. - 3. Post-processing of output to attach metadata (like parent resource info). - - Args: - opts: Options including `uri` or `template` for matching. - fn: The function performing the resource retrieval. - - Returns: - An `Action` configured as a resource. + """Create a resource Action that matches URIs and executes the given function.""" + if not inspect.iscoroutinefunction(fn): + raise TypeError('fn must be an async function') - Raises: - ValueError: If neither `uri` nor `template` is provided in options. - """ uri = opts.get('uri') or opts.get('template') if not uri: raise ValueError('must specify either uri or template options') @@ -215,18 +124,17 @@ async def wrapped_fn(input_data: ResourceInput, ctx: ActionRunContext) -> Resour raise ValueError(f'input {input_data} did not match template {uri}') sig = inspect.signature(fn) - afn = ensure_async(fn) n_params = len(sig.parameters) if n_params == 0: - parts = await afn() + parts = await fn() elif n_params == 1: - parts = await afn(input_data) + parts = await fn(input_data) else: - parts = await afn(input_data, ctx) + parts = await fn(input_data, ctx) # Post-processing parts to add metadata - content_list = parts.content if hasattr(parts, 'content') else parts.get('content', []) + content_list = parts.content if isinstance(parts, ResourceOutput) else parts.get('content', []) for p in content_list: if isinstance(p, Part): @@ -236,17 +144,23 @@ async def wrapped_fn(input_data: ResourceInput, ctx: ActionRunContext) -> Resour if p.metadata is None: # Different Part types have different metadata types (Metadata or dict) # dict works for both types at runtime - p.metadata = {} # pyright: ignore[reportAttributeAccessIssue] + # pyrefly:ignore[bad-assignment] + p.metadata = {} # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[invalid-assignment] if isinstance(p.metadata, Metadata): p_metadata = p.metadata.root elif isinstance(p.metadata, dict): p_metadata = p.metadata else: # dict works for both Part types at runtime - p.metadata = {} # pyright: ignore[reportAttributeAccessIssue] + # pyrefly:ignore[bad-assignment] + p.metadata = {} # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[invalid-assignment] p_metadata = p.metadata template = opts.get('template') + # p_metadata is guaranteed to be dict here due to isinstance checks above, + # but type checkers can't narrow the union type. Use cast to inform them. + p_metadata = cast(dict[str, Any], p_metadata) + if 'resource' in p_metadata: if 'parent' not in p_metadata['resource']: p_metadata['resource']['parent'] = {'uri': input_data.uri} @@ -278,7 +192,7 @@ async def wrapped_fn(input_data: ResourceInput, ctx: ActionRunContext) -> Resour act = Action( name=name, - kind=cast(ActionKind, ActionKind.RESOURCE), + kind=ActionKind.RESOURCE, fn=wrapped_fn, metadata={ 'resource': { @@ -295,15 +209,7 @@ async def wrapped_fn(input_data: ResourceInput, ctx: ActionRunContext) -> Resour def create_matcher(uri: str | None, template: str | None) -> Callable[[object], bool]: - """Creates a matching function for resource validation. - - Args: - uri: Optional fixed URI string. - template: Optional URI template string. - - Returns: - A callable that takes an object (expected to be ResourceInput) and returns True if it matches. - """ + """Create a matcher function for URI or template matching.""" def matcher(input_data: object) -> bool: if not isinstance(input_data, ResourceInput): @@ -318,33 +224,12 @@ def matcher(input_data: object) -> bool: def is_dynamic_resource_action(action: Action) -> bool: - """Checks if an action is a dynamic resource (not registered). - - Args: - action: The action to check. - - Returns: - True if the action is a dynamic resource, False otherwise. - """ + """Check if an action is a dynamic (unregistered) resource.""" return action.kind == ActionKind.RESOURCE and bool(action.metadata.get('dynamic', True)) def matches_uri_template(template: str, uri: str) -> dict[str, str] | None: - """Check if a URI matches a template and extract parameters. - - Args: - template: URI template with {param} placeholders (e.g., "file://{path}"). - uri: The URI to match against the template. - - Returns: - Dictionary of extracted parameters if match, None otherwise. - - Examples: - >>> matches_uri_template('file://{path}', 'file:///home/user/doc.txt') - {'path': '/home/user/doc.txt'} - >>> matches_uri_template('user://{id}/profile', 'user://123/profile') - {'id': '123'} - """ + """Match URI against template, returning extracted params or None.""" # Split template into parts: text and {param} placeholders parts = re.split(r'(\{[\w\+]+\})', template) pattern_parts = [] @@ -372,38 +257,23 @@ def matches_uri_template(template: str, uri: str) -> dict[str, str] | None: async def find_matching_resource( registry: Registry, dynamic_resources: list[Action] | None, input_data: ResourceInput ) -> Action | None: - """Finds a matching resource action. - - Checks dynamic resources first, then the registry. - - Args: - registry: The registry to search. - dynamic_resources: Optional list of dynamic resource actions to check first. - input_data: The resource input containing the URI matched against. - - Returns: - The matching Action or None. - """ + """Find a matching resource action from dynamic resources or registry.""" if dynamic_resources: for action in dynamic_resources: - if ( - hasattr(action, 'matches') and callable(action.matches) and action.matches(input_data) # type: ignore[attr-defined] - ): + if hasattr(action, 'matches') and callable(action.matches) and action.matches(input_data): return action # Try exact match in registry - resource = await registry.resolve_action(cast(ActionKind, ActionKind.RESOURCE), input_data.uri) + resource = await registry.resolve_action(ActionKind.RESOURCE, input_data.uri) if resource: return resource # Iterate all resources to check for matches (e.g. templates) # This is less efficient but necessary for template matching if not optimized - resources = await registry.resolve_actions_by_kind(cast(ActionKind, ActionKind.RESOURCE)) + resources = await registry.resolve_actions_by_kind(ActionKind.RESOURCE) for action in resources.values(): - if ( - hasattr(action, 'matches') and callable(action.matches) and action.matches(input_data) # type: ignore[attr-defined] - ): + if hasattr(action, 'matches') and callable(action.matches) and action.matches(input_data): return action return None diff --git a/py/packages/genkit/src/genkit/ai/_runtime.py b/py/packages/genkit/src/genkit/_ai/_runtime.py similarity index 97% rename from py/packages/genkit/src/genkit/ai/_runtime.py rename to py/packages/genkit/src/genkit/_ai/_runtime.py index 388321cf49..278944fd0a 100644 --- a/py/packages/genkit/src/genkit/ai/_runtime.py +++ b/py/packages/genkit/src/genkit/_ai/_runtime.py @@ -26,10 +26,9 @@ from pathlib import Path from types import TracebackType -from genkit.core.constants import DEFAULT_GENKIT_VERSION -from genkit.core.logging import get_logger - -from ._server import ServerSpec +from genkit._core._constants import GENKIT_VERSION +from genkit._core._logger import get_logger +from genkit._core._reflection import ServerSpec logger = get_logger(__name__) @@ -103,7 +102,7 @@ def _create_and_write_runtime_file(runtime_dir: Path, spec: ServerSpec) -> Path: 'reflectionApiSpecVersion': 1, 'id': runtime_id, 'pid': pid, - 'genkitVersion': 'py/' + DEFAULT_GENKIT_VERSION, + 'genkitVersion': 'py/' + GENKIT_VERSION, 'reflectionServerUrl': spec.url, 'timestamp': current_datetime.isoformat(), }) diff --git a/py/packages/genkit/src/genkit/_ai/_testing.py b/py/packages/genkit/src/genkit/_ai/_testing.py new file mode 100644 index 0000000000..f33560c809 --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_testing.py @@ -0,0 +1,397 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use it except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Internal testing utilities for Genkit AI (mock models, test_models).""" + +import json +from copy import deepcopy +from typing import Any, TypedDict + +from pydantic import BaseModel, Field + +from genkit._core._action import Action, ActionKind, ActionRunContext +from genkit._core._tracing import run_in_new_span +from genkit._core._typing import ( + Media, + MediaPart, + ModelInfo, + Part, + Role, + SpanMetadata, + TextPart, +) +from genkit.model import Message, ModelRequest, ModelResponse, ModelResponseChunk + +from ._aio import Genkit + + +class ProgrammableModel: + """A configurable model implementation for testing.""" + + def __init__(self) -> None: + self._request_idx: int = 0 + self.request_count: int = 0 + self.responses: list[ModelResponse] = [] + self.chunks: list[list[ModelResponseChunk]] | None = None + self.last_request: ModelRequest | None = None + + def reset(self) -> None: + self._request_idx = 0 + self.request_count = 0 + self.responses = [] + self.chunks = None + self.last_request = None + + async def model_fn( + self, + request: ModelRequest, + ctx: ActionRunContext, + ) -> ModelResponse: + self.last_request = deepcopy(request) + self.request_count += 1 + + response = self.responses[self._request_idx] + if self.chunks and self._request_idx < len(self.chunks): + for chunk in self.chunks[self._request_idx]: + ctx.send_chunk(chunk) + self._request_idx += 1 + return response + + +def define_programmable_model( + ai: Genkit, + name: str = 'programmableModel', +) -> tuple[ProgrammableModel, Action]: + pm = ProgrammableModel() + + async def model_fn( + request: ModelRequest, + ctx: ActionRunContext, + ) -> ModelResponse: + return await pm.model_fn(request, ctx) + + action = ai.define_model(name=name, fn=model_fn) + + return (pm, action) + + +class EchoModel: + """A model implementation that echoes back the input with metadata.""" + + def __init__(self, stream_countdown: bool = False) -> None: + self.last_request: ModelRequest | None = None + self.stream_countdown: bool = stream_countdown + + async def model_fn( + self, + request: ModelRequest, + ctx: ActionRunContext, + ) -> ModelResponse: + self.last_request = request + + merged_txt = '' + messages = request.messages.root if hasattr(request.messages, 'root') else request.messages # pyright: ignore[reportAttributeAccessIssue] + for m in messages: # ty: ignore[not-iterable] + merged_txt += f' {m.role}: ' + ','.join( + json.dumps(p.root.text) if p.root.text is not None else '""' for p in m.content + ) + echo_resp = f'[ECHO]{merged_txt}' + + if request.config: + if hasattr(request.config, 'model_dump_json'): + config_json = request.config.model_dump_json() + else: + config_json = json.dumps(request.config, separators=(',', ':')) + else: + config_json = '{}' + if request.config and config_json != '{}': + echo_resp += f' {config_json}' + tools_list = request.tools.root if hasattr(request.tools, 'root') else request.tools # pyright: ignore[reportAttributeAccessIssue,reportOptionalMemberAccess] + if tools_list: + echo_resp += f' tools={",".join(t.name for t in tools_list)}' # ty: ignore[not-iterable] + if request.tool_choice is not None: + echo_resp += f' tool_choice={request.tool_choice}' + output_dict: dict[str, object] = {} + if request.output_format: + output_dict['format'] = request.output_format + if request.output_schema: + output_dict['schema'] = request.output_schema + if request.output_constrained is not None: + output_dict['constrained'] = request.output_constrained + if request.output_content_type: + output_dict['contentType'] = request.output_content_type + output_json = json.dumps(output_dict, separators=(',', ':')) if output_dict else '{}' + if output_dict and output_json != '{}': + echo_resp += f' output={output_json}' + + if self.stream_countdown: + for i, countdown in enumerate(['3', '2', '1']): + ctx.send_chunk( + ModelResponseChunk(role=Role.MODEL, index=i, content=[Part(root=TextPart(text=countdown))]) + ) + + return ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=echo_resp))])) + + +def define_echo_model( + ai: Genkit, + name: str = 'echoModel', + stream_countdown: bool = False, +) -> tuple[EchoModel, Action]: + echo = EchoModel(stream_countdown=stream_countdown) + + async def model_fn( + request: ModelRequest, + ctx: ActionRunContext, + ) -> ModelResponse: + return await echo.model_fn(request, ctx) + + action = ai.define_model(name=name, fn=model_fn) + + return (echo, action) + + +class StaticResponseModel: + """A model that always returns the same static response.""" + + def __init__(self, message: dict[str, Any]) -> None: + self.response_message: Message = Message.model_validate(message) + self.last_request: ModelRequest | None = None + self.request_count: int = 0 + + async def model_fn( + self, + request: ModelRequest, + _ctx: ActionRunContext, + ) -> ModelResponse: + self.last_request = request + self.request_count += 1 + return ModelResponse(message=self.response_message) + + +def define_static_response_model( + ai: Genkit, + message: dict[str, Any], + name: str = 'staticModel', +) -> tuple[StaticResponseModel, Action]: + static = StaticResponseModel(message) + + async def model_fn( + request: ModelRequest, + ctx: ActionRunContext, + ) -> ModelResponse: + return await static.model_fn(request, ctx) + + action = ai.define_model(name=name, fn=model_fn) + + return (static, action) + + +class SkipTestError(Exception): + """Exception raised to skip a test case.""" + + +def skip() -> None: + raise SkipTestError() + + +class ModelTestError(TypedDict, total=False): + message: str + stack: str | None + + +class ModelTestResult(TypedDict, total=False): + name: str + passed: bool + skipped: bool + error: ModelTestError + + +class TestCaseReport(TypedDict): + description: str + models: list[ModelTestResult] + + +TestReport = list[TestCaseReport] + + +class GablorkenInput(BaseModel): + value: float = Field(..., description='The value to calculate gablorken for') + + +async def test_models(ai: Genkit, models: list[str]) -> TestReport: + """Run a standard test suite against one or more models.""" + + @ai.tool(name='gablorkenTool') + async def gablorken_tool(input: GablorkenInput) -> float: + """Calculate the gablorken of a value.""" + return (input.value**3) + 1.407 + + async def get_model_info(model_name: str) -> ModelInfo | None: + model_action = await ai.registry.resolve_action(ActionKind.MODEL, model_name) + if model_action and model_action.metadata: + info_obj = model_action.metadata.get('model') + if isinstance(info_obj, ModelInfo): + return info_obj + return None + + async def test_basic_hi(model: str) -> None: + response = await ai.generate(model=model, prompt='just say "Hi", literally') + got = response.text.strip() + assert 'hi' in got.lower(), f'Expected "Hi" in response, got: {got}' + + async def test_multimodal(model: str) -> None: + info = await get_model_info(model) + if not (info and info.supports and info.supports.media): + skip() + + test_image = ( + 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAIAAACQkWg2' + 'AAABhGlDQ1BJQ0MgcHJvZmlsZQAAKJF9kT1Iw0AcxV9TpSoVETOIOGSoulgQFXHU' + 'KhShQqgVWnUwufRDaNKQtLg4Cq4FBz8Wqw4uzro6uAqC4AeIs4OToouU+L+k0CLG' + 'g+N+vLv3uHsHCLUi0+22MUA3ylYyHpPSmRUp9IpOhCCiFyMKs81ZWU7Ad3zdI8DX' + 'uyjP8j/35+jWsjYDAhLxDDOtMvE68dRm2eS8TyyygqIRnxOPWnRB4keuqx6/cc67' + 'LPBM0Uol54hFYinfwmoLs4KlE08SRzTdoHwh7bHGeYuzXqywxj35C8NZY3mJ6zQH' + 'EccCFiFDgooKNlBEGVFaDVJsJGk/5uMfcP0yuVRybYCRYx4l6FBcP/gf/O7Wzk2M' + 'e0nhGND+4jgfQ0BoF6hXHef72HHqJ0DwGbgymv5SDZj+JL3a1CJHQM82cHHd1NQ9' + '4HIH6H8yFUtxpSBNIZcD3s/omzJA3y3Qter11tjH6QOQoq4SN8DBITCcp+w1n3d3' + 'tPb275lGfz9aC3Kd0jYiSQAAAAlwSFlzAAAuIwAALiMBeKU/dgAAAAd0SU1FB+gJ' + 'BxQRO1/5qB8AAAAZdEVYdENvbW1lbnQAQ3JlYXRlZCB3aXRoIEdJTVBXgQ4XAAAA' + 'sUlEQVQoz61SMQqEMBDcO5SYToUE/IBPyRMCftAH+INUviApUwYjNkKCVcTiQK7I' + 'HSw45czODrMswCOQUkopEQZjzDiOWemdZfu+b5oGYYgx1nWNMPwB2vACAK01Y4wQ' + '8qGqqirL8jzPlNI9t64r55wQUgBA27be+xDCfaJhGJxzSqnv3UKIn7ne+2VZEB2s' + 'tZRSRLN93+d5RiRs28Y5RySEEI7jyEpFlp2mqeu6Zx75ApQwPdsIcq0ZAAAAAElF' + 'TkSuQmCC' + ) + + response = await ai.generate( + model=model, + prompt=[ + Part(root=MediaPart(media=Media(url=test_image))), + Part(root=TextPart(text='what math operation is this? plus, minus, multiply or divide?')), + ], + ) + got = response.text.strip().lower() + assert 'plus' in got, f'Expected "plus" in response, got: {got}' + + async def test_history(model: str) -> None: + info = await get_model_info(model) + if not (info and info.supports and info.supports.multiturn): + skip() + + response1 = await ai.generate(model=model, prompt='My name is Glorb') + response2 = await ai.generate( + model=model, + prompt="What's my name?", + messages=response1.messages, + ) + got = response2.text.strip() + assert 'Glorb' in got, f'Expected "Glorb" in response, got: {got}' + + async def test_system_prompt(model: str) -> None: + response = await ai.generate( + model=model, + prompt='Hi', + messages=[ + Message.model_validate({ + 'role': 'system', + 'content': [{'text': 'If the user says "Hi", just say "Bye"'}], + }), + ], + ) + got = response.text.strip() + assert 'Bye' in got, f'Expected "Bye" in response, got: {got}' + + async def test_structured_output(model: str) -> None: + class PersonInfo(BaseModel): + name: str + occupation: str + + response = await ai.generate( + model=model, + prompt='extract data as json from: Jack was a Lumberjack', + output_schema=PersonInfo, + ) + got = response.output + assert got is not None, 'Expected structured output' + if isinstance(got, BaseModel): + got = got.model_dump() + + assert isinstance(got, dict), f'Expected output to be a dict or BaseModel, got {type(got)}' + assert got.get('name') == 'Jack', f"Expected name='Jack', got: {got.get('name')}" + assert got.get('occupation') == 'Lumberjack', f"Expected occupation='Lumberjack', got: {got.get('occupation')}" + + async def test_tool_calling(model: str) -> None: + info = await get_model_info(model) + if not (info and info.supports and info.supports.tools): + skip() + + response = await ai.generate( + model=model, + prompt='what is a gablorken of 2? use provided tool', + tools=['gablorkenTool'], + ) + got = response.text.strip() + assert '9.407' in got, f'Expected "9.407" in response, got: {got}' + + tests: dict[str, Any] = { + 'basic hi': test_basic_hi, + 'multimodal': test_multimodal, + 'history': test_history, + 'system prompt': test_system_prompt, + 'structured output': test_structured_output, + 'tool calling': test_tool_calling, + } + + report: TestReport = [] + + with run_in_new_span(SpanMetadata(name='testModels'), labels={'genkit:type': 'testSuite'}): + for test_name, test_fn in tests.items(): + with run_in_new_span(SpanMetadata(name=test_name), labels={'genkit:type': 'testCase'}): + case_report: TestCaseReport = { + 'description': test_name, + 'models': [], + } + + for model in models: + model_result: ModelTestResult = { + 'name': model, + 'passed': True, + } + + try: + await test_fn(model) + except SkipTestError: + model_result['passed'] = False + model_result['skipped'] = True + except AssertionError as e: + model_result['passed'] = False + model_result['error'] = { + 'message': str(e), + 'stack': None, + } + except Exception as e: + model_result['passed'] = False + model_result['error'] = { + 'message': str(e), + 'stack': None, + } + + case_report['models'].append(model_result) + + report.append(case_report) + + return report diff --git a/py/packages/genkit/src/genkit/_ai/_tools.py b/py/packages/genkit/src/genkit/_ai/_tools.py new file mode 100644 index 0000000000..f534cf9362 --- /dev/null +++ b/py/packages/genkit/src/genkit/_ai/_tools.py @@ -0,0 +1,151 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tool-specific types and utilities for the Genkit framework.""" + +import inspect +from collections.abc import Callable +from functools import wraps +from typing import Any, NoReturn, ParamSpec, TypeVar, cast + +from genkit._core._action import ActionKind, ActionRunContext +from genkit._core._registry import Registry +from genkit._core._typing import Metadata, Part, ToolRequest, ToolRequestPart, ToolResponse, ToolResponsePart + +P = ParamSpec('P') +T = TypeVar('T') + + +class ToolRunContext(ActionRunContext): + """Tool execution context with interrupt support.""" + + def __init__( + self, + ctx: ActionRunContext, + ) -> None: + """Initialize from parent ActionRunContext.""" + super().__init__(context=ctx.context) + + def interrupt(self, metadata: dict[str, Any] | None = None) -> NoReturn: + """Raise ToolInterruptError to pause execution (e.g., for user input).""" + raise ToolInterruptError(metadata=metadata) + + +# TODO(#4346): make this extend GenkitError once it has INTERRUPTED status +class ToolInterruptError(Exception): + """Controlled interruption of tool execution (e.g., to request user input).""" + + def __init__(self, metadata: dict[str, Any] | None = None) -> None: + """Initialize with optional interrupt metadata.""" + super().__init__() + self.metadata: dict[str, Any] = metadata or {} + + +def tool_response( + interrupt: Part | ToolRequestPart, + response_data: object | None = None, + metadata: dict[str, Any] | None = None, +) -> Part: + """Create a ToolResponse Part for an interrupted tool request.""" + # TODO(#4347): validate against tool schema + tool_request = interrupt.root.tool_request if isinstance(interrupt, Part) else interrupt.tool_request + + interrupt_metadata = True + if isinstance(metadata, Metadata): + interrupt_metadata = metadata.root + elif metadata: + interrupt_metadata = metadata + + tr = cast(ToolRequest, tool_request) + return Part( + root=ToolResponsePart( + tool_response=ToolResponse( + name=tr.name, + ref=tr.ref, + output=response_data, + ), + metadata=Metadata( + root={ + 'interruptResponse': interrupt_metadata, + } + ), + ) + ) + + +def _get_func_description(func: Callable[..., Any], description: str | None = None) -> str: + """Return description if provided, otherwise use the function's docstring.""" + if description is not None: + return description + if func.__doc__ is not None: + return func.__doc__ + return '' + + +def define_tool( + registry: Registry, + func: Callable[P, T], + name: str | None = None, + description: str | None = None, +) -> Callable[P, T]: + """Register a function as a tool. + + Args: + registry: The registry to register the tool in. + func: The async function to register as a tool. Must be a coroutine function. + name: Optional name for the tool. Defaults to the function name. + description: Optional description. Defaults to the function's docstring. + + Raises: + TypeError: If func is not an async function. + """ + # All Python functions have __name__, but ty is strict about Callable protocol + if not inspect.iscoroutinefunction(func): + raise TypeError(f'Tool function must be async. Got sync function: {func.__name__}') # ty: ignore[unresolved-attribute] + + tool_name = name if name is not None else getattr(func, '__name__', 'unnamed_tool') + tool_description = _get_func_description(func, description) + + input_spec = inspect.getfullargspec(func) + + func_any = cast(Callable[..., Any], func) + + async def tool_fn_wrapper(*args: Any) -> Any: # noqa: ANN401 + # Dynamic dispatch based on function signature - pyright can't verify ParamSpec here + match len(input_spec.args): + case 0: + return await func_any() + case 1: + return await func_any(args[0]) + case 2: + return await func_any(args[0], ToolRunContext(cast(ActionRunContext, args[1]))) + case _: + raise ValueError('tool must have 0-2 args...') + + action = registry.register_action( + name=tool_name, + kind=ActionKind.TOOL, + description=tool_description, + fn=tool_fn_wrapper, + metadata_fn=func, + ) + + @wraps(func) + async def wrapper(*args: P.args, **kwargs: P.kwargs) -> Any: # noqa: ANN401 + action_any = cast(Any, action) + return (await action_any.run(*args, **kwargs)).response + + return cast(Callable[P, T], wrapper) diff --git a/py/packages/genkit/src/genkit/_core/_action.py b/py/packages/genkit/src/genkit/_core/_action.py new file mode 100644 index 0000000000..4b985f90a1 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_action.py @@ -0,0 +1,601 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Action module for defining and managing remotely callable functions.""" + +import asyncio +import inspect +import json +import time +from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, ClassVar, Generic, cast, get_type_hints + +from opentelemetry.trace import Span +from opentelemetry.util import types as otel_types +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic.alias_generators import to_camel +from typing_extensions import Never, TypeVar + +from genkit._core._channel import Channel +from genkit._core._compat import StrEnum +from genkit._core._error import GenkitError +from genkit._core._trace._path import build_path +from genkit._core._tracing import tracer + +# ============================================================================= +# Span attribute types and tracing helpers +# ============================================================================= + +# Type alias for span attribute values +SpanAttributeValue = otel_types.AttributeValue + +# Context variable to track parent path across nested spans +_parent_path_context: ContextVar[str] = ContextVar('genkit_parent_path', default='') + + +@contextmanager +def _save_parent_path() -> Generator[None, None, None]: + """Context manager to save and restore parent path.""" + saved = _parent_path_context.get() + try: + yield + finally: + _parent_path_context.set(saved) + + +def _record_input_metadata( + span: Span, + kind: str, + name: str, + span_metadata: dict[str, SpanAttributeValue] | None, + input: object | None, +) -> None: + """Records input metadata onto an OpenTelemetry span for a Genkit action.""" + span.set_attribute('genkit:type', 'action') + span.set_attribute('genkit:metadata:subtype', kind) + span.set_attribute('genkit:name', name) + if input is not None: + input_json = input.model_dump_json() if isinstance(input, BaseModel) else json.dumps(input) + span.set_attribute('genkit:input', input_json) + + # Build and set path attributes (qualified path with full annotations) + parent_path = _parent_path_context.get() + qualified_path = build_path(name, parent_path, 'action', kind) + + span.set_attribute('genkit:path', qualified_path) + span.set_attribute('genkit:qualifiedPath', qualified_path) + + # Update context for nested spans + _parent_path_context.set(qualified_path) + + if span_metadata is not None: + for meta_key in span_metadata: + span.set_attribute(meta_key, span_metadata[meta_key]) + + +def _record_output_metadata(span: Span, output: object) -> None: + """Records output metadata onto an OpenTelemetry span for a Genkit action.""" + span.set_attribute('genkit:state', 'success') + try: + output_json = output.model_dump_json() if isinstance(output, BaseModel) else json.dumps(output) + span.set_attribute('genkit:output', output_json) + except Exception: + # Fallback for non-serializable output + span.set_attribute('genkit:output', str(output)) + + +# ============================================================================= +# Action types +# ============================================================================= + +# Type alias for action name. +ActionName = str + + +class ActionKind(StrEnum): + """Types of actions that can be registered.""" + + BACKGROUND_MODEL = 'background-model' + CANCEL_OPERATION = 'cancel-operation' + CHECK_OPERATION = 'check-operation' + CUSTOM = 'custom' + DYNAMIC_ACTION_PROVIDER = 'dynamic-action-provider' + EMBEDDER = 'embedder' + EVALUATOR = 'evaluator' + EXECUTABLE_PROMPT = 'executable-prompt' + FLOW = 'flow' + INDEXER = 'indexer' + MODEL = 'model' + PROMPT = 'prompt' + RERANKER = 'reranker' + RESOURCE = 'resource' + RETRIEVER = 'retriever' + TOOL = 'tool' + UTIL = 'util' + + +ResponseT = TypeVar('ResponseT') + + +class ActionResponse(BaseModel, Generic[ResponseT]): + """Response from an action with trace ID.""" + + model_config: ClassVar[ConfigDict] = ConfigDict( + extra='forbid', populate_by_name=True, alias_generator=to_camel, arbitrary_types_allowed=True + ) + + response: ResponseT + trace_id: str + span_id: str = '' + + +ChunkT_co = TypeVar('ChunkT_co', covariant=True) +OutputT_co = TypeVar('OutputT_co', covariant=True) + + +class StreamResponse(Generic[ChunkT_co, OutputT_co]): + """Wrapper for streaming action results.""" + + def __init__( + self, + stream: AsyncIterator[ChunkT_co], + response: Awaitable[OutputT_co], + ) -> None: + self._stream = stream + self._response = response + + @property + def stream(self) -> AsyncIterator[ChunkT_co]: + return self._stream + + @property + def response(self) -> Awaitable[OutputT_co]: + return self._response + + +class ActionMetadataKey(StrEnum): + """Keys for action metadata.""" + + INPUT_KEY = 'inputSchema' + OUTPUT_KEY = 'outputSchema' + RETURN = 'return' + + +# ============================================================================= +# Action utilities +# ============================================================================= + + +def noop_streaming_callback(_chunk: Any) -> None: # noqa: ANN401 + pass + + +def get_func_description(func: Callable[..., Any], description: str | None = None) -> str: + """Get description from explicit param or function docstring.""" + if description is not None: + return description + return func.__doc__ or '' + + +def parse_plugin_name_from_action_name(name: str) -> str | None: + """Extract plugin namespace from 'plugin/action' format.""" + tokens = name.split('/') + if len(tokens) > 1: + return tokens[0] + return None + + +def extract_action_args_and_types( + input_spec: inspect.FullArgSpec, + annotations: Mapping[str, Any] | None = None, +) -> tuple[list[str], list[Any]]: + """Extract argument names and types from a function spec.""" + arg_types = [] + action_args = input_spec.args.copy() + resolved_annotations = annotations or input_spec.annotations + + # Special case when using a method as an action, we ignore first "self" + # arg. (Note: The original condition `len(action_args) <= 3` is preserved + # from the source snippet). + if len(action_args) > 0 and len(action_args) <= 3 and action_args[0] == 'self': + del action_args[0] + + for arg in action_args: + arg_types.append(resolved_annotations.get(arg, Any)) + + return action_args, arg_types + + +# ============================================================================= +# Action key utilities +# ============================================================================= + + +def parse_action_key(key: str) -> tuple[ActionKind, str]: + """Parse '//' key into (ActionKind, name).""" + tokens = key.split('/') + if len(tokens) < 3 or not tokens[1] or not tokens[2]: + msg = f'Invalid action key format: `{key}`.Expected format: `//`' + raise ValueError(msg) + + kind_str = tokens[1] + name = '/'.join(tokens[2:]) + try: + kind = ActionKind(kind_str) + except ValueError as e: + msg = f'Invalid action kind: `{kind_str}`' + raise ValueError(msg) from e + # pyrefly: ignore[bad-return] - ActionKind is StrEnum subclass, pyrefly doesn't narrow properly + return kind, name + + +def create_action_key(kind: ActionKind, name: str) -> str: + """Create '//' key.""" + return f'/{kind}/{name}' + + +# ============================================================================= +# Action core +# ============================================================================= + +InputT = TypeVar('InputT', default=Any) +OutputT = TypeVar('OutputT', default=Any) +ChunkT = TypeVar('ChunkT', default=Never) + +# Generic streaming callback - use Callable[[ChunkT], None] for typed chunks +# This untyped version is for internal use where chunk type is unknown +StreamingCallback = Callable[[object], None] + +_action_context: ContextVar[dict[str, object] | None] = ContextVar('context') +_ = _action_context.set(None) + + +class ActionRunContext: + """Execution context for an action. + + Provides read-only access to action context (auth, metadata) and streaming support. + """ + + def __init__( + self, + context: dict[str, object] | None = None, + streaming_callback: StreamingCallback | None = None, + ) -> None: + self._context: dict[str, object] = context if context is not None else {} + self._streaming_callback = streaming_callback + + @property + def context(self) -> dict[str, object]: + return self._context + + @property + def is_streaming(self) -> bool: + """Returns True if a streaming callback is registered.""" + return self._streaming_callback is not None + + @property + def streaming_callback(self) -> StreamingCallback | None: + """Returns the streaming callback, if any. + + Use this when you need to pass the callback to another action. + For sending chunks directly, use send_chunk() instead. + """ + return self._streaming_callback + + def send_chunk(self, chunk: object) -> None: + """Send a streaming chunk to the client. + + Args: + chunk: The chunk data to stream. + """ + if self._streaming_callback is not None: + self._streaming_callback(chunk) + + @staticmethod + def _current_context() -> dict[str, object] | None: + return _action_context.get(None) + + +class Action(Generic[InputT, OutputT, ChunkT]): + """A named, traced, remotely callable function.""" + + def __init__( + self, + kind: ActionKind, + name: str, + fn: Callable[..., Awaitable[OutputT]], + metadata_fn: Callable[..., object] | None = None, + description: str | None = None, + metadata: dict[str, object] | None = None, + span_metadata: dict[str, SpanAttributeValue] | None = None, + ) -> None: + self._kind: ActionKind = kind + self._name: str = name + self._metadata: dict[str, object] = metadata if metadata else {} + self._description: str | None = description + # Optional matcher function for resource actions + self.matches: Callable[[object], bool] | None = None + + # All action handlers must be async + if not inspect.iscoroutinefunction(fn): + raise TypeError(f"Action handlers must be async functions. Got sync function for '{name}'.") + + input_spec = inspect.getfullargspec(metadata_fn if metadata_fn else fn) + try: + resolved_annotations = get_type_hints(metadata_fn if metadata_fn else fn) + except (NameError, TypeError, AttributeError): + resolved_annotations = input_spec.annotations + action_args, arg_types = extract_action_args_and_types(input_spec, resolved_annotations) + n_action_args = len(action_args) + self._fn = _make_tracing_wrapper(name, kind, span_metadata or {}, n_action_args, fn) + self._initialize_io_schemas(action_args, arg_types, resolved_annotations, input_spec) + + @property + def kind(self) -> ActionKind: + return self._kind + + @property + def name(self) -> str: + return self._name + + @property + def description(self) -> str | None: + return self._description + + @property + def metadata(self) -> dict[str, object]: + return self._metadata + + @property + def input_type(self) -> TypeAdapter[InputT] | None: + return self._input_type + + @property + def input_schema(self) -> dict[str, object]: + return self._input_schema + + @input_schema.setter + def input_schema(self, value: dict[str, object]) -> None: + self._input_schema = value + self._metadata[ActionMetadataKey.INPUT_KEY] = value + + @property + def output_schema(self) -> dict[str, object]: + return self._output_schema + + @output_schema.setter + def output_schema(self, value: dict[str, object]) -> None: + self._output_schema = value + self._metadata[ActionMetadataKey.OUTPUT_KEY] = value + + async def __call__(self, input: InputT | None = None) -> OutputT: + """Call the action directly, returning just the response value.""" + return (await self.run(input)).response + + async def run( + self, + input: InputT | None = None, + on_chunk: Callable[[ChunkT], None] | None = None, + context: dict[str, object] | None = None, + on_trace_start: Callable[[str, str], None] | None = None, + telemetry_labels: dict[str, object] | None = None, + ) -> ActionResponse[OutputT]: + """Execute the action with optional input validation. + + Args: + input: The input to the action. Will be validated against the input schema. + on_chunk: Optional streaming callback for chunked responses. + context: Optional context dict for the action. + on_trace_start: Optional callback invoked when trace starts. + telemetry_labels: Custom labels to set as direct span attributes. + + Returns: + ActionResponse containing the result and trace metadata. + + Raises: + GenkitError: If input validation fails (INVALID_ARGUMENT status). + """ + # Validate input if we have a schema + if self._input_type is not None: + try: + input = self._input_type.validate_python(input) + except ValidationError as e: + if input is None: + raise GenkitError( + message=( + f"Action '{self.name}' requires input but none was provided. " + 'Please supply a valid input payload.' + ), + status='INVALID_ARGUMENT', + ) from e + raise GenkitError( + message=f"Invalid input for action '{self.name}': {e}", + status='INVALID_ARGUMENT', + cause=e, + ) from e + + if context: + _ = _action_context.set(context) + + streaming_cb = cast(StreamingCallback, on_chunk) if on_chunk else None + + return await self._fn( + input, + ActionRunContext( + context=_action_context.get(None), + streaming_callback=streaming_cb, + ), + streaming_cb, + on_trace_start, + telemetry_labels, + ) + + def stream( + self, + input: InputT | None = None, + context: dict[str, object] | None = None, + telemetry_labels: dict[str, object] | None = None, + timeout: float | None = None, + ) -> StreamResponse[ChunkT, OutputT]: + """Execute and return a StreamResponse with .stream and .response properties.""" + channel: Channel[ChunkT, ActionResponse[OutputT]] = Channel(timeout=timeout) + + def send_chunk(c: ChunkT) -> None: + channel.send(c) + + resp = self.run( + input=input, + context=context, + telemetry_labels=telemetry_labels, + on_chunk=send_chunk, + ) + channel.set_close_future(asyncio.create_task(resp)) + + result_future: asyncio.Future[OutputT] = asyncio.Future() + channel.closed.add_done_callback(lambda _: result_future.set_result(channel.closed.result().response)) + + return StreamResponse(stream=channel, response=result_future) + + def _initialize_io_schemas( + self, + action_args: list[str], + arg_types: list[type], + annotations: dict[str, Any], + _input_spec: inspect.FullArgSpec, + ) -> None: + # Allow up to 2 args: (input, ctx) - use ctx.send_chunk() for streaming + if len(action_args) > 2: + raise TypeError(f'can only have up to 2 args: {action_args}') + + if len(action_args) > 0: + type_adapter = TypeAdapter(arg_types[0]) + self._input_schema: dict[str, object] = type_adapter.json_schema() + self._input_type: TypeAdapter[InputT] | None = cast(TypeAdapter[InputT], type_adapter) + self._metadata[ActionMetadataKey.INPUT_KEY] = self._input_schema + else: + self._input_schema = TypeAdapter(object).json_schema() + self._input_type = None + self._metadata[ActionMetadataKey.INPUT_KEY] = self._input_schema + + if ActionMetadataKey.RETURN in annotations: + type_adapter = TypeAdapter(annotations[ActionMetadataKey.RETURN]) + self._output_schema: dict[str, object] = type_adapter.json_schema() + self._metadata[ActionMetadataKey.OUTPUT_KEY] = self._output_schema + else: + self._output_schema = TypeAdapter(object).json_schema() + self._metadata[ActionMetadataKey.OUTPUT_KEY] = self._output_schema + + +class ActionMetadata(BaseModel): + """Action metadata for registry and reflection.""" + + kind: ActionKind + name: str + description: str | None = None + input_schema: object | None = None + input_json_schema: dict[str, object] | None = None + output_schema: object | None = None + output_json_schema: dict[str, object] | None = None + stream_schema: object | None = None + metadata: dict[str, object] | None = None + + +def _make_tracing_wrapper( + name: str, + kind: ActionKind, + span_metadata: dict[str, SpanAttributeValue], + n_action_args: int, + fn: Callable[..., Awaitable[Any]], +) -> Callable[ + [ + object | None, + ActionRunContext, + StreamingCallback | None, + Callable[[str, str], None] | None, + dict[str, object] | None, + ], + Awaitable[ActionResponse[Any]], +]: + """Create a tracing wrapper for an async action function.""" + + def _record_latency(output: object, start_time: float) -> object: + latency_ms = (time.perf_counter() - start_time) * 1000 + if hasattr(output, 'latency_ms'): + try: + cast(Any, output).latency_ms = latency_ms + except (TypeError, ValidationError, AttributeError): + # If immutable (e.g. Pydantic model with frozen=True), try model_copy + if hasattr(output, 'model_copy'): + output = cast(Any, output).model_copy(update={'latency_ms': latency_ms}) + return output + + async def tracing_wrapper( + input: object | None, + ctx: ActionRunContext, + on_chunk: StreamingCallback | None, + on_trace_start: Callable[[str, str], None] | None, + telemetry_labels: dict[str, object] | None, + ) -> ActionResponse[Any]: + start_time = time.perf_counter() + + with _save_parent_path(): + with tracer.start_as_current_span(name) as span: + # Format trace_id and span_id as hex strings (OpenTelemetry standard format) + trace_id = format(span.get_span_context().trace_id, '032x') + span_id = format(span.get_span_context().span_id, '016x') + if on_trace_start: + on_trace_start(trace_id, span_id) + + # Set telemetry labels as direct span attributes (matches JS/Go behavior) + if telemetry_labels: + for key, value in telemetry_labels.items(): + span.set_attribute(key, str(value)) + + _record_input_metadata( + span=span, + kind=kind, + name=name, + span_metadata=span_metadata, + input=input, + ) + + try: + match n_action_args: + case 0: + output = await fn() + case 1: + output = await fn(input) + case 2: + output = await fn(input, ctx) + case _: + raise ValueError('action fn must have 0-2 args') + except Exception as e: + # Re-raise existing GenkitError instances to avoid double-wrapping + if isinstance(e, GenkitError): + raise + raise GenkitError( + cause=e, + message=f'Error while running action {name}', + trace_id=trace_id, + ) from e + + output = _record_latency(output, start_time) + _record_output_metadata(span, output=output) + return ActionResponse(response=output, trace_id=trace_id, span_id=span_id) + + return tracing_wrapper diff --git a/py/packages/genkit/src/genkit/blocks/background_model.py b/py/packages/genkit/src/genkit/_core/_background.py similarity index 60% rename from py/packages/genkit/src/genkit/blocks/background_model.py rename to py/packages/genkit/src/genkit/_core/_background.py index dce288a618..aacfb3bf8a 100644 --- a/py/packages/genkit/src/genkit/blocks/background_model.py +++ b/py/packages/genkit/src/genkit/_core/_background.py @@ -14,109 +14,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Background model definitions for the Genkit framework. - -Background models are long-running AI operations that don't complete immediately. -They are used for tasks like video generation (Veo) or large image generation (Imagen) -that may take seconds or minutes to complete. - -Why Background Models? - Regular models return results synchronously - you call generate() and get a - response. But some AI tasks (video generation, complex rendering) can take - minutes. Background models solve this by: - - 1. Returning immediately with an operation ID - 2. Allowing you to poll for completion - 3. Optionally supporting cancellation - -Architecture: - A background model consists of three actions registered together. - The naming convention matches the JS implementation: - - ``` - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ BackgroundAction β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ start_action β”‚ β”‚ check_action β”‚ β”‚ cancel_action β”‚ β”‚ - β”‚ β”‚ /background- β”‚ β”‚ /check-operation β”‚ β”‚ /cancel-operationβ”‚ β”‚ - β”‚ β”‚ model/{name} β”‚ β”‚ /{name}/check β”‚ β”‚ /{name}/cancel β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ β”‚ (optional) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ β”‚ β”‚ - β–Ό β–Ό β–Ό - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ StartModelOpFn β”‚ β”‚ CheckModelOpFn β”‚ β”‚ CancelModelOpFn β”‚ - β”‚ (user-provided) β”‚ β”‚ (user-provided) β”‚ β”‚ (user-provided) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - ``` - -Key Concepts: - +---------------+----------------------------------------------------------------+ - | Term | Description | - +---------------+----------------------------------------------------------------+ - | Operation | A long-running task with an ID, status, and eventual result | - | start() | Initiates the background operation, returns an Operation | - | check() | Polls the operation status, returns updated Operation | - | cancel() | Attempts to cancel a running operation (if supported) | - +---------------+----------------------------------------------------------------+ - -Workflow: - ``` - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Client β”‚ β”‚ Background β”‚ β”‚ Backend β”‚ - β”‚ Code β”‚ β”‚ Model β”‚ β”‚ Service β”‚ - β””β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”˜ - β”‚ β”‚ β”‚ - β”‚ 1. start(request) β”‚ β”‚ - │────────────────────►│ submit job β”‚ - β”‚ │──────────────────────►│ - β”‚ β”‚ job_id β”‚ - β”‚ Operation(id=X) │◄──────────────────────│ - │◄────────────────────│ β”‚ - β”‚ β”‚ β”‚ - β”‚ 2. check(op) β”‚ β”‚ - │────────────────────►│ get status(X) β”‚ - β”‚ │──────────────────────►│ - β”‚ β”‚ status: processing β”‚ - β”‚ Operation(done=F) │◄──────────────────────│ - │◄────────────────────│ β”‚ - β”‚ β”‚ β”‚ - β”‚ ... (poll loop) β”‚ β”‚ - β”‚ β”‚ β”‚ - β”‚ 3. check(op) β”‚ β”‚ - │────────────────────►│ get status(X) β”‚ - β”‚ │──────────────────────►│ - β”‚ β”‚ status: complete β”‚ - β”‚ Operation(done=T, │◄──────────────────────│ - β”‚ output=result) β”‚ β”‚ - │◄────────────────────│ β”‚ - β”‚ β”‚ β”‚ - ``` - -Example: - >>> # Define a background model for video generation - >>> async def start_video(request: GenerateRequest, ctx) -> Operation: - ... job_id = await video_api.submit(request.messages[0].content[0].text) - ... return Operation(id=job_id, done=False) - >>> async def check_video(op: Operation, ctx) -> Operation: - ... status = await video_api.get_status(op.id) - ... if status.complete: - ... return Operation(id=op.id, done=True, output={...}) - ... return Operation(id=op.id, done=False) - >>> ai.define_background_model( - ... name='my-video-model', - ... start=start_video, - ... check=check_video, - ... ) - -See Also: - - JS implementation: js/core/src/background-action.ts - - JS model wrapper: js/ai/src/model.ts (defineBackgroundModel) - - Sample: py/samples/background-model-demo/ -""" +"""Background model definitions for the Genkit framework.""" from __future__ import annotations @@ -126,14 +24,11 @@ from pydantic import BaseModel -from genkit.codec import dump_dict -from genkit.core.action import Action, ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry -from genkit.core.schema import to_json_schema -from genkit.core.typing import ( - GenerateRequest, - GenerateResponse, +from genkit._core._action import Action, ActionKind, ActionRunContext +from genkit._core._model import ModelRequest, ModelResponse +from genkit._core._registry import Registry +from genkit._core._schema import to_json_schema +from genkit._core._typing import ( ModelInfo, Operation, ) @@ -157,7 +52,7 @@ def _make_action_key(action_type: ActionKind | str, name: str) -> str: # Type aliases for background model functions matching JS signatures # JS: start: (input, options) => Promise> -StartModelOpFn = Callable[[GenerateRequest, ActionRunContext], Awaitable[Operation]] +StartModelOpFn = Callable[[ModelRequest, ActionRunContext], Awaitable[Operation]] # JS: check: (input: Operation) => Promise> CheckModelOpFn = Callable[[Operation], Awaitable[Operation]] # JS: cancel?: (input: Operation) => Promise> @@ -218,7 +113,7 @@ def supports_cancel(self) -> bool: async def start( self, - input: GenerateRequest | None = None, + input: ModelRequest | None = None, options: dict[str, Any] | None = None, ) -> Operation: """Start a background operation. @@ -232,7 +127,7 @@ async def start( Returns: An Operation with an ID to track the job. """ - result = await self.start_action.arun(input) + result = await self.start_action.run(input) return _ensure_operation(result.response) async def check(self, operation: Operation) -> Operation: @@ -246,7 +141,7 @@ async def check(self, operation: Operation) -> Operation: Returns: Updated Operation with current status. """ - result = await self.check_action.arun(operation) + result = await self.check_action.run(operation) return _ensure_operation(result.response) async def cancel(self, operation: Operation) -> Operation: @@ -266,7 +161,7 @@ async def cancel(self, operation: Operation) -> Operation: if self.cancel_action is None: # Match JS behavior: return operation unchanged if cancel not supported return operation - result = await self.cancel_action.arun(operation) + result = await self.cancel_action.run(operation) return _ensure_operation(result.response) @@ -310,7 +205,7 @@ def define_background_model( config_schema: type | dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, description: str | None = None, -) -> BackgroundAction[GenerateResponse]: +) -> BackgroundAction[ModelResponse]: """Define and register a background model. This matches the JS defineBackgroundModel function from js/ai/src/model.ts. @@ -355,9 +250,7 @@ def define_background_model( model_options: dict[str, Any] = {} if info: - info_dict = dump_dict(info) - if isinstance(info_dict, dict): - model_options.update(info_dict) # type: ignore[arg-type] + model_options.update(info.model_dump()) model_options['label'] = label if config_schema: @@ -366,11 +259,11 @@ def define_background_model( model_meta['model'] = model_options # Build output schema metadata (matching JS) - output_schema_meta = to_json_schema(GenerateResponse) + output_schema_meta = to_json_schema(ModelResponse) model_meta['outputSchema'] = output_schema_meta # Wrap the start function to add the action key and timing (matching JS) - async def wrapped_start(request: GenerateRequest, ctx: ActionRunContext) -> Operation: + async def wrapped_start(request: ModelRequest, ctx: ActionRunContext) -> Operation: start_time = time.perf_counter() op = await start(request, ctx) # Set action key matching JS format: /{actionType}/{name} @@ -441,7 +334,7 @@ async def wrapped_cancel(op: Operation, ctx: ActionRunContext) -> Operation: async def lookup_background_action( registry: Registry, key: str, -) -> BackgroundAction[GenerateResponse] | None: +) -> BackgroundAction[ModelResponse] | None: """Look up a background action by its action key. Matches JS lookupBackgroundAction from js/core/src/background-action.ts. diff --git a/py/packages/genkit/src/genkit/_core/_base.py b/py/packages/genkit/src/genkit/_core/_base.py new file mode 100644 index 0000000000..d4436a463b --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_base.py @@ -0,0 +1,65 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Base model with correct serialization defaults for Genkit types.""" + +from __future__ import annotations + +import base64 +from typing import Any, ClassVar + +from pydantic import BaseModel, ConfigDict +from pydantic.alias_generators import to_camel + + +def _default_serializer(obj: object) -> object: + """Default serializer for objects not handled by json.dumps.""" + if isinstance(obj, bytes): + try: + return base64.b64encode(obj).decode('utf-8') + except Exception: + return '' + return str(obj) + + +class GenkitModel(BaseModel): + """Base model with correct serialization defaults. + + All Genkit types inherit from this to ensure consistent serialization: + - by_alias=True: Use camelCase field names (matching JS SDK) + - exclude_none=True: Omit null fields (cleaner JSON) + - fallback=_default_serializer: Handle bytes and other edge cases + """ + + model_config: ClassVar[ConfigDict] = ConfigDict( + alias_generator=to_camel, + extra='forbid', + populate_by_name=True, + ) + + def model_dump(self, **kwargs: Any) -> dict[str, Any]: + """Dump model with Genkit defaults (by_alias=True, exclude_none=True).""" + kwargs.setdefault('by_alias', True) + kwargs.setdefault('exclude_none', True) + kwargs.setdefault('fallback', _default_serializer) + return super().model_dump(**kwargs) + + def model_dump_json(self, **kwargs: Any) -> str: + """Dump model to JSON with Genkit defaults.""" + kwargs.setdefault('by_alias', True) + kwargs.setdefault('exclude_none', True) + kwargs.setdefault('fallback', _default_serializer) + return super().model_dump_json(**kwargs) diff --git a/py/packages/genkit/src/genkit/_core/_channel.py b/py/packages/genkit/src/genkit/_core/_channel.py new file mode 100644 index 0000000000..7f3665f813 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_channel.py @@ -0,0 +1,117 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Channel for async streaming with final value, and uvloop-aware runner.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Coroutine +from typing import Any, Generic, TypeVar + +from typing_extensions import TypeVar as TypeVarExt + +from genkit._core._logger import get_logger + +from ._compat import wait_for + +logger = get_logger(__name__) + +T = TypeVar('T') +T_co = TypeVarExt('T_co') +R = TypeVarExt('R', default=Any) + + +class Channel(Generic[T_co, R]): + """Async channel for streaming values with a final result when closed.""" + + def __init__(self, timeout: float | int | None = None) -> None: + if timeout is not None and timeout < 0: + raise ValueError('Timeout must be non-negative') + self.queue: asyncio.Queue[T_co] = asyncio.Queue() + self.closed: asyncio.Future[R] = asyncio.Future() + self._close_future: asyncio.Future[R] | None = None + self._timeout = timeout + + def __aiter__(self) -> AsyncIterator[T_co]: + return self + + async def __anext__(self) -> T_co: + if not self.queue.empty(): + return self.queue.get_nowait() + + pop_task = asyncio.ensure_future(self._pop()) + if not self._close_future: + return await wait_for(pop_task, timeout=self._timeout) + + finished, _ = await asyncio.wait( + [pop_task, self._close_future], + return_when=asyncio.FIRST_COMPLETED, + timeout=self._timeout, + ) + + if not finished: + _ = pop_task.cancel() + raise TimeoutError('Channel timeout exceeded') + + if pop_task in finished: + return pop_task.result() + + if self._close_future in finished: + _ = pop_task.cancel() + raise StopAsyncIteration + + return await wait_for(pop_task, timeout=self._timeout) + + def send(self, value: T_co) -> None: + """Send a value into the channel.""" + self.queue.put_nowait(value) + + def set_close_future(self, future: asyncio.Future[R]) -> None: + """Set a future that closes the channel when completed.""" + if future is None: # pyright: ignore[reportUnnecessaryComparison] + raise ValueError('Cannot set a None future') # pyright: ignore[reportUnreachable] + + def _handle_done(v: asyncio.Future[R]) -> None: + if v.cancelled(): + _ = self.closed.cancel() + elif (exc := v.exception()) is not None: + self.closed.set_exception(exc) + else: + self.closed.set_result(v.result()) + + self._close_future = asyncio.ensure_future(future) + if self._close_future is not None: # pyright: ignore[reportUnnecessaryComparison] + self._close_future.add_done_callback(_handle_done) + + async def _pop(self) -> T_co: + r = await self.queue.get() + self.queue.task_done() + if r is None: + raise StopAsyncIteration + return r + + +def run_loop(coro: Coroutine[object, object, T], *, debug: bool | None = None) -> T: + """Run a coroutine using uvloop if available, otherwise asyncio.""" + try: + import uvloop # noqa: PLC0415 + + logger.debug('Using uvloop (recommended)') + return uvloop.run(coro, debug=debug) + except ImportError as e: + logger.debug('Using asyncio (install uvloop for better performance)', error=e) + return asyncio.run(coro, debug=debug) diff --git a/py/packages/genkit/src/genkit/aio/_compat.py b/py/packages/genkit/src/genkit/_core/_compat.py similarity index 55% rename from py/packages/genkit/src/genkit/aio/_compat.py rename to py/packages/genkit/src/genkit/_core/_compat.py index 338c1f973e..64dc052002 100644 --- a/py/packages/genkit/src/genkit/aio/_compat.py +++ b/py/packages/genkit/src/genkit/_core/_compat.py @@ -14,16 +14,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Compatibility layer for asyncio. - -This module provides a compatibility layer for asyncio. - -The asyncio.wait_for function was changed in Python 3.11 to raise a TimeoutError -instead of an asyncio.TimeoutError. This module provides a compatibility layer -for this change among others. - -See: https://docs.python.org/3/library/asyncio-task.html#asyncio.wait_for -""" +"""Compatibility layer for asyncio.""" import asyncio import sys @@ -31,25 +22,18 @@ T = TypeVar('T') +# StrEnum - use strenum package for cross-version compatibility +# Note: StrEnum was added to stdlib in Python 3.11, but we use strenum for 3.10 compat +# override decorator - use typing_extensions for consistency across Python versions +# Note: override was added to typing in Python 3.12, but typing_extensions has it for all versions +from typing import overload as overload # noqa: E402 -async def wait_for_310(fut: asyncio.Future[T], timeout: float | None = None) -> T: - """Wait for a future to complete. - - This is a compatibility layer for asyncio.wait_for that raises a TimeoutError - instead of an asyncio.TimeoutError. - - This is necessary because the behavior of asyncio.wait_for changed in Python - 3.11. +from strenum import StrEnum as StrEnum # noqa: E402 +from typing_extensions import override as override # noqa: E402 - See: https://docs.python.org/3/library/asyncio-task.html#asyncio.wait_for - Args: - fut: The future to wait for. - timeout: The timeout in seconds. - - Returns: - The result of the future. - """ +async def wait_for_310(fut: asyncio.Future[T], timeout: float | None = None) -> T: + """Python 3.10 compat: raises TimeoutError instead of asyncio.TimeoutError.""" try: return await asyncio.wait_for(fut, timeout) except asyncio.TimeoutError as e: diff --git a/py/packages/genkit/src/genkit/core/constants.py b/py/packages/genkit/src/genkit/_core/_constants.py similarity index 84% rename from py/packages/genkit/src/genkit/core/constants.py rename to py/packages/genkit/src/genkit/_core/_constants.py index 77cbb4c17d..bc103645cc 100644 --- a/py/packages/genkit/src/genkit/core/constants.py +++ b/py/packages/genkit/src/genkit/_core/_constants.py @@ -17,9 +17,7 @@ """Module containing various core constants.""" # The version of Genkit sent over HTTP in the headers. -DEFAULT_GENKIT_VERSION = '0.3.2' - # TODO(#4349): make this dynamic -GENKIT_VERSION = DEFAULT_GENKIT_VERSION +GENKIT_VERSION = '0.3.2' -GENKIT_CLIENT_HEADER = f'genkit-python/{DEFAULT_GENKIT_VERSION}' +GENKIT_CLIENT_HEADER = f'genkit-python/{GENKIT_VERSION}' diff --git a/py/packages/genkit/src/genkit/core/context.py b/py/packages/genkit/src/genkit/_core/_context.py similarity index 95% rename from py/packages/genkit/src/genkit/core/context.py rename to py/packages/genkit/src/genkit/_core/_context.py index add64e457c..dc629cd3bd 100644 --- a/py/packages/genkit/src/genkit/core/context.py +++ b/py/packages/genkit/src/genkit/_core/_context.py @@ -50,7 +50,7 @@ class RequestData(Generic[T]): Action. If middleware throws an error, that error will fail the request and the Action will not be called. -Expected cases should return a UserFacingError, which allows the request handler to +Expected cases should return a PublicError, which allows the request handler to know what data is safe to return to end users. Middleware can provide validation in addition to parsing. For example, an auth middleware can have policies for validating auth in addition to passing auth context to the Action. diff --git a/py/packages/genkit/src/genkit/_core/_dap.py b/py/packages/genkit/src/genkit/_core/_dap.py new file mode 100644 index 0000000000..5d509786e5 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_dap.py @@ -0,0 +1,154 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Dynamic Action Provider (DAP) support for Genkit.""" + +import asyncio +import time +from collections.abc import Awaitable, Callable, Mapping +from typing import Any + +from genkit._core._action import Action, ActionKind +from genkit._core._registry import Registry + +ActionMetadataLike = Mapping[str, object] +DapValue = dict[str, list[Action[Any, Any]]] +DapFn = Callable[[], Awaitable[DapValue]] +DapMetadata = dict[str, list[ActionMetadataLike]] + +# Default cache TTL in milliseconds +_DEFAULT_CACHE_TTL_MS = 3000 + + +class DynamicActionProvider: + """Lazily resolves actions from an external source with TTL caching.""" + + def __init__( + self, + action: Action[Any, Any], + dap_fn: DapFn, + cache_ttl_millis: int | None = None, + ) -> None: + self.action = action + self._dap_fn = dap_fn + self._value: DapValue | None = None + self._expires_at: float | None = None + self._fetch_task: asyncio.Task[DapValue] | None = None + self._ttl_millis = ( + _DEFAULT_CACHE_TTL_MS if cache_ttl_millis is None or cache_ttl_millis == 0 else cache_ttl_millis + ) + + def invalidate_cache(self) -> None: + self._value = None + self._expires_at = None + + async def _get_or_fetch(self, skip_trace: bool = False) -> DapValue: + """Get cached value or fetch fresh data, coalescing concurrent fetches.""" + is_stale = ( + self._value is None + or self._expires_at is None + or self._ttl_millis < 0 + or time.time() * 1000 > self._expires_at + ) + if not is_stale and self._value is not None: + return self._value + + if self._fetch_task is not None: + return await self._fetch_task + + self._fetch_task = asyncio.create_task(self._do_fetch(skip_trace)) + try: + return await self._fetch_task + finally: + self._fetch_task = None + + async def _do_fetch(self, skip_trace: bool) -> DapValue: + try: + self._value = await self._dap_fn() + self._expires_at = time.time() * 1000 + self._ttl_millis + if not skip_trace: + metadata = {k: [a.metadata or {} for a in v] for k, v in self._value.items()} + await self.action.run(metadata) + return self._value + except Exception: + self.invalidate_cache() + raise + + async def get_action(self, action_type: str, action_name: str) -> Action[Any, Any] | None: + result = await self._get_or_fetch() + for action in result.get(action_type, []): + if action.name == action_name: + return action + return None + + async def list_action_metadata(self, action_type: str, action_name: str) -> list[ActionMetadataLike]: + """List metadata matching pattern: '*'=all, 'prefix*'=prefix match, else exact.""" + result = await self._get_or_fetch() + actions = result.get(action_type, []) + if not actions: + return [] + + metadata_list: list[ActionMetadataLike] = [action.metadata or {} for action in actions] + + if action_name == '*': + return metadata_list + if action_name.endswith('*'): + prefix = action_name[:-1] + return [m for m in metadata_list if str(m.get('name', '')).startswith(prefix)] + return [m for m in metadata_list if m.get('name') == action_name] + + async def get_action_metadata_record(self, dap_prefix: str) -> dict[str, ActionMetadataLike]: + """Get all actions as metadata record for reflection API.""" + result = await self._get_or_fetch(skip_trace=True) + dap_actions: dict[str, ActionMetadataLike] = {} + for action_type, actions in result.items(): + for action in actions: + if not action.name: + raise ValueError(f'Invalid metadata from {dap_prefix} - name required') + dap_actions[f'{dap_prefix}:{action_type}/{action.name}'] = action.metadata or {} + return dap_actions + + +def is_dynamic_action_provider(obj: object) -> bool: + if isinstance(obj, DynamicActionProvider): + return True + metadata = getattr(obj, 'metadata', None) + return isinstance(metadata, dict) and metadata.get('type') == 'dynamic-action-provider' + + +def define_dynamic_action_provider( + registry: Registry, + name: str, + fn: DapFn, + *, + description: str | None = None, + cache_ttl_millis: int | None = None, + metadata: dict[str, Any] | None = None, +) -> DynamicActionProvider: + """Define and register a Dynamic Action Provider for lazy action resolution.""" + + async def dap_action(input: DapMetadata) -> DapMetadata: + return input + + action = registry.register_action( + name=name, + kind=ActionKind.DYNAMIC_ACTION_PROVIDER, + description=description, + fn=dap_action, + metadata={**(metadata or {}), 'type': 'dynamic-action-provider'}, + ) + + return DynamicActionProvider(action, fn, cache_ttl_millis) diff --git a/py/packages/genkit/src/genkit/_core/_environment.py b/py/packages/genkit/src/genkit/_core/_environment.py new file mode 100644 index 0000000000..8b6b621f02 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_environment.py @@ -0,0 +1,44 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Environment detection for Genkit runtime.""" + +import os + +from genkit._core._compat import StrEnum + +# Environment variable name +GENKIT_ENV = 'GENKIT_ENV' + + +class GenkitEnvironment(StrEnum): + """Genkit runtime environments.""" + + DEV = 'dev' + PROD = 'prod' + + +def is_dev_environment() -> bool: + """Check if running in development mode (GENKIT_ENV=dev).""" + return os.getenv(GENKIT_ENV) == GenkitEnvironment.DEV + + +def get_current_environment() -> GenkitEnvironment: + """Get current environment, defaults to PROD.""" + env = os.getenv(GENKIT_ENV) + if env == GenkitEnvironment.DEV: + return GenkitEnvironment.DEV + return GenkitEnvironment.PROD diff --git a/py/packages/genkit/src/genkit/core/error.py b/py/packages/genkit/src/genkit/_core/_error.py similarity index 50% rename from py/packages/genkit/src/genkit/core/error.py rename to py/packages/genkit/src/genkit/_core/_error.py index de3c1aa5cc..e383f08fd6 100644 --- a/py/packages/genkit/src/genkit/core/error.py +++ b/py/packages/genkit/src/genkit/_core/_error.py @@ -14,95 +14,109 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Error classes and utilities for the Genkit framework. - -This module defines the error hierarchy and utilities for handling errors -in Genkit applications. It provides structured error types with status codes, -trace IDs, and serialization for HTTP responses. - -Overview: - Genkit uses a structured error system based on gRPC-style status codes. - The base ``GenkitError`` class provides rich error context including - status codes, trace IDs, and stack traces for debugging. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Error Class Hierarchy β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ Exception β”‚ - β”‚ β”‚ β”‚ - β”‚ └── GenkitError β”‚ - β”‚ β”‚ β”‚ - β”‚ β”œβ”€β”€ UserFacingError (safe to return to users) β”‚ - β”‚ β”‚ β”‚ - β”‚ └── UnstableApiError (beta/alpha API misuse) β”‚ - β”‚ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Terminology: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Term β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ GenkitError β”‚ Base error with status, trace_id, and details β”‚ - β”‚ UserFacingError β”‚ Error safe to return in HTTP responses β”‚ - β”‚ StatusName β”‚ gRPC status name (e.g., 'NOT_FOUND', 'INTERNAL') β”‚ - β”‚ StatusCodes β”‚ Enum mapping status names to numeric codes β”‚ - β”‚ http_code β”‚ HTTP status code derived from StatusName β”‚ - β”‚ trace_id β”‚ Unique ID linking error to trace spans β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Functions: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Function β”‚ Purpose β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ get_http_status() β”‚ Get HTTP status code from any error β”‚ - β”‚ get_callable_json() β”‚ Serialize error for callable HTTP responses β”‚ - β”‚ get_error_message() β”‚ Extract message string from any error β”‚ - β”‚ get_error_stack() β”‚ Extract stack trace from an exception β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Raising and handling errors: - - ```python - from genkit.core.error import GenkitError, UserFacingError, get_http_status - - # Raise a structured error - raise GenkitError( - message='Model not found', - status='NOT_FOUND', - trace_id='abc123', - ) +"""Error classes and utilities for the Genkit framework.""" - # User-facing error (safe to return in HTTP response) - raise UserFacingError( - status='INVALID_ARGUMENT', - message='Invalid prompt: too long', - ) +from enum import IntEnum +from typing import Any, ClassVar, Literal - # Get HTTP status for any error - try: - await ai.generate(...) - except Exception as e: - status_code = get_http_status(e) # 404 for NOT_FOUND, 500 otherwise - ``` +from pydantic import BaseModel, ConfigDict, Field +from pydantic.alias_generators import to_camel -Caveats: - - Only ``UserFacingError`` messages are safe to return to end users - - Other ``GenkitError`` messages may contain internal details - - Use ``get_callable_json()`` for Genkit callable serialization format -See Also: - - gRPC status codes: https://grpc.io/docs/guides/status-codes/ - - genkit.core.status_types: Status code definitions -""" +class StatusCodes(IntEnum): + """gRPC-style status codes. See _STATUS_CODE_MAP for HTTP mappings.""" + + OK = 0 + CANCELLED = 1 + UNKNOWN = 2 + INVALID_ARGUMENT = 3 + DEADLINE_EXCEEDED = 4 + NOT_FOUND = 5 + ALREADY_EXISTS = 6 + PERMISSION_DENIED = 7 + RESOURCE_EXHAUSTED = 8 + FAILED_PRECONDITION = 9 + ABORTED = 10 + OUT_OF_RANGE = 11 + UNIMPLEMENTED = 12 + INTERNAL = 13 + UNAVAILABLE = 14 + DATA_LOSS = 15 + UNAUTHENTICATED = 16 + + +# Type alias for status names +StatusName = Literal[ + 'OK', + 'CANCELLED', + 'UNKNOWN', + 'INVALID_ARGUMENT', + 'DEADLINE_EXCEEDED', + 'NOT_FOUND', + 'ALREADY_EXISTS', + 'PERMISSION_DENIED', + 'UNAUTHENTICATED', + 'RESOURCE_EXHAUSTED', + 'FAILED_PRECONDITION', + 'ABORTED', + 'OUT_OF_RANGE', + 'UNIMPLEMENTED', + 'INTERNAL', + 'UNAVAILABLE', + 'DATA_LOSS', +] + +# Mapping of status names to HTTP status codes +_STATUS_CODE_MAP: dict[StatusName, int] = { + 'OK': 200, + 'CANCELLED': 499, + 'UNKNOWN': 500, + 'INVALID_ARGUMENT': 400, + 'DEADLINE_EXCEEDED': 504, + 'NOT_FOUND': 404, + 'ALREADY_EXISTS': 409, + 'PERMISSION_DENIED': 403, + 'UNAUTHENTICATED': 401, + 'RESOURCE_EXHAUSTED': 429, + 'FAILED_PRECONDITION': 400, + 'ABORTED': 409, + 'OUT_OF_RANGE': 400, + 'UNIMPLEMENTED': 501, + 'INTERNAL': 500, + 'UNAVAILABLE': 503, + 'DATA_LOSS': 500, +} + + +def http_status_code(status: StatusName) -> int: + """Gets the HTTP status code for a given status name. -from typing import Any, ClassVar + Args: + status: The status name to get the HTTP code for. -from pydantic import BaseModel, ConfigDict -from pydantic.alias_generators import to_camel + Returns: + The corresponding HTTP status code. + """ + return _STATUS_CODE_MAP[status] + + +class Status(BaseModel): + """Represents a status with a name and optional message.""" + + model_config: ClassVar[ConfigDict] = ConfigDict( + frozen=True, + validate_assignment=True, + extra='forbid', + populate_by_name=True, + ) + + name: StatusName + message: str = Field(default='') -from genkit.core.status_types import StatusCodes, StatusName, http_status_code + +# ============================================================================= +# Error Classes +# ============================================================================= class ReflectionErrorDetails(BaseModel): @@ -215,25 +229,7 @@ def to_serializable(self) -> ReflectionError: ) -class UnstableApiError(GenkitError): - """Error raised when using unstable APIs from a more stable instance.""" - - def __init__(self, level: str = 'beta', message: str | None = None) -> None: - """Initialize an UnstableApiError. - - Args: - level: The stability level required. - message: Optional message describing which feature is not allowed. - """ - msg_prefix = f'{message} ' if message else '' - super().__init__( - status='FAILED_PRECONDITION', - message=f"{msg_prefix}This API requires '{level}' stability level.\n\n" - + f'To use this feature, initialize Genkit using `from genkit.{level} import genkit`.', - ) - - -class UserFacingError(GenkitError): +class PublicError(GenkitError): """Error class for issues to be returned to users. Using this error allows a web framework handler (e.g. FastAPI, Flask) to know it @@ -243,7 +239,7 @@ class UserFacingError(GenkitError): """ def __init__(self, status: StatusName, message: str, details: Any = None) -> None: # noqa: ANN401 - """Initialize a UserFacingError. + """Initialize a PublicError. Args: status: The status name for this error. @@ -303,20 +299,6 @@ def get_callable_json(error: object) -> HttpErrorWireFormat: ) -def get_error_message(error: object) -> str: - """Extract error message from an error object. - - Args: - error: The error to get the message from. - - Returns: - The error message string. - """ - if isinstance(error, Exception): - return str(error) - return str(error) - - def get_error_stack(error: object) -> str | None: """Extract stack trace from an error object. diff --git a/py/packages/genkit/src/genkit/_core/_extract_json.py b/py/packages/genkit/src/genkit/_core/_extract_json.py new file mode 100644 index 0000000000..c564a9653c --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_extract_json.py @@ -0,0 +1,144 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Utility functions for extracting JSON data from text and markdown.""" + +from dataclasses import dataclass +from typing import Any + +import json5 +from partial_json_parser import loads + +CHAR_NON_BREAKING_SPACE = '\u00a0' + + +def parse_partial_json(json_string: str) -> Any: # noqa: ANN401 + """Parse a partially complete JSON string.""" + return loads(json_string) + + +def extract_json(text: str, throw_on_bad_json: bool = True) -> Any: # noqa: ANN401 + """Extract JSON from text with lenient parsing (handles trailing commas, partial JSON, etc.).""" + if not text.strip(): + return None + + opening_char: str | None = None + closing_char: str | None = None + start_pos: int | None = None + nesting_count = 0 + in_string = False + escape_next = False + + for i in range(len(text)): + char = text[i].replace(CHAR_NON_BREAKING_SPACE, ' ') + + if escape_next: + escape_next = False + continue + + if char == '\\': + escape_next = True + continue + + if char == '"': + in_string = not in_string + continue + + if in_string: + continue + + if not opening_char and char in '{[': + opening_char = char + closing_char = '}' if char == '{' else ']' + start_pos = i + nesting_count += 1 + elif char == opening_char: + nesting_count += 1 + elif char == closing_char: + nesting_count -= 1 + if not nesting_count: + return json5.loads(text[start_pos or 0 : i + 1]) + + # Handle incomplete JSON structure + if start_pos is not None and nesting_count > 0: + try: + return parse_partial_json(text[start_pos:]) + except Exception as e: + if throw_on_bad_json: + raise ValueError(f'Invalid JSON extracted from model output: {text}') from e + return None + + if throw_on_bad_json: + raise ValueError(f'Invalid JSON extracted from model output: {text}') + return None + + +@dataclass +class ExtractItemsResult: + """Result of extracting JSON items from text.""" + + items: list + cursor: int + + +def extract_json_array_from_text(text: str, cursor: int = 0) -> ExtractItemsResult: + """Extract complete JSON objects from the first array found in text.""" + items: list = [] + current_cursor = cursor + + if cursor == 0: + array_start = text.find('[') + if array_start == -1: + return ExtractItemsResult(items=[], cursor=len(text)) + current_cursor = array_start + 1 + + object_start = -1 + brace_count = 0 + in_string = False + escape_next = False + + for i in range(current_cursor, len(text)): + char = text[i] + + if escape_next: + escape_next = False + continue + if char == '\\': + escape_next = True + continue + if char == '"': + in_string = not in_string + continue + if in_string: + continue + + if char == '{': + if brace_count == 0: + object_start = i + brace_count += 1 + elif char == '}': + brace_count -= 1 + if brace_count == 0 and object_start != -1: + try: + items.append(json5.loads(text[object_start : i + 1])) + current_cursor = i + 1 + object_start = -1 + except Exception: # noqa: S110 + pass + elif char == ']' and brace_count == 0: + break + + return ExtractItemsResult(items=items, cursor=current_cursor) diff --git a/py/packages/genkit/src/genkit/_core/_flow.py b/py/packages/genkit/src/genkit/_core/_flow.py new file mode 100644 index 0000000000..2958832caf --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_flow.py @@ -0,0 +1,79 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Flow registration for Genkit.""" + +from __future__ import annotations + +import inspect +from collections.abc import Awaitable, Callable +from typing import Any, overload + +from typing_extensions import TypeVar + +from genkit._core._action import Action, ActionKind, ActionRunContext, get_func_description +from genkit._core._registry import Registry + +InputT = TypeVar('InputT') +OutputT = TypeVar('OutputT') + + +@overload +def define_flow( + registry: Registry, + func: Callable[[], Awaitable[OutputT]], + name: str | None = None, + description: str | None = None, +) -> Action[None, OutputT]: ... + + +@overload +def define_flow( + registry: Registry, + func: Callable[[InputT], Awaitable[OutputT]], + name: str | None = None, + description: str | None = None, +) -> Action[InputT, OutputT]: ... + + +@overload +def define_flow( + registry: Registry, + func: Callable[[InputT, ActionRunContext], Awaitable[OutputT]], + name: str | None = None, + description: str | None = None, +) -> Action[InputT, OutputT]: ... + + +def define_flow( + registry: Registry, + func: Callable[..., Awaitable[Any]], + name: str | None = None, + description: str | None = None, +) -> Action[Any, Any]: + """Register an async function as a flow action.""" + # All Python functions have __name__, but ty is strict about Callable protocol + if not inspect.iscoroutinefunction(func): + raise TypeError(f'Flow must be async: {func.__name__}') # ty: ignore[unresolved-attribute] + + flow_name = name or func.__name__ # ty: ignore[unresolved-attribute] + return registry.register_action( + name=flow_name, + kind=ActionKind.FLOW, + fn=func, + description=get_func_description(func, description), + span_metadata={'genkit:metadata:flow:name': flow_name}, + ) diff --git a/py/packages/genkit/src/genkit/_core/_http_client.py b/py/packages/genkit/src/genkit/_core/_http_client.py new file mode 100644 index 0000000000..5ee9c782d4 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_http_client.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Shared HTTP client utilities for Genkit plugins.""" + +from typing import Any + +import httpx + +from genkit._core._logger import get_logger +from genkit._core._loop_cache import _loop_local_client + +logger = get_logger(__name__) + +_get_store = _loop_local_client(dict) + + +def get_cached_client( + cache_key: str, + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | float | None = None, + **httpx_kwargs: Any, +) -> httpx.AsyncClient: + """Get or create a cached httpx.AsyncClient for the current event loop.""" + d = _get_store() + if cache_key not in d or d[cache_key].is_closed: + if timeout is None: + timeout = httpx.Timeout(60.0, connect=10.0) + elif isinstance(timeout, (int, float)): + timeout = httpx.Timeout(float(timeout)) + d[cache_key] = httpx.AsyncClient(headers=headers or {}, timeout=timeout, **httpx_kwargs) + return d[cache_key] + + +async def close_cached_clients(cache_key: str | None = None) -> None: + """Close and remove cached clients for the current event loop.""" + try: + d = _get_store() + except RuntimeError: + return + + clients_to_close: dict[str, httpx.AsyncClient] = {} + + if cache_key is not None: + if cache_key in d: + clients_to_close[cache_key] = d.pop(cache_key) + else: + clients_to_close.update(d) + d.clear() + + for key, client in clients_to_close.items(): + try: + await client.aclose() + except Exception as e: + logger.warning('Failed to close cached client', cache_key=key, error=e) + + +def clear_client_cache() -> None: + """Clear all cached clients (for testing). Does NOT close clients.""" + try: + d = _get_store() + d.clear() + except RuntimeError: + pass diff --git a/py/packages/genkit/src/genkit/_core/_logger.py b/py/packages/genkit/src/genkit/_core/_logger.py new file mode 100644 index 0000000000..f4f8a3636b --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_logger.py @@ -0,0 +1,8 @@ +# Copyright 2025 Google LLC +# SPDX-License-Identifier: Apache-2.0 + +"""Internal logger for genkit core. Not part of public API.""" + +import structlog + +get_logger = structlog.get_logger diff --git a/py/packages/genkit/src/genkit/core/_loop_local.py b/py/packages/genkit/src/genkit/_core/_loop_cache.py similarity index 95% rename from py/packages/genkit/src/genkit/core/_loop_local.py rename to py/packages/genkit/src/genkit/_core/_loop_cache.py index d9974c00db..73e9fda310 100644 --- a/py/packages/genkit/src/genkit/core/_loop_local.py +++ b/py/packages/genkit/src/genkit/_core/_loop_cache.py @@ -14,7 +14,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Internal loop-local cache for async resources.""" +"""Per-event-loop resource caching for async HTTP clients.""" import asyncio import threading diff --git a/py/packages/genkit/src/genkit/_core/_model.py b/py/packages/genkit/src/genkit/_core/_model.py new file mode 100644 index 0000000000..9ef80e54da --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_model.py @@ -0,0 +1,518 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Model veneer types for the Genkit framework. + +This module contains the hand-written wrapper classes that provide convenient +properties and methods on top of the generated wire types. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Sequence +from copy import deepcopy +from functools import cached_property +from typing import Any, ClassVar, Generic, cast + +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_serializer +from pydantic.alias_generators import to_camel +from typing_extensions import TypeVar + +from genkit._core._base import GenkitModel +from genkit._core._extract_json import extract_json +from genkit._core._typing import ( + Candidate, + DocumentData, + DocumentPart, + FinishReason, + GenerateResponseChunk, + GenerationCommonConfig, + GenerationUsage, + Media, + MediaModel, + MediaPart, + MessageData, + Operation, + Part, + Text, + TextPart, + ToolChoice, + ToolDefinition, + ToolRequestPart, +) + +ModelConfig = GenerationCommonConfig # public name for GenerationCommonConfig +ModelUsage = GenerationUsage # public name for GenerationUsage + +# TypeVars for generic types +OutputT = TypeVar('OutputT', default=object) +ConfigT = TypeVar('ConfigT', bound=ModelConfig, default=ModelConfig) + + +class ModelRef(BaseModel): + """Reference to a model with configuration.""" + + name: str + config_schema: object | None = None + info: object | None = None + version: str | None = None + config: dict[str, object] | None = None + + +class Message(MessageData): + """Message wrapper with utility properties for text and tool requests.""" + + def __init__( + self, + message: MessageData | None = None, + **kwargs: object, + ) -> None: + """Initialize from MessageData or keyword arguments.""" + if message is not None: + super().__init__( + role=message.role, + content=message.content, + metadata=message.metadata, + ) + else: + super().__init__(**kwargs) # type: ignore[arg-type] + + def __eq__(self, other: object) -> bool: + """Compare messages by role, content, and metadata.""" + if isinstance(other, MessageData): + return self.role == other.role and self.content == other.content and self.metadata == other.metadata + return super().__eq__(other) + + def __hash__(self) -> int: + """Return identity-based hash.""" + return hash(id(self)) + + @cached_property + def text(self) -> str: + """All text parts concatenated into a single string.""" + return text_from_message(self) + + @cached_property + def tool_requests(self) -> list[ToolRequestPart]: + """All tool request parts in this message.""" + return [p.root for p in self.content if isinstance(p.root, ToolRequestPart)] + + @cached_property + def interrupts(self) -> list[ToolRequestPart]: + """Tool requests marked as interrupted.""" + return [p for p in self.tool_requests if p.metadata and p.metadata.root.get('interrupt')] + + +_TEXT_DATA_TYPE: str = 'text' + + +class Document(DocumentData): + """Multi-part document that can be embedded, indexed, or retrieved.""" + + def __init__( + self, + content: list[DocumentPart], + metadata: dict[str, Any] | None = None, + ) -> None: + """Initialize with content parts and optional metadata.""" + doc_content = deepcopy(content) + doc_metadata = deepcopy(metadata) + super().__init__(content=doc_content, metadata=doc_metadata) + + @staticmethod + def from_text(text: str, metadata: dict[str, Any] | None = None) -> Document: + """Create a document from a text string.""" + return Document(content=[DocumentPart(root=TextPart(text=text))], metadata=metadata) + + @staticmethod + def from_media( + url: str, + content_type: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> Document: + """Create a document from a media URL.""" + return Document( + content=[DocumentPart(root=MediaPart(media=Media(url=url, content_type=content_type)))], + metadata=metadata, + ) + + @staticmethod + def from_data( + data: str, + data_type: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> Document: + """Create a document from data, inferring text vs media from data_type.""" + if data_type == _TEXT_DATA_TYPE: + return Document.from_text(data, metadata) + return Document.from_media(data, data_type, metadata) + + @cached_property + def text(self) -> str: + """Concatenate all text parts.""" + texts = [] + for p in self.content: + part = p.root if hasattr(p, 'root') else p + text_val = getattr(part, 'text', None) + if isinstance(text_val, str): + texts.append(text_val) + return ''.join(texts) + + @cached_property + def media(self) -> list[Media]: + """Get all media parts.""" + return [ + part.root.media for part in self.content if isinstance(part.root, MediaPart) and part.root.media is not None + ] + + @cached_property + def data(self) -> str: + """Primary data: text if available, otherwise first media URL.""" + if self.text: + return self.text + if self.media: + return self.media[0].url + return '' + + @cached_property + def data_type(self) -> str | None: + """Type of primary data: 'text' or first media's content type.""" + if self.text: + return _TEXT_DATA_TYPE + if self.media and self.media[0].content_type: + return self.media[0].content_type + return None + + +class ModelRequest(GenkitModel, Generic[ConfigT]): + """Hand-written model request with flat output fields and veneer types. + + Output config is inlined as flat fields (output_format, output_schema, etc.) + instead of a nested OutputConfig object. Messages and docs use veneer types + (Message, Document) for convenience methods like .text. + + Example: + class GeminiConfig(ModelConfig): + safety_settings: dict[str, str] | None = None + + def gemini_model(request: ModelRequest[GeminiConfig]) -> ModelResponse: + temp = request.config.temperature # inherited from ModelConfig + for msg in request.messages: + print(msg.text) # Message veneer property + if request.output_format == 'json': + schema = request.output_schema + """ + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + # Veneer types for IDE/typing (validators wrap MessageData->Message, DocumentData->Document) + messages: list[Message] # pyright: ignore[reportIncompatibleVariableOverride] + docs: list[Document] | None = None # pyright: ignore[reportIncompatibleVariableOverride] + config: ConfigT | None = None + tools: list[ToolDefinition] | None = None + tool_choice: ToolChoice | None = Field(default=None) + # Flat output fields (no nested OutputConfig) + output_format: str | None = None + output_schema: dict[str, Any] | None = None + output_constrained: bool | None = None + output_content_type: str | None = None + + @field_validator('messages', mode='before') + @classmethod + def _wrap_messages(cls, v: list[MessageData]) -> list[Message]: + """Wrap MessageData in Message veneer for convenience methods.""" + # pyrefly: ignore[bad-return] + return [m if isinstance(m, Message) else Message(m) for m in v] + + @field_validator('docs', mode='before') + @classmethod + def _wrap_docs(cls, v: list[DocumentData] | None) -> list[Document] | None: + """Wrap DocumentData in Document veneer for convenience methods.""" + if v is None: + return None + # pyrefly: ignore[bad-return] + return [d if isinstance(d, Document) else Document(d.content, d.metadata) for d in v] + + @model_serializer(mode='wrap') + def _serialize_for_spec(self, serializer: Callable[..., dict[str, Any]]) -> dict[str, Any]: + """Serialize to spec wire format with nested output (matches JS/Go).""" + data = serializer(self) + # Build nested output from flat fields - spec expects output key always present + output: dict[str, Any] = {} + if self.output_format is not None: + output['format'] = self.output_format + if self.output_schema is not None: + output['schema'] = self.output_schema + if self.output_constrained is not None: + output['constrained'] = self.output_constrained + if self.output_content_type is not None: + output['contentType'] = self.output_content_type + # Remove flat fields, add nested output + data.pop('outputFormat', None) + data.pop('outputSchema', None) + data.pop('outputConstrained', None) + data.pop('outputContentType', None) + data['output'] = output + return data + + +class ModelResponse(GenkitModel, Generic[OutputT]): + """Model response with utilities for text extraction, output parsing, and validation.""" + + # _message_parser and _schema_type are set by the framework after construction + # when output format parsing or schema validation is needed. + _message_parser: Callable[[Message], object] | None = PrivateAttr(None) + _schema_type: type[BaseModel] | None = PrivateAttr(None) + # Wire fields (must be declared for extra='forbid' to accept wire responses) + message: Message | None = None + finish_reason: FinishReason | None = None + finish_message: str | None = None + latency_ms: float | None = None + usage: GenerationUsage | None = None + custom: dict[str, Any] | None = None + raw: dict[str, Any] | None = None + request: ModelRequest | None = None + operation: Operation | None = None + candidates: list[Candidate] | None = None + + def model_post_init(self, __context: object) -> None: + """Initialize default usage and custom dict if not provided.""" + if self.usage is None: + self.usage = GenerationUsage() + if self.custom is None: + self.custom = {} + + def assert_valid(self) -> None: + """Validate response structure. (TODO: not yet implemented).""" + # TODO(#4343): implement + pass + + def assert_valid_schema(self) -> None: + """Validate response conforms to output schema. (TODO: not yet implemented).""" + # TODO(#4343): implement + pass + + def __eq__(self, other: object) -> bool: + """Compare responses by message and finish_reason.""" + if isinstance(other, ModelResponse): + return self.message == other.message and self.finish_reason == other.finish_reason + return super().__eq__(other) + + def __hash__(self) -> int: + """Return identity-based hash.""" + return hash(id(self)) + + @cached_property + def text(self) -> str: + """All text parts concatenated into a single string.""" + if self.message is None: + return '' + return self.message.text + + @cached_property + def output(self) -> OutputT: + """Parsed JSON output from the response text, validated against schema if set.""" + if self._message_parser and self.message is not None: + parsed = self._message_parser(self.message) + else: + parsed = extract_json(self.text) + + # If we have a schema type and the parsed output is a dict, validate and + # return a proper Pydantic instance. Skip if parsed is already the correct + # type or if it's not a dict (e.g., custom formats may return strings). + if self._schema_type is not None and parsed is not None and isinstance(parsed, dict): + return cast(OutputT, self._schema_type.model_validate(parsed)) + + return cast(OutputT, parsed) + + @cached_property + def messages(self) -> list[Message]: + """All messages including request history and the response message.""" + if self.message is None: + return [Message(m) for m in self.request.messages] if self.request else [] + return [ + *(Message(m) for m in (self.request.messages if self.request else [])), + self.message, + ] + + @cached_property + def tool_requests(self) -> list[ToolRequestPart]: + """All tool request parts in the response message.""" + if self.message is None: + return [] + return self.message.tool_requests + + @cached_property + def interrupts(self) -> list[ToolRequestPart]: + """Tool requests marked as interrupted.""" + if self.message is None: + return [] + return self.message.interrupts + + +class ModelResponseChunk(GenerateResponseChunk, Generic[OutputT]): + """Streaming chunk with text, accumulated text, and output parsing.""" + + # Field(exclude=True) means these fields are not included in serialization + previous_chunks: list[ModelResponseChunk[Any]] = Field(default_factory=list, exclude=True) + chunk_parser: Callable[[ModelResponseChunk[Any]], object] | None = Field(None, exclude=True) + + def __init__( + self, + chunk: ModelResponseChunk[Any] | None = None, + previous_chunks: list[ModelResponseChunk[Any]] | None = None, + index: int | float | None = None, + chunk_parser: Callable[[ModelResponseChunk[Any]], object] | None = None, + **kwargs: Any, # noqa: ANN401 + ) -> None: + """Initialize from a chunk or keyword arguments.""" + if chunk is not None: + # Framework wrapping mode + super().__init__( + role=chunk.role, + index=index, + content=chunk.content, + custom=chunk.custom, + aggregated=chunk.aggregated, + ) + else: + # No source chunk β€” caller passes fields (role, content, etc.) as kwargs directly + super().__init__(**kwargs) + self.previous_chunks = previous_chunks or [] + self.chunk_parser = chunk_parser + + def __eq__(self, other: object) -> bool: + """Check equality.""" + if isinstance(other, ModelResponseChunk): + return self.role == other.role and self.content == other.content + return super().__eq__(other) + + def __hash__(self) -> int: + """Return hash.""" + return hash(id(self)) + + @cached_property + def text(self) -> str: + """Text content of this chunk.""" + parts: list[str] = [] + for p in self.content: + text_val = p.root.text + if text_val is not None: + # Handle Text RootModel (access .root) or plain str + if isinstance(text_val, Text): + parts.append(str(text_val.root) if text_val.root is not None else '') + else: + parts.append(str(text_val)) + return ''.join(parts) + + @cached_property + def accumulated_text(self) -> str: + """Text from all previous chunks plus this chunk.""" + parts: list[str] = [] + if self.previous_chunks: + for chunk in self.previous_chunks: + for p in chunk.content: + text_val = p.root.text + if text_val: + # Handle Text RootModel (access .root) or plain str + if isinstance(text_val, Text): + parts.append(str(text_val.root) if text_val.root is not None else '') + else: + parts.append(str(text_val)) + return ''.join(parts) + self.text + + @cached_property + def output(self) -> OutputT: + """Parsed JSON output from accumulated text.""" + if self.chunk_parser: + return cast(OutputT, self.chunk_parser(self)) + return cast(OutputT, extract_json(self.accumulated_text)) + + +def text_from_message(msg: Message) -> str: + """Concatenate text from all parts of a message.""" + return text_from_content(msg.content) + + +def text_from_content(content: Sequence[Part | DocumentPart]) -> str: + """Concatenate text from a list of parts.""" + return ''.join(str(p.root.text) for p in content if hasattr(p.root, 'text') and p.root.text is not None) + + +def get_basic_usage_stats(input_: list[Message], response: Message) -> GenerationUsage: + """Calculate usage stats (characters, media counts) from messages.""" + request_parts: list[Part] = [] + for msg in input_: + request_parts.extend(msg.content) + + response_parts = response.content + + def count_parts(parts: list[Part]) -> tuple[int, int, int, int]: + """Count characters, images, videos, audio in parts.""" + characters = 0 + images = 0 + videos = 0 + audio = 0 + + for part in parts: + text_val = part.root.text + if text_val: + if isinstance(text_val, Text): + characters += len(str(text_val.root)) if text_val.root else 0 + else: + characters += len(str(text_val)) + + media = part.root.media + if media: + if isinstance(media, Media): + content_type = media.content_type or '' + url = media.url or '' + elif isinstance(media, MediaModel) and hasattr(media.root, 'content_type'): + content_type = getattr(media.root, 'content_type', '') or '' + url = getattr(media.root, 'url', '') or '' + else: + content_type = '' + url = '' + + if content_type.startswith('image') or url.startswith('data:image'): + images += 1 + elif content_type.startswith('video') or url.startswith('data:video'): + videos += 1 + elif content_type.startswith('audio') or url.startswith('data:audio'): + audio += 1 + + return characters, images, videos, audio + + in_chars, in_imgs, in_vids, in_audio = count_parts(request_parts) + out_chars, out_imgs, out_vids, out_audio = count_parts(response_parts) + + return GenerationUsage( + input_characters=in_chars, + input_images=in_imgs, + input_videos=in_vids, + input_audio_files=in_audio, + output_characters=out_chars, + output_images=out_imgs, + output_videos=out_vids, + output_audio_files=out_audio, + ) + + +# Type aliases for model middleware (Any is intentional - middleware is type-agnostic) +# Middleware can have two signatures: +# Simple (3 params): (req, ctx, next) -> response +# Streaming (4 params): (req, ctx, on_chunk, next) -> response +# The framework detects which signature is used based on parameter count. +ModelMiddleware = Callable[..., Awaitable[ModelResponse[Any]]] diff --git a/py/packages/genkit/src/genkit/_core/_plugin.py b/py/packages/genkit/src/genkit/_core/_plugin.py new file mode 100644 index 0000000000..74d0c22566 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_plugin.py @@ -0,0 +1,52 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Abstract base class for Genkit plugins.""" + +import abc + +from genkit._core._action import Action, ActionKind, ActionMetadata + + +class Plugin(abc.ABC): + """Abstract base class for Genkit plugins.""" + + name: str # plugin namespace + + @abc.abstractmethod + async def init(self) -> list[Action]: + """Lazy warm-up called once per plugin; return actions to pre-register.""" + ... + + @abc.abstractmethod + async def resolve(self, action_type: ActionKind, name: str) -> Action | None: + """Resolve a single action by kind and namespaced name.""" + ... + + @abc.abstractmethod + async def list_actions(self) -> list[ActionMetadata]: + """Return advertised actions for dev UI/reflection listing.""" + ... + + async def model(self, name: str) -> Action | None: + """Resolve a model action by name (local or namespaced).""" + target = name if '/' in name else f'{self.name}/{name}' + return await self.resolve(ActionKind.MODEL, target) + + async def embedder(self, name: str) -> Action | None: + """Resolve an embedder action by name (local or namespaced).""" + target = name if '/' in name else f'{self.name}/{name}' + return await self.resolve(ActionKind.EMBEDDER, target) diff --git a/py/packages/genkit/src/genkit/core/_plugins.py b/py/packages/genkit/src/genkit/_core/_plugins.py similarity index 98% rename from py/packages/genkit/src/genkit/core/_plugins.py rename to py/packages/genkit/src/genkit/_core/_plugins.py index ec79f753b0..140377f303 100644 --- a/py/packages/genkit/src/genkit/core/_plugins.py +++ b/py/packages/genkit/src/genkit/_core/_plugins.py @@ -66,6 +66,3 @@ def extend_plugin_namespace() -> None: if plugins_path_str not in existing_paths: genkit_plugins.__path__.append(plugins_path_str) existing_paths.add(plugins_path_str) - - -__all__ = ['extend_plugin_namespace'] diff --git a/py/packages/genkit/src/genkit/_core/_reflection.py b/py/packages/genkit/src/genkit/_core/_reflection.py new file mode 100644 index 0000000000..76bc1253e3 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_reflection.py @@ -0,0 +1,264 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Reflection API server for Genkit Dev UI.""" + +from __future__ import annotations + +import asyncio +import json +import os +import signal +import threading +from collections.abc import AsyncGenerator, Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +import uvicorn +from pydantic import BaseModel +from starlette.applications import Starlette +from starlette.middleware import Middleware +from starlette.middleware.cors import CORSMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse, Response, StreamingResponse +from starlette.routing import Route + +from genkit._core._action import Action, ActionKind +from genkit._core._constants import GENKIT_VERSION +from genkit._core._error import get_reflection_json +from genkit._core._logger import get_logger +from genkit._core._registry import Registry + +logger = get_logger(__name__) + +LifecycleHook = Callable[[], Awaitable[None]] + + +@dataclass +class ServerSpec: + port: int + scheme: str = 'http' + host: str = 'localhost' + + @property + def url(self) -> str: + return f'{self.scheme}://{self.host}:{self.port}' + + +@dataclass +class ActionRunner: + """Encapsulates state for running an action with streaming support.""" + + action: Action + payload: dict[str, Any] + stream: bool + active_actions: dict[str, asyncio.Task[Any]] + + queue: asyncio.Queue[str | None] = field(default_factory=asyncio.Queue) + trace_ready: asyncio.Event = field(default_factory=asyncio.Event) + trace_id: str | None = None + span_id: str | None = None + + def on_trace_start(self, tid: str, sid: str) -> None: + self.trace_id, self.span_id = tid, sid + if task := asyncio.current_task(): + self.active_actions[tid] = task + self.trace_ready.set() + + async def execute(self) -> None: + try: + on_chunk = (lambda c: self.queue.put_nowait(f'{c.model_dump_json()}\n')) if self.stream else None + output = await self.action.run( + input=self.payload.get('input'), + on_chunk=on_chunk, + context=self.payload.get('context', {}), + on_trace_start=self.on_trace_start, + ) + result = output.response.model_dump() if isinstance(output.response, BaseModel) else output.response + self.queue.put_nowait( + json.dumps({ + 'result': result, + 'telemetry': {'traceId': output.trace_id, 'spanId': output.span_id}, + }) + ) + except asyncio.CancelledError: + raise + except Exception as e: + logger.exception('Error executing action') + self.queue.put_nowait(json.dumps({'error': get_reflection_json(e).model_dump(by_alias=True)})) + finally: + self.trace_ready.set() + self.queue.put_nowait(None) + if self.trace_id: + self.active_actions.pop(self.trace_id, None) + + async def stream_response(self, version: str) -> StreamingResponse: + task = asyncio.create_task(self.execute()) + await self.trace_ready.wait() + + headers = {'x-genkit-version': version, 'Transfer-Encoding': 'chunked'} + if self.trace_id: + headers['X-Genkit-Trace-Id'] = self.trace_id + if self.span_id: + headers['X-Genkit-Span-Id'] = self.span_id + + async def gen() -> AsyncGenerator[str, None]: + try: + while (chunk := await self.queue.get()) is not None: + yield chunk + finally: + task.cancel() + + return StreamingResponse(gen(), media_type='text/plain' if self.stream else 'application/json', headers=headers) + + +async def _get_actions_payload(registry: Registry) -> dict[str, dict[str, Any]]: + actions: dict[str, dict[str, Any]] = {} + + for kind in ActionKind.__members__.values(): + for name, action in (await registry.resolve_actions_by_kind(kind)).items(): + key = f'/{kind}/{name}' + actions[key] = { + 'key': key, + 'name': action.name, + 'type': action.kind, + 'description': action.description, + 'inputSchema': action.input_schema, + 'outputSchema': action.output_schema, + 'metadata': action.metadata, + } + + for meta in await registry.list_actions() or []: + try: + key = f'/{meta.kind}/{meta.name}' + except Exception as exc: + logger.warning('Skipping invalid plugin metadata: %s', exc) + continue + + advertised = { + 'key': key, + 'name': meta.name, + 'type': meta.kind, + 'description': getattr(meta, 'description', None), + 'inputSchema': getattr(meta, 'input_json_schema', None), + 'outputSchema': getattr(meta, 'output_json_schema', None), + 'metadata': getattr(meta, 'metadata', None), + } + + if key not in actions: + actions[key] = advertised + else: + existing = actions[key] + for f in ('description', 'inputSchema', 'outputSchema'): + if not existing.get(f) and advertised.get(f): + existing[f] = advertised[f] + if isinstance(existing.get('metadata'), dict) and isinstance(advertised.get('metadata'), dict): + # isinstance checks above guarantee both are dicts, but ty can't narrow .get() results + existing['metadata'] = {**advertised['metadata'], **existing['metadata']} # ty: ignore[invalid-argument-type] + + return actions + + +def create_reflection_asgi_app( + registry: Registry, + on_startup: LifecycleHook | None = None, + on_shutdown: LifecycleHook | None = None, + version: str = GENKIT_VERSION, +) -> Starlette: + active_actions: dict[str, asyncio.Task[Any]] = {} + + async def health(_: Request) -> JSONResponse: + return JSONResponse({'status': 'OK'}) + + async def terminate(_: Request) -> JSONResponse: + logger.info('Shutting down...') + asyncio.get_running_loop().call_soon(os.kill, os.getpid(), signal.SIGTERM) + return JSONResponse({'status': 'OK'}) + + async def actions(_: Request) -> JSONResponse: + return JSONResponse(await _get_actions_payload(registry), headers={'x-genkit-version': version}) + + async def values(req: Request) -> JSONResponse: + if req.query_params.get('type') != 'defaultModel': + return JSONResponse({'error': 'Only type=defaultModel supported'}, status_code=400) + return JSONResponse(registry.list_values('defaultModel')) + + async def envs(_: Request) -> JSONResponse: + return JSONResponse(['dev']) + + async def notify(_: Request) -> JSONResponse: + return JSONResponse({}, headers={'x-genkit-version': version}) + + async def cancel(req: Request) -> JSONResponse: + trace_id = (await req.json()).get('traceId') + if not trace_id: + return JSONResponse({'error': 'traceId required'}, status_code=400) + if task := active_actions.get(trace_id): + task.cancel() + return JSONResponse({'message': 'Cancelled'}) + return JSONResponse({'message': 'Not found'}, status_code=404) + + async def run(req: Request) -> Response: + payload = await req.json() + action = await registry.resolve_action_by_key(payload['key']) + if not action: + return JSONResponse({'error': f'Action not found: {payload["key"]}'}, status_code=404) + + runner = ActionRunner( + action=action, + payload=payload, + stream=req.headers.get('accept') == 'text/event-stream' or req.query_params.get('stream') == 'true', + active_actions=active_actions, + ) + return await runner.stream_response(version) + + app = Starlette( + routes=[ + Route('/api/__health', health, methods=['GET']), + Route('/api/__quitquitquit', terminate, methods=['GET', 'POST']), + Route('/api/actions', actions, methods=['GET']), + Route('/api/values', values, methods=['GET']), + Route('/api/envs', envs, methods=['GET']), + Route('/api/notify', notify, methods=['POST']), + Route('/api/runAction', run, methods=['POST']), + Route('/api/cancelAction', cancel, methods=['POST']), + ], + middleware=[ + Middleware( + CORSMiddleware, # type: ignore[arg-type] + allow_origins=['*'], + allow_methods=['*'], + allow_headers=['*'], + expose_headers=['X-Genkit-Trace-Id', 'X-Genkit-Span-Id', 'x-genkit-version'], + ) + ], + on_startup=[on_startup] if on_startup else [], + on_shutdown=[on_shutdown] if on_shutdown else [], + ) + app.active_actions = active_actions # type: ignore[attr-defined] + return app + + +class ReflectionServer(uvicorn.Server): + def __init__(self, config: uvicorn.Config, ready: threading.Event) -> None: + super().__init__(config) + self._ready = ready + + async def startup(self, sockets: list | None = None) -> None: + try: + await super().startup(sockets=sockets) + finally: + self._ready.set() diff --git a/py/packages/genkit/src/genkit/core/registry.py b/py/packages/genkit/src/genkit/_core/_registry.py similarity index 90% rename from py/packages/genkit/src/genkit/core/registry.py rename to py/packages/genkit/src/genkit/_core/_registry.py index 34e4f1a146..b4731900be 100644 --- a/py/packages/genkit/src/genkit/core/registry.py +++ b/py/packages/genkit/src/genkit/_core/_registry.py @@ -14,18 +14,7 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Registry for managing Genkit resources and actions. - -This module provides the Registry class, which is the central repository for -storing and managing various Genkit resources such as actions, flows, -plugins, and schemas. The registry enables dynamic registration and lookup -of these resources during runtime. - -Example: - >>> registry = Registry() - >>> registry.register_action('', 'my_action', ...) - >>> action = await registry.resolve_action('', 'my_action') -""" +"""Registry for managing Genkit resources and actions.""" import asyncio import threading @@ -36,28 +25,27 @@ from pydantic import BaseModel from typing_extensions import Never, TypeVar -from genkit.core.action import ( +from genkit._core._action import ( Action, + ActionKind, ActionMetadata, + ActionName, ActionRunContext, SpanAttributeValue, parse_action_key, ) -from genkit.core.action.types import ActionKind, ActionName -from genkit.core.logging import get_logger -from genkit.core.plugin import Plugin -from genkit.core.typing import ( +from genkit._core._logger import get_logger +from genkit._core._model import ( + ModelRequest, + ModelResponse, + ModelResponseChunk, +) +from genkit._core._plugin import Plugin +from genkit._core._typing import ( EmbedRequest, EmbedResponse, EvalRequest, EvalResponse, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - RerankerRequest, - RerankerResponse, - RetrieverRequest, - RetrieverResponse, ) logger = get_logger(__name__) @@ -441,7 +429,7 @@ async def resolve_action(self, kind: ActionKind, name: str) -> Action | None: if len(successes) > 1: plugin_names = [p for p, _ in successes] raise ValueError( - f"Ambiguous {kind.value} action name '{name}'. " + f"Ambiguous {kind} action name '{name}'. " + f"Matches plugins: {plugin_names}. Use 'plugin/{name}'." ) @@ -462,7 +450,7 @@ async def resolve_action(self, kind: ActionKind, name: str) -> Action | None: providers = list(providers_dict.values()) for provider in providers: try: - response = await provider.arun({'kind': kind, 'name': name}) + response = await provider.run({'kind': kind, 'name': name}) if response.response: self.register_action_instance(response.response) return await self._trigger_lazy_loading(response.response) @@ -580,20 +568,6 @@ def lookup_schema_type(self, name: str) -> type[BaseModel] | None: # They wrap resolve_action() with appropriate casts to preserve generic # type parameters that would otherwise be erased. - async def resolve_retriever(self, name: str) -> Action[RetrieverRequest, RetrieverResponse, Never] | None: - """Resolve a retriever action by name with full type information. - - Args: - name: The retriever name (e.g., "my-retriever" or "plugin/retriever"). - - Returns: - A fully typed retriever action, or None if not found. - """ - action = await self.resolve_action(ActionKind.RETRIEVER, name) - if action is None: - return None - return cast(Action[RetrieverRequest, RetrieverResponse, Never], action) - async def resolve_embedder(self, name: str) -> Action[EmbedRequest, EmbedResponse, Never] | None: """Resolve an embedder action by name with full type information. @@ -608,21 +582,7 @@ async def resolve_embedder(self, name: str) -> Action[EmbedRequest, EmbedRespons return None return cast(Action[EmbedRequest, EmbedResponse, Never], action) - async def resolve_reranker(self, name: str) -> Action[RerankerRequest, RerankerResponse, Never] | None: - """Resolve a reranker action by name with full type information. - - Args: - name: The reranker name (e.g., "my-reranker" or "plugin/reranker"). - - Returns: - A fully typed reranker action, or None if not found. - """ - action = await self.resolve_action(ActionKind.RERANKER, name) - if action is None: - return None - return cast(Action[RerankerRequest, RerankerResponse, Never], action) - - async def resolve_model(self, name: str) -> Action[GenerateRequest, GenerateResponse, GenerateResponseChunk] | None: + async def resolve_model(self, name: str) -> Action[ModelRequest, ModelResponse, ModelResponseChunk] | None: """Resolve a model action by name with full type information. Args: @@ -635,7 +595,7 @@ async def resolve_model(self, name: str) -> Action[GenerateRequest, GenerateResp if action is None: return None return cast( - Action[GenerateRequest, GenerateResponse, GenerateResponseChunk], + Action[ModelRequest, ModelResponse, ModelResponseChunk], action, ) diff --git a/py/samples/framework-restaurant-demo/src/case_04/__init__.py b/py/packages/genkit/src/genkit/_core/_schema.py similarity index 57% rename from py/samples/framework-restaurant-demo/src/case_04/__init__.py rename to py/packages/genkit/src/genkit/_core/_schema.py index 0105fe7c98..0a60e61102 100644 --- a/py/samples/framework-restaurant-demo/src/case_04/__init__.py +++ b/py/packages/genkit/src/genkit/_core/_schema.py @@ -14,5 +14,18 @@ # # SPDX-License-Identifier: Apache-2.0 +"""Functions for working with schema.""" -"""Case 04 package.""" +from typing import Any + +from pydantic import TypeAdapter + + +def to_json_schema(schema: type | dict[str, Any] | str | None) -> dict[str, Any]: + """Convert a Python type to JSON schema. Pass-through if already a dict.""" + if schema is None: + return {'type': 'null'} + if isinstance(schema, dict): + return schema + type_adapter = TypeAdapter(schema) + return type_adapter.json_schema() diff --git a/py/packages/genkit/src/genkit/_core/_trace/_adjusting_exporter.py b/py/packages/genkit/src/genkit/_core/_trace/_adjusting_exporter.py new file mode 100644 index 0000000000..3cfedafcd9 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_trace/_adjusting_exporter.py @@ -0,0 +1,154 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Adjusting trace exporter for PII redaction and span enhancement.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from typing import Any, ClassVar + +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult +from opentelemetry.trace import StatusCode + +from genkit._core._compat import override + + +def _copy_attrs(span: ReadableSpan) -> dict[str, Any]: + """Return a mutable copy of span attributes.""" + return dict(span.attributes) if span.attributes else {} + + +class RedactedSpan(ReadableSpan): + """A span wrapper that overrides attributes while delegating everything else.""" + + # pyrefly:ignore[bad-override] + _attributes: dict[str, Any] + + def __init__(self, span: ReadableSpan, attributes: dict[str, Any]) -> None: + self._span = span + self._attributes = attributes + + def __getattr__(self, name: str) -> Any: # noqa: ANN401 + return getattr(self._span, name) + + @property + def attributes(self) -> dict[str, Any]: + """Return the modified attributes.""" + # pyrefly: ignore[bad-return] - dict[str, Any] is compatible with Mapping at runtime + return self._attributes + + +class AdjustingTraceExporter(SpanExporter): + """Wraps a SpanExporter to redact PII and enhance spans for cloud plugins (GCP, AWS).""" + + REDACTED: ClassVar[str] = '' + + def __init__( + self, + exporter: SpanExporter, + log_input_and_output: bool = False, + project_id: str | None = None, + error_handler: Callable[[Exception], None] | None = None, + ) -> None: + self._exporter = exporter + self._log_input_and_output = log_input_and_output + self._project_id = project_id + self._error_handler = error_handler + + @property + def project_id(self) -> str | None: + return self._project_id + + @property + def log_input_and_output(self) -> bool: + return self._log_input_and_output + + @override + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + adjusted = [self._adjust(span) for span in spans] + try: + return self._exporter.export(adjusted) + except Exception as e: + if self._error_handler: + self._error_handler(e) + raise + + @override + def shutdown(self) -> None: + self._exporter.shutdown() + + @override + def force_flush(self, timeout_millis: int = 30000) -> bool: + if hasattr(self._exporter, 'force_flush'): + return self._exporter.force_flush(timeout_millis) + return True + + def _adjust(self, span: ReadableSpan) -> ReadableSpan: + """Apply all adjustments to a span.""" + span = self._redact_pii(span) + span = self._mark_error(span) + span = self._mark_failure_source(span) + span = self._mark_feature(span) + span = self._mark_model(span) + span = self._normalize_labels(span) + return span + + def _redact_pii(self, span: ReadableSpan) -> ReadableSpan: + if self._log_input_and_output: + return span + attrs = _copy_attrs(span) + keys_to_redact = [k for k in ('genkit:input', 'genkit:output') if k in attrs] + if not keys_to_redact: + return span + for key in keys_to_redact: + attrs[key] = self.REDACTED + return RedactedSpan(span, attrs) + + def _mark_error(self, span: ReadableSpan) -> ReadableSpan: + if not span.status or span.status.status_code != StatusCode.ERROR: + return span + attrs = _copy_attrs(span) + attrs['/http/status_code'] = '599' + return RedactedSpan(span, attrs) + + def _mark_failure_source(self, span: ReadableSpan) -> ReadableSpan: + attrs = _copy_attrs(span) + if not attrs.get('genkit:isFailureSource'): + return span + attrs['genkit:failedSpan'] = attrs.get('genkit:name', '') + attrs['genkit:failedPath'] = attrs.get('genkit:path', '') + return RedactedSpan(span, attrs) + + def _mark_feature(self, span: ReadableSpan) -> ReadableSpan: + attrs = _copy_attrs(span) + if not attrs.get('genkit:isRoot') or not attrs.get('genkit:name'): + return span + attrs['genkit:feature'] = attrs['genkit:name'] + return RedactedSpan(span, attrs) + + def _mark_model(self, span: ReadableSpan) -> ReadableSpan: + attrs = _copy_attrs(span) + if attrs.get('genkit:metadata:subtype') != 'model' or not attrs.get('genkit:name'): + return span + attrs['genkit:model'] = attrs['genkit:name'] + return RedactedSpan(span, attrs) + + def _normalize_labels(self, span: ReadableSpan) -> ReadableSpan: + attrs = _copy_attrs(span) + normalized = {k.replace(':', '/'): v for k, v in attrs.items()} + return RedactedSpan(span, normalized) diff --git a/py/packages/genkit/src/genkit/_core/_trace/_default_exporter.py b/py/packages/genkit/src/genkit/_core/_trace/_default_exporter.py new file mode 100644 index 0000000000..fc23c64ec8 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_trace/_default_exporter.py @@ -0,0 +1,128 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Telemetry and tracing default exporter for the Genkit framework.""" + +from __future__ import annotations + +import os +from collections.abc import Sequence +from typing import Any, cast +from urllib.parse import urljoin + +import httpx +from opentelemetry import trace as trace_api +from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor +from opentelemetry.sdk.trace.export import ( + BatchSpanProcessor, + SpanExporter, + SpanExportResult, +) +from opentelemetry.trace import SpanContext + +from genkit._core._compat import override +from genkit._core._environment import is_dev_environment +from genkit._core._logger import get_logger + +from ._realtime_processor import RealtimeSpanProcessor + +logger = get_logger(__name__) + +INSTRUMENTATION = {'name': 'genkit-tracer', 'version': 'v1'} +TRACE_HEADERS = {'Content-Type': 'application/json', 'Accept': 'application/json'} + + +def _ns_to_ms(ns: int | None) -> float: + return ns / 1_000_000 if ns is not None else 0 + + +def extract_span_data(span: ReadableSpan) -> dict[str, Any]: + """Convert ReadableSpan to Genkit telemetry server JSON format.""" + ctx = cast(SpanContext, span.context) + trace_id = format(ctx.trace_id, '032x') + span_id = format(ctx.span_id, '016x') + parent_id = format(span.parent.span_id, '016x') if span.parent else None + start = _ns_to_ms(span.start_time) + end = _ns_to_ms(span.end_time) + + span_entry: dict[str, Any] = { + 'spanId': span_id, + 'traceId': trace_id, + 'startTime': start, + 'endTime': end, + 'attributes': dict(span.attributes or {}), + 'displayName': span.name, + 'spanKind': trace_api.SpanKind(span.kind).name, + 'instrumentationLibrary': {'name': 'genkit-tracer', 'version': 'v1'}, + } + if parent_id: + span_entry['parentSpanId'] = parent_id + if span.status: + span_entry['status'] = { + 'code': trace_api.StatusCode(span.status.status_code).value, + 'description': span.status.description, + } + + result: dict[str, Any] = {'traceId': trace_id, 'spans': {span_id: span_entry}} + if not span.parent: + result['displayName'] = span.name + result['startTime'] = start + result['endTime'] = end + + return result + + +class TraceServerExporter(SpanExporter): + """Exports spans to Genkit telemetry server (DevUI).""" + + def __init__( + self, + telemetry_server_url: str, + telemetry_server_endpoint: str = '/api/traces', + ) -> None: + self.telemetry_server_url = telemetry_server_url + self.telemetry_server_endpoint = telemetry_server_endpoint + + @override + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + url = urljoin(self.telemetry_server_url, self.telemetry_server_endpoint) + headers = {'Content-Type': 'application/json', 'Accept': 'application/json'} + with httpx.Client() as client: + for span in spans: + client.post(url, json=extract_span_data(span), headers=headers) + return SpanExportResult.SUCCESS + + @override + def force_flush(self, timeout_millis: int = 30000) -> bool: + return True + + +def init_telemetry_server_exporter() -> SpanExporter | None: + """Return TraceServerExporter if GENKIT_TELEMETRY_SERVER is set, else None.""" + url = os.environ.get('GENKIT_TELEMETRY_SERVER') + if not url: + logger.warn( + 'GENKIT_TELEMETRY_SERVER is not set. If running with `genkit start`, make sure `genkit-cli` is up to date.' + ) + return None + return TraceServerExporter(telemetry_server_url=url) + + +def create_span_processor(exporter: SpanExporter) -> SpanProcessor: + """RealtimeSpanProcessor in dev, BatchSpanProcessor in production.""" + if is_dev_environment(): + return RealtimeSpanProcessor(exporter) + return BatchSpanProcessor(exporter) diff --git a/py/packages/genkit/src/genkit/_core/_trace/_path.py b/py/packages/genkit/src/genkit/_core/_trace/_path.py new file mode 100644 index 0000000000..e6539a3db3 --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_trace/_path.py @@ -0,0 +1,65 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Path utilities for Genkit trace paths. Format: /{name,t:type,s:subtype}.""" + +import re +from urllib.parse import quote + +_PATH_SEGMENT_RE = re.compile(r'\{([^,}]+),[^}]+\}') + + +def build_path( + name: str, + parent_path: str, + type_str: str, + subtype: str | None = None, +) -> str: + """Build hierarchical path: /{name,t:type,s:subtype}.""" + segment = quote(name, safe='') + if type_str: + segment = f'{segment},t:{type_str}' + if subtype: + segment = f'{segment},s:{subtype}' + return f'{parent_path}/{{{segment}}}' + + +def _has_subtype(segment_inner: str) -> bool: + parts = segment_inner.split(',')[1:] # skip name, check annotations only + return any(p.strip().startswith('s:') for p in parts) + + +def decorate_path_with_subtype(path: str, subtype: str) -> str: + """Add subtype to leaf node. Idempotent if subtype already present.""" + if not path or not subtype: + return path + start = path.rfind('{') + if start == -1: + return path + end = path.find('}', start) + if end == -1: + return path + inner = path[start + 1 : end] + if _has_subtype(inner): + return path + return f'{path[: start + 1]}{inner},s:{subtype}{path[end:]}' + + +def to_display_path(qualified_path: str) -> str: + """Convert /{a,t:flow}/{b,t:step} to 'a > b'.""" + if not qualified_path: + return '' + return ' > '.join(_PATH_SEGMENT_RE.findall(qualified_path)) diff --git a/py/packages/genkit/src/genkit/_core/_trace/_realtime_processor.py b/py/packages/genkit/src/genkit/_core/_trace/_realtime_processor.py new file mode 100644 index 0000000000..51d9e9cdce --- /dev/null +++ b/py/packages/genkit/src/genkit/_core/_trace/_realtime_processor.py @@ -0,0 +1,46 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Realtime span processor for live trace visualization.""" + +from opentelemetry.context import Context +from opentelemetry.sdk.trace import Span +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +from genkit._core._compat import override +from genkit._core._logger import get_logger + +logger = get_logger(__name__) + + +class RealtimeSpanProcessor(SimpleSpanProcessor): + """Exports spans on start (real-time) and on end, unlike SimpleSpanProcessor (end only).""" + + @override + def on_start(self, span: Span, parent_context: Context | None = None) -> None: + """Export span immediately so DevUI can show in-progress traces.""" + try: + self.span_exporter.export([span]) + except ConnectionError: + logger.debug( + 'RealtimeSpanProcessor: export failed on_start (collector unreachable)', + exc_info=True, + ) + except Exception: # noqa: BLE001 β€” must never crash the caller + logger.warning( + 'RealtimeSpanProcessor: unexpected error during export on_start', + exc_info=True, + ) diff --git a/py/packages/genkit/src/genkit/core/tracing.py b/py/packages/genkit/src/genkit/_core/_tracing.py similarity index 74% rename from py/packages/genkit/src/genkit/core/tracing.py rename to py/packages/genkit/src/genkit/_core/_tracing.py index 239a32ef1a..529df8cd6d 100644 --- a/py/packages/genkit/src/genkit/core/tracing.py +++ b/py/packages/genkit/src/genkit/_core/_tracing.py @@ -15,17 +15,7 @@ # SPDX-License-Identifier: Apache-2.0 -"""Telemetry and tracing functionality for the Genkit framework. - -This module provides functionality for collecting and exporting telemetry data -from Genkit operations. It uses OpenTelemetry for tracing and exports span -data to a telemetry server for monitoring and debugging purposes. - -The module includes: - - A custom span exporter for sending trace data to a telemetry server - - Utility functions for converting and formatting trace attributes - - Configuration for development environment tracing -""" +"""Telemetry and tracing functionality for the Genkit framework.""" import traceback from collections.abc import Generator @@ -36,13 +26,11 @@ from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SpanExporter -from genkit.core.environment import is_dev_environment -from genkit.core.logging import get_logger -from genkit.core.trace.default_exporter import create_span_processor, init_telemetry_server_exporter -from genkit.core.trace.types import GenkitSpan -from genkit.core.typing import SpanMetadata +from genkit._core._environment import is_dev_environment +from genkit._core._logger import get_logger +from genkit._core._trace._default_exporter import create_span_processor, init_telemetry_server_exporter +from genkit._core._typing import SpanMetadata -ATTR_PREFIX = 'genkit' logger = get_logger(__name__) @@ -99,22 +87,31 @@ def run_in_new_span( metadata: SpanMetadata, labels: dict[str, str] | None = None, links: list[trace_api.Link] | None = None, -) -> Generator[GenkitSpan, None, None]: +) -> Generator[trace_api.Span, None, None]: """Starts a new span context under the current trace. - This method provides a contexmanager for working with Genkit spans. The - context object is a `GenkitSpan`, which is a light wrapper on OpenTelemetry - span object, with handling for genkit attributes. + This method provides a context manager for working with OpenTelemetry spans. + The yielded span is a standard OpenTelemetry Span. Use span.set_attribute() + with 'genkit:' prefix for Genkit-specific attributes. + + Args: + metadata: Span metadata containing the span name. + labels: Optional labels to set as span attributes. + links: Optional span links. + + Yields: + The OpenTelemetry Span object. """ - with tracer.start_as_current_span(name=metadata.name, links=links) as ot_span: - span = GenkitSpan(ot_span, labels) + with tracer.start_as_current_span(name=metadata.name, links=links) as span: + if labels is not None: + span.set_attributes(labels) try: yield span - span.set_genkit_attribute('status', 'success') + span.set_attribute('genkit:state', 'success') except Exception as e: logger.debug(f'Error in run_in_new_span: {e!s}') logger.debug(traceback.format_exc()) - span.set_genkit_attribute('status', 'error') + span.set_attribute('genkit:state', 'error') span.set_status(status=trace_api.StatusCode.ERROR, description=str(e)) span.record_exception(e) raise e diff --git a/py/packages/genkit/src/genkit/core/typing.py b/py/packages/genkit/src/genkit/_core/_typing.py similarity index 81% rename from py/packages/genkit/src/genkit/core/typing.py rename to py/packages/genkit/src/genkit/_core/_typing.py index f60ed30169..0263ed537b 100644 --- a/py/packages/genkit/src/genkit/core/typing.py +++ b/py/packages/genkit/src/genkit/_core/_typing.py @@ -25,12 +25,22 @@ from __future__ import annotations +import warnings from typing import Any, ClassVar, Literal -from pydantic import BaseModel, ConfigDict, Field, RootModel +from pydantic import ConfigDict, Field, RootModel from pydantic.alias_generators import to_camel +from strenum import StrEnum -from genkit.core._compat import StrEnum +from genkit._core._base import GenkitModel + +# 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, +) class Model(RootModel[Any]): @@ -39,7 +49,7 @@ class Model(RootModel[Any]): root: Any -class Embedding(BaseModel): +class Embedding(GenkitModel): """Model for embedding data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -47,7 +57,7 @@ class Embedding(BaseModel): metadata: dict[str, Any] | None = None -class BaseDataPoint(BaseModel): +class BaseDataPoint(GenkitModel): """Model for basedatapoint data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -59,7 +69,7 @@ class BaseDataPoint(BaseModel): trace_ids: list[str] | None = Field(default=None) -class EvalRequest(BaseModel): +class EvalRequest(GenkitModel): """Model for evalrequest data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -76,14 +86,14 @@ class EvalStatusEnum(StrEnum): FAIL = 'FAIL' -class Details(BaseModel): +class Details(GenkitModel): """Model for details data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) reasoning: str | None = None -class Score(BaseModel): +class Score(GenkitModel): """Model for score data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -94,7 +104,7 @@ class Score(BaseModel): details: Details | None = None -class GenkitErrorDetails(BaseModel): +class GenkitErrorDetails(GenkitModel): """Model for genkiterrordetails data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -102,7 +112,7 @@ class GenkitErrorDetails(BaseModel): trace_id: str = Field(...) -class Data(BaseModel): +class Data(GenkitModel): """Model for data data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -110,7 +120,7 @@ class Data(BaseModel): genkit_error_details: GenkitErrorDetails | None = Field(default=None) -class GenkitError(BaseModel): +class GenkitError(GenkitModel): """Model for genkiterror data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -128,7 +138,7 @@ class Code(StrEnum): UNKNOWN = 'unknown' -class CandidateError(BaseModel): +class CandidateError(GenkitModel): """Model for candidateerror data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -137,7 +147,7 @@ class CandidateError(BaseModel): message: str | None = None -class CustomPart(BaseModel): +class CustomPart(GenkitModel): """Model for custompart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -171,7 +181,7 @@ class ToolChoice(StrEnum): NONE = 'none' -class GenerateActionOutputConfig(BaseModel): +class GenerateActionOutputConfig(GenkitModel): """Model for generateactionoutputconfig data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -184,7 +194,7 @@ class GenerateActionOutputConfig(BaseModel): schema_type: Any = Field(default=None, exclude=True) -class GenerationCommonConfig(BaseModel): +class GenerationCommonConfig(GenkitModel): """Model for generationcommonconfig data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) @@ -212,7 +222,7 @@ class GenerationCommonConfig(BaseModel): ) -class GenerationUsage(BaseModel): +class GenerationUsage(GenkitModel): """Model for generationusage data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -232,7 +242,7 @@ class GenerationUsage(BaseModel): cached_content_tokens: float | None = Field(default=None) -class MiddlewareDesc(BaseModel): +class MiddlewareDesc(GenkitModel): """Model for middlewaredesc data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -242,7 +252,7 @@ class MiddlewareDesc(BaseModel): metadata: dict[str, Any] | None = None -class MiddlewareRef(BaseModel): +class MiddlewareRef(GenkitModel): """Model for middlewareref data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -258,7 +268,7 @@ class Constrained(StrEnum): NO_TOOLS = 'no-tools' -class Supports(BaseModel): +class Supports(GenkitModel): """Model for supports data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -284,7 +294,7 @@ class Stage(StrEnum): DEPRECATED = 'deprecated' -class ModelInfo(BaseModel): +class ModelInfo(GenkitModel): """Model for modelinfo data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -295,7 +305,7 @@ class ModelInfo(BaseModel): stage: Stage | None = None -class ModelReference(BaseModel): +class ModelReference(GenkitModel): """Model for modelreference data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -306,14 +316,14 @@ class ModelReference(BaseModel): config: Any | None = None -class Error(BaseModel): +class Error(GenkitModel): """Model for error data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) message: str -class Operation(BaseModel): +class Operation(GenkitModel): """Model for operation data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -325,7 +335,7 @@ class Operation(BaseModel): metadata: dict[str, Any] | None = None -class OutputConfig(BaseModel): +class OutputConfig(GenkitModel): """Model for outputconfig data.""" model_config: ClassVar[ConfigDict] = ConfigDict( @@ -339,7 +349,7 @@ class OutputConfig(BaseModel): content_type: str | None = Field(default=None) -class Resource1(BaseModel): +class Resource1(GenkitModel): """Model for resource1 data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -355,7 +365,7 @@ class Role(StrEnum): TOOL = 'tool' -class ToolDefinition(BaseModel): +class ToolDefinition(GenkitModel): """Model for tooldefinition data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -370,7 +380,7 @@ class ToolDefinition(BaseModel): metadata: dict[str, Any] | None = Field(default=None, description='additional metadata for this tool definition') -class Media(BaseModel): +class Media(GenkitModel): """Model for media data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -378,7 +388,7 @@ class Media(BaseModel): url: str -class ToolRequest(BaseModel): +class ToolRequest(GenkitModel): """Model for toolrequest data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -388,7 +398,7 @@ class ToolRequest(BaseModel): partial: bool | None = None -class ToolResponse(BaseModel): +class ToolResponse(GenkitModel): """Model for toolresponse data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -398,28 +408,7 @@ class ToolResponse(BaseModel): content: list[Any] | None = None -class CommonRerankerOptions(BaseModel): - """Model for commonrerankeroptions data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - k: float | None = Field(default=None, description='Number of documents to rerank') - - -class RankedDocumentMetadata(BaseModel): - """Model for rankeddocumentmetadata data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) - score: float - - -class CommonRetrieverOptions(BaseModel): - """Model for commonretrieveroptions data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - k: float | None = Field(default=None, description='Number of documents to retrieve') - - -class InstrumentationLibrary(BaseModel): +class InstrumentationLibrary(GenkitModel): """Model for instrumentationlibrary data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -428,7 +417,7 @@ class InstrumentationLibrary(BaseModel): schema_url: str | None = Field(default=None) -class PathMetadata(BaseModel): +class PathMetadata(GenkitModel): """Model for pathmetadata data.""" model_config: ClassVar[ConfigDict] = ConfigDict( @@ -440,7 +429,7 @@ class PathMetadata(BaseModel): latency: float -class SpanContext(BaseModel): +class SpanContext(GenkitModel): """Model for spancontext data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -450,7 +439,7 @@ class SpanContext(BaseModel): trace_flags: float = Field(...) -class SameProcessAsParentSpan(BaseModel): +class SameProcessAsParentSpan(GenkitModel): """Model for sameprocessasparentspan data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -464,7 +453,7 @@ class State(StrEnum): ERROR = 'error' -class SpanMetadata(BaseModel): +class SpanMetadata(GenkitModel): """Model for spanmetadata data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -477,7 +466,7 @@ class SpanMetadata(BaseModel): path: str | None = None -class SpanStatus(BaseModel): +class SpanStatus(GenkitModel): """Model for spanstatus data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -485,7 +474,7 @@ class SpanStatus(BaseModel): message: str | None = None -class Annotation(BaseModel): +class Annotation(GenkitModel): """Model for annotation data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -493,7 +482,7 @@ class Annotation(BaseModel): description: str -class TimeEvent(BaseModel): +class TimeEvent(GenkitModel): """Model for timeevent data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -501,7 +490,7 @@ class TimeEvent(BaseModel): annotation: Annotation -class TraceMetadata(BaseModel): +class TraceMetadata(GenkitModel): """Model for tracemetadata data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -600,12 +589,6 @@ class Config(RootModel[Any]): root: Any -class OutputModel(RootModel[OutputConfig]): - """Root model for outputmodel.""" - - root: OutputConfig - - class Tools(RootModel[list[ToolDefinition]]): """Root model for tools.""" @@ -660,14 +643,14 @@ class TraceId(RootModel[str]): root: str -class EmbedResponse(BaseModel): +class EmbedResponse(GenkitModel): """Model for embedresponse data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) embeddings: list[Embedding] -class BaseEvalDataPoint(BaseModel): +class BaseEvalDataPoint(GenkitModel): """Model for baseevaldatapoint data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -679,7 +662,7 @@ class BaseEvalDataPoint(BaseModel): trace_ids: TraceIds | None = Field(default=None) -class EvalFnResponse(BaseModel): +class EvalFnResponse(GenkitModel): """Model for evalfnresponse data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -696,7 +679,7 @@ class EvalResponse(RootModel[list[EvalFnResponse]]): root: list[EvalFnResponse] -class DataPart(BaseModel): +class DataPart(GenkitModel): """Model for datapart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -711,7 +694,7 @@ class DataPart(BaseModel): resource: Resource | None = None -class MediaPart(BaseModel): +class MediaPart(GenkitModel): """Model for mediapart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -726,7 +709,7 @@ class MediaPart(BaseModel): resource: Resource | None = None -class ReasoningPart(BaseModel): +class ReasoningPart(GenkitModel): """Model for reasoningpart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -741,7 +724,7 @@ class ReasoningPart(BaseModel): resource: Resource | None = None -class ResourcePart(BaseModel): +class ResourcePart(GenkitModel): """Model for resourcepart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -756,7 +739,7 @@ class ResourcePart(BaseModel): resource: Resource1 -class TextPart(BaseModel): +class TextPart(GenkitModel): """Model for textpart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -771,7 +754,7 @@ class TextPart(BaseModel): resource: Resource | None = None -class ToolRequestPart(BaseModel): +class ToolRequestPart(GenkitModel): """Model for toolrequestpart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -786,7 +769,7 @@ class ToolRequestPart(BaseModel): resource: Resource | None = None -class ToolResponsePart(BaseModel): +class ToolResponsePart(GenkitModel): """Model for toolresponsepart data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -801,7 +784,7 @@ class ToolResponsePart(BaseModel): resource: Resource | None = None -class Link(BaseModel): +class Link(GenkitModel): """Model for link data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -810,14 +793,14 @@ class Link(BaseModel): dropped_attributes_count: float | None = Field(default=None) -class TimeEvents(BaseModel): +class TimeEvents(GenkitModel): """Model for timeevents data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) time_event: list[TimeEvent] | None = Field(default=None) -class SpanData(BaseModel): +class SpanData(GenkitModel): """Model for spandata data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -837,7 +820,7 @@ class SpanData(BaseModel): truncated: bool | None = None -class SpanEndEvent(BaseModel): +class SpanEndEvent(GenkitModel): """Model for spanendevent data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -846,7 +829,7 @@ class SpanEndEvent(BaseModel): type: Literal['span_end'] -class SpanStartEvent(BaseModel): +class SpanStartEvent(GenkitModel): """Model for spanstartevent data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -855,7 +838,7 @@ class SpanStartEvent(BaseModel): type: Literal['span_start'] -class SpantEventBase(BaseModel): +class SpantEventBase(GenkitModel): """Model for spanteventbase data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -863,7 +846,7 @@ class SpantEventBase(BaseModel): span: SpanData -class TraceData(BaseModel): +class TraceData(GenkitModel): """Model for tracedata data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -886,7 +869,7 @@ class DocumentPart(RootModel[TextPart | MediaPart]): root: TextPart | MediaPart -class Resume(BaseModel): +class Resume(GenkitModel): """Model for resume data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -907,28 +890,13 @@ class Part( ) -class RankedDocumentData(BaseModel): - """Model for rankeddocumentdata data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - content: list[DocumentPart] - metadata: RankedDocumentMetadata - - -class RerankerResponse(BaseModel): - """Model for rerankerresponse data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - documents: list[RankedDocumentData] - - class Content(RootModel[list[Part]]): """Root model for content.""" root: list[Part] -class DocumentData(BaseModel): +class DocumentData(GenkitModel): """Model for documentdata data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -936,7 +904,7 @@ class DocumentData(BaseModel): metadata: dict[str, Any] | None = None -class EmbedRequest(BaseModel): +class EmbedRequest(GenkitModel): """Model for embedrequest data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -944,7 +912,7 @@ class EmbedRequest(BaseModel): options: Any | None = None -class GenerateResponseChunk(BaseModel): +class GenerateResponseChunk(GenkitModel): """Model for generateresponsechunk data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -955,8 +923,8 @@ class GenerateResponseChunk(BaseModel): aggregated: bool | None = None -class Message(BaseModel): - """Model for message data.""" +class MessageData(GenkitModel): + """Model for messagedata data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) role: Role | str @@ -964,7 +932,7 @@ class Message(BaseModel): metadata: dict[str, Any] | None = None -class ModelResponseChunk(BaseModel): +class ModelResponseChunk(GenkitModel): """Model for modelresponsechunk data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -975,7 +943,7 @@ class ModelResponseChunk(BaseModel): aggregated: Aggregated | None = None -class MultipartToolResponse(BaseModel): +class MultipartToolResponse(GenkitModel): """Model for multiparttoolresponse data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) @@ -984,61 +952,37 @@ class MultipartToolResponse(BaseModel): metadata: dict[str, Any] | None = None -class RerankerRequest(BaseModel): - """Model for rerankerrequest data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - query: DocumentData - documents: list[DocumentData] - options: Any | None = None - - -class RetrieverRequest(BaseModel): - """Model for retrieverrequest data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - query: DocumentData - options: Any | None = None - - -class RetrieverResponse(BaseModel): - """Model for retrieverresponse data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - documents: list[DocumentData] - - class Docs(RootModel[list[DocumentData]]): """Root model for docs.""" root: list[DocumentData] -class Messages(RootModel[list[Message]]): +class Messages(RootModel[list[MessageData]]): """Root model for messages.""" - root: list[Message] + root: list[MessageData] -class Candidate(BaseModel): +class Candidate(GenkitModel): """Model for candidate data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) index: float - message: Message + message: MessageData usage: GenerationUsage | None = None finish_reason: FinishReason = Field(...) finish_message: str | None = Field(default=None) custom: Any | None = None -class GenerateActionOptions(BaseModel): +class GenerateActionOptions(GenkitModel): """Model for generateactionoptions data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) model: str | None = None docs: list[DocumentData] | None = None - messages: list[Message] + messages: list[MessageData] tools: list[str] | None = None resources: list[str] | None = None tool_choice: ToolChoice | None = Field(default=None) @@ -1049,65 +993,3 @@ class GenerateActionOptions(BaseModel): max_turns: float | None = Field(default=None) step_name: str | None = Field(default=None) use: list[MiddlewareRef] | None = None - - -class GenerateRequest(BaseModel): - """Model for generaterequest data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - messages: list[Message] - config: Any | None = None - tools: list[ToolDefinition] | None = None - tool_choice: ToolChoice | None = Field(default=None) - output: OutputConfig | None = None - docs: list[DocumentData] | None = None - candidates: float | None = None - - -class GenerateResponse(BaseModel): - """Model for generateresponse data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - message: Message | None = None - finish_reason: FinishReason | None = Field(default=None) - finish_message: str | None = Field(default=None) - latency_ms: float | None = Field(default=None) - usage: GenerationUsage | None = None - custom: Any | None = None - raw: Any | None = None - request: GenerateRequest | None = None - operation: Operation | None = None - candidates: list[Candidate] | None = None - - -class ModelRequest(BaseModel): - """Model for modelrequest data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - messages: Messages - config: Config | None = None - tools: Tools | None = None - tool_choice: ToolChoice | None = Field(default=None) - output: OutputModel | None = None - docs: Docs | None = None - - -class Request(RootModel[GenerateRequest]): - """Root model for request.""" - - root: GenerateRequest - - -class ModelResponse(BaseModel): - """Model for modelresponse data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - message: Message | None = None - finish_reason: FinishReason = Field(...) - finish_message: FinishMessage | None = Field(default=None) - latency_ms: LatencyMs | None = Field(default=None) - usage: Usage | None = None - custom: CustomModel | None = None - raw: Raw | None = None - request: Request | None = None - operation: Operation | None = None diff --git a/py/packages/genkit/src/genkit/ai/__init__.py b/py/packages/genkit/src/genkit/ai/__init__.py deleted file mode 100644 index c09261ffbe..0000000000 --- a/py/packages/genkit/src/genkit/ai/__init__.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Genkit AI module - core Genkit class and related utilities. - -This module provides the main Genkit class for building AI applications, -along with essential types for actions, sessions, prompts, and tools. - -Overview: - The ``genkit.ai`` module exposes the primary entry points for building - AI-powered applications. The ``Genkit`` class is the main orchestrator - that manages plugins, models, prompts, flows, and sessions. - -Terminology: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Term β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Genkit β”‚ Main orchestrator class for AI applications β”‚ - β”‚ Plugin β”‚ Extension that adds models, embedders, or other β”‚ - β”‚ β”‚ capabilities (e.g., GoogleAI, Anthropic) β”‚ - β”‚ Flow β”‚ A durable, traceable function for AI workflows β”‚ - β”‚ Action β”‚ A strongly-typed, remotely callable function β”‚ - β”‚ Tool β”‚ A function callable by models during generation β”‚ - β”‚ Session/Chat β”‚ State management for multi-turn conversations β”‚ - β”‚ Prompt β”‚ Reusable template for model interactions β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Components: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Component β”‚ Purpose β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Genkit β”‚ Main class for building AI applications β”‚ - β”‚ GenkitRegistry β”‚ Internal registry for actions and resources β”‚ - β”‚ ExecutablePrompt β”‚ Callable prompt with template rendering β”‚ - β”‚ Chat β”‚ Stateful multi-turn conversation interface β”‚ - β”‚ ActionRunContext β”‚ Execution context for actions (streaming) β”‚ - β”‚ ToolRunContext β”‚ Execution context for tools (with interrupt) β”‚ - β”‚ GenerateResponseWrapper β”‚ Enhanced response with helper methods β”‚ - β”‚ GenerateStreamResponse β”‚ Streaming response with chunks and final resp β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Basic usage: - - ```python - from genkit import Genkit - from genkit.plugins.google_genai import GoogleAI - - # Initialize with plugins - ai = Genkit(plugins=[GoogleAI()], model='googleai/gemini-2.0-flash') - - - # Define a flow - @ai.flow() - async def hello(name: str) -> str: - response = await ai.generate(prompt=f'Say hello to {name}') - return response.text - - - # Define a tool - @ai.tool() - def get_weather(city: str) -> str: - return f'Weather in {city}: Sunny, 72Β°F' - - - # Create a chat session - chat = ai.chat(system='You are a helpful assistant.') - response = await chat.send('What is the weather in Paris?') - ``` - -See Also: - - Genkit documentation: https://genkit.dev/ - - JavaScript SDK: https://github.com/firebase/genkit -""" - -from genkit.blocks.document import Document -from genkit.blocks.interfaces import Input -from genkit.blocks.model import GenerateResponseWrapper -from genkit.blocks.prompt import ( - ExecutablePrompt, - GenerateStreamResponse, - OutputOptions, - PromptGenerateOptions, - ResumeOptions, -) -from genkit.blocks.tools import ToolRunContext, tool_response -from genkit.core import GENKIT_CLIENT_HEADER, GENKIT_VERSION -from genkit.core.action import ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.plugin import Plugin - -from ._aio import Genkit, Output -from ._registry import FlowWrapper, GenkitRegistry, SimpleRetrieverOptions - -__all__ = [ - # Version info - 'GENKIT_CLIENT_HEADER', - 'GENKIT_VERSION', - # Main class - 'Genkit', - 'Input', - 'Output', - # Actions - 'ActionKind', - 'ActionRunContext', - # Document - 'Document', - # Prompts - 'ExecutablePrompt', - 'OutputOptions', - 'PromptGenerateOptions', - 'ResumeOptions', - # Registry and flow - 'FlowWrapper', - 'GenkitRegistry', - 'SimpleRetrieverOptions', - # Response types - 'GenerateResponseWrapper', - 'GenerateStreamResponse', - # Tools - 'ToolRunContext', - 'tool_response', - # Plugin - 'Plugin', -] diff --git a/py/packages/genkit/src/genkit/ai/_aio.py b/py/packages/genkit/src/genkit/ai/_aio.py deleted file mode 100644 index fb8e6ee65a..0000000000 --- a/py/packages/genkit/src/genkit/ai/_aio.py +++ /dev/null @@ -1,1164 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""User-facing asyncio API for Genkit. - -This module provides the primary entry point for using Genkit in an asynchronous -environment. The `Genkit` class coordinates plugins, registry, and execution -of AI actions like generation, embedding, and retrieval. - -Key features provided by the `Genkit` class: -- **Generation**: Interface for unified model interaction via `generate` and `generate_stream`. -- **Flow Control**: Execution of granular steps with tracing via `run`. -- **Dynamic Extensibility**: On-the-fly creation of tools via `dynamic_tool`. -- **Observability**: Specialized methods for managing trace context and flushing telemetry. -""" - -from __future__ import annotations - -import asyncio -import uuid -from collections.abc import AsyncIterator, Awaitable, Callable -from pathlib import Path -from typing import Any, TypeVar, cast, overload - -from opentelemetry import trace as trace_api -from opentelemetry.sdk.trace import TracerProvider - -from genkit.aio._util import ensure_async -from genkit.aio.channel import Channel -from genkit.blocks.background_model import ( - check_operation as check_operation_impl, - lookup_background_action, -) -from genkit.blocks.document import Document -from genkit.blocks.embedding import EmbedderRef -from genkit.blocks.evaluator import EvaluatorRef -from genkit.blocks.generate import ( - StreamingCallback as ModelStreamingCallback, - generate_action, -) -from genkit.blocks.interfaces import Output, OutputConfigDict -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - GenerateResponseWrapper, - ModelMiddleware, -) -from genkit.blocks.prompt import PromptConfig, load_prompt_folder, to_generate_action_options -from genkit.blocks.retriever import IndexerRef, IndexerRequest, RetrieverRef -from genkit.core.action import Action, ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.error import GenkitError -from genkit.core.plugin import Plugin -from genkit.core.tracing import run_in_new_span -from genkit.core.typing import ( - BaseDataPoint, - Embedding, - EmbedRequest, - EvalRequest, - EvalResponse, - Operation, - SpanMetadata, -) -from genkit.types import ( - DocumentData, - GenerationCommonConfig, - Message, - OutputConfig, - Part, - RetrieverRequest, - RetrieverResponse, - ToolChoice, -) - -from ._base_async import GenkitBase -from ._server import ServerSpec - -T = TypeVar('T') -InputT = TypeVar('InputT') -OutputT = TypeVar('OutputT') - - -class Genkit(GenkitBase): - """Genkit asyncio user-facing API.""" - - def __init__( - self, - plugins: list[Plugin] | None = None, - model: str | None = None, - prompt_dir: str | Path | None = None, - reflection_server_spec: ServerSpec | None = None, - ) -> None: - """Initialize a new Genkit instance. - - Args: - plugins: List of plugins to initialize. - model: Model name to use. - prompt_dir: Directory to automatically load prompts from. - If not provided, defaults to loading from './prompts' if it exists. - reflection_server_spec: Server spec for the reflection - server. - """ - super().__init__(plugins=plugins, model=model, reflection_server_spec=reflection_server_spec) - - load_path = prompt_dir - if load_path is None: - default_prompts_path = Path('./prompts') - if default_prompts_path.is_dir(): - load_path = default_prompts_path - - if load_path: - load_prompt_folder(self.registry, dir_path=load_path) - - def _resolve_embedder_name(self, embedder: str | EmbedderRef | None) -> str: - """Resolve embedder name from string or EmbedderRef. - - Args: - embedder: The embedder specified as a string name or EmbedderRef. - - Returns: - The resolved embedder name. - - Raises: - ValueError: If embedder is not specified or is of invalid type. - """ - if isinstance(embedder, EmbedderRef): - return embedder.name - elif isinstance(embedder, str): - return embedder - else: - raise ValueError('Embedder must be specified as a string name or an EmbedderRef.') - - @overload - async def generate( - self, - model: str | None = None, - prompt: str | Part | list[Part] | None = None, - system: str | Part | list[Part] | None = None, - messages: list[Message] | None = None, - tools: list[str] | None = None, - return_tool_requests: bool | None = None, - tool_choice: ToolChoice | None = None, - tool_responses: list[Part] | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - max_turns: int | None = None, - on_chunk: ModelStreamingCallback | None = None, - context: dict[str, object] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_constrained: bool | None = None, - *, - output: Output[OutputT], - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | None = None, - ) -> GenerateResponseWrapper[OutputT]: ... - - @overload - async def generate( - self, - model: str | None = None, - prompt: str | Part | list[Part] | None = None, - system: str | Part | list[Part] | None = None, - messages: list[Message] | None = None, - tools: list[str] | None = None, - return_tool_requests: bool | None = None, - tool_choice: ToolChoice | None = None, - tool_responses: list[Part] | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - max_turns: int | None = None, - on_chunk: ModelStreamingCallback | None = None, - context: dict[str, object] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_constrained: bool | None = None, - output: OutputConfig | OutputConfigDict | Output[Any] | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | None = None, - ) -> GenerateResponseWrapper[Any]: ... - - async def generate( - self, - model: str | None = None, - prompt: str | Part | list[Part] | None = None, - system: str | Part | list[Part] | None = None, - messages: list[Message] | None = None, - tools: list[str] | None = None, - return_tool_requests: bool | None = None, - tool_choice: ToolChoice | None = None, - tool_responses: list[Part] | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - max_turns: int | None = None, - on_chunk: ModelStreamingCallback | None = None, - context: dict[str, object] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_constrained: bool | None = None, - output: OutputConfig | OutputConfigDict | Output[Any] | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | None = None, - ) -> GenerateResponseWrapper[Any]: - """Generates text or structured data using a language model. - - This function provides a flexible interface for interacting with various - language models, supporting both simple text generation and more complex - interactions involving tools and structured conversations. - - Args: - model: Optional. The name of the model to use for generation. If not - provided, a default model may be used. - prompt: Optional. A single prompt string, a `Part` object, or a list - of `Part` objects to provide as input to the model. This is used - for simple text generation. - system: Optional. A system message string, a `Part` object, or a - list of `Part` objects to provide context or instructions to - the model, especially for chat-based models. - messages: Optional. A list of `Message` objects representing a - conversation history. This is used for chat-based models to - maintain context. - tools: Optional. A list of tool names (strings) that the model can - use. - return_tool_requests: Optional. If `True`, the model will return - tool requests instead of executing them directly. - tool_choice: Optional. A `ToolChoice` object specifying how the - model should choose which tool to use. - tool_responses: Optional. tool_responses should contain a list of - tool response parts corresponding to interrupt tool request - parts from the most recent model message. Each entry must have - a matching `name` and `ref` (if supplied) for its tool request - counterpart. - config: Optional. A `GenerationCommonConfig` object or a dictionary - containing configuration parameters for the generation process. - This allows fine-tuning the model's behavior. - max_turns: Optional. The maximum number of turns in a conversation. - on_chunk: Optional. A callback function of type - `ModelStreamingCallback` that is called for each chunk of - generated text during streaming. - context: Optional. A dictionary containing additional context - information that can be used during generation. - output_format: Optional. The format to use for the output (e.g., - 'json'). - output_content_type: Optional. The content type of the output. - output_instructions: Optional. Instructions for formatting the - output. - output_constrained: Optional. Whether to constrain the output to the - schema. - output: Optional. Use `Output(schema=YourSchema)` for typed responses. - Can also be an `OutputConfig` object or dictionary. - use: Optional. A list of `ModelMiddleware` functions to apply to the - generation process. Middleware can be used to intercept and - modify requests and responses. - docs: Optional. A list of documents to be used for grounding. - - - Returns: - A `GenerateResponseWrapper` object containing the model's response, - which may include generated text, tool requests, or other relevant - information. - - Note: - - The `tools`, `return_tool_requests`, and `tool_choice` arguments - are used for models that support tool usage. - - The `on_chunk` argument enables streaming responses, allowing you - to process the generated content as it becomes available. - """ - # Initialize output_schema - extracted from output parameter - output_schema: type | dict[str, object] | None = None - - # Unpack output config if provided - if output: - if isinstance(output, Output): - # Handle typed Output[T] - extract values from the typed wrapper - if output_format is None: - output_format = output.format - if output_content_type is None: - output_content_type = output.content_type - if output_instructions is None: - output_instructions = output.instructions - if output_schema is None: - output_schema = output.schema - if output_constrained is None: - output_constrained = output.constrained - elif isinstance(output, dict): - # Handle dict input - extract values directly - if output_format is None: - output_format = output.get('format') - if output_content_type is None: - output_content_type = output.get('content_type') - if output_instructions is None: - output_instructions = output.get('instructions') - if output_schema is None: - output_schema = output.get('schema') - if output_constrained is None: - output_constrained = output.get('constrained') - else: - # Handle OutputConfig object - use getattr for safety since - # OutputConfig is auto-generated and may not have all fields. - if output_format is None: - output_format = getattr(output, 'format', None) - if output_content_type is None: - output_content_type = getattr(output, 'content_type', None) - if output_instructions is None: - output_instructions = getattr(output, 'instructions', None) - if output_schema is None: - output_schema = getattr(output, 'schema', None) - if output_constrained is None: - output_constrained = getattr(output, 'constrained', None) - - return await generate_action( - self.registry, - await to_generate_action_options( - self.registry, - PromptConfig( - model=model, - prompt=prompt, - system=system, - messages=messages, - tools=tools, - return_tool_requests=return_tool_requests, - tool_choice=tool_choice, - tool_responses=tool_responses, - config=config, - max_turns=max_turns, - output_format=output_format, - output_content_type=output_content_type, - output_instructions=output_instructions, - output_schema=output_schema, - output_constrained=output_constrained, - docs=docs, - ), - ), - on_chunk=on_chunk, - middleware=use, - context=context if context else ActionRunContext._current_context(), # pyright: ignore[reportPrivateUsage] - ) - - @overload - def generate_stream( - self, - model: str | None = None, - prompt: str | Part | list[Part] | None = None, - system: str | Part | list[Part] | None = None, - messages: list[Message] | None = None, - tools: list[str] | None = None, - return_tool_requests: bool | None = None, - tool_choice: ToolChoice | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - max_turns: int | None = None, - context: dict[str, object] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_constrained: bool | None = None, - *, - output: Output[OutputT], - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | None = None, - timeout: float | None = None, - ) -> tuple[ - AsyncIterator[GenerateResponseChunkWrapper], - asyncio.Future[GenerateResponseWrapper[OutputT]], - ]: ... - - @overload - def generate_stream( - self, - model: str | None = None, - prompt: str | Part | list[Part] | None = None, - system: str | Part | list[Part] | None = None, - messages: list[Message] | None = None, - tools: list[str] | None = None, - return_tool_requests: bool | None = None, - tool_choice: ToolChoice | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - max_turns: int | None = None, - context: dict[str, object] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_constrained: bool | None = None, - output: OutputConfig | OutputConfigDict | Output[Any] | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | None = None, - timeout: float | None = None, - ) -> tuple[ - AsyncIterator[GenerateResponseChunkWrapper], - asyncio.Future[GenerateResponseWrapper[Any]], - ]: ... - - def generate_stream( - self, - model: str | None = None, - prompt: str | Part | list[Part] | None = None, - system: str | Part | list[Part] | None = None, - messages: list[Message] | None = None, - tools: list[str] | None = None, - return_tool_requests: bool | None = None, - tool_choice: ToolChoice | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - max_turns: int | None = None, - context: dict[str, object] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_constrained: bool | None = None, - output: OutputConfig | OutputConfigDict | Output[Any] | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | None = None, - timeout: float | None = None, - ) -> tuple[ - AsyncIterator[GenerateResponseChunkWrapper], - asyncio.Future[GenerateResponseWrapper[Any]], - ]: - """Streams generated text or structured data using a language model. - - This function provides a flexible interface for interacting with various - language models, supporting both simple text generation and more complex - interactions involving tools and structured conversations. - - Args: - model: Optional. The name of the model to use for generation. If not - provided, a default model may be used. - prompt: Optional. A single prompt string, a `Part` object, or a list - of `Part` objects to provide as input to the model. This is used - for simple text generation. - system: Optional. A system message string, a `Part` object, or a - list of `Part` objects to provide context or instructions to the - model, especially for chat-based models. - messages: Optional. A list of `Message` objects representing a - conversation history. This is used for chat-based models to - maintain context. - tools: Optional. A list of tool names (strings) that the model can - use. - return_tool_requests: Optional. If `True`, the model will return - tool requests instead of executing them directly. - tool_choice: Optional. A `ToolChoice` object specifying how the - model should choose which tool to use. - config: Optional. A `GenerationCommonConfig` object or a dictionary - containing configuration parameters for the generation process. - This allows fine-tuning the model's behavior. - max_turns: Optional. The maximum number of turns in a conversation. - context: Optional. A dictionary containing additional context - information that can be used during generation. - output_format: Optional. The format to use for the output (e.g., - 'json'). - output_content_type: Optional. The content type of the output. - output_instructions: Optional. Instructions for formatting the - output. - output_constrained: Optional. Whether to constrain the output to the - schema. - output: Optional. Use `Output(schema=YourSchema)` for typed responses. - Can also be an `OutputConfig` object or dictionary. - use: Optional. A list of `ModelMiddleware` functions to apply to the - generation process. Middleware can be used to intercept and - modify requests and responses. - docs: Optional. A list of documents to be used for grounding. - timeout: Optional. The timeout for the streaming action. - - Returns: - A `GenerateResponseWrapper` object containing the model's response, - which may include generated text, tool requests, or other relevant - information. - - Note: - - The `tools`, `return_tool_requests`, and `tool_choice` arguments - are used for models that support tool usage. - - The `on_chunk` argument enables streaming responses, allowing you - to process the generated content as it becomes available. - """ - stream: Channel[GenerateResponseChunkWrapper, GenerateResponseWrapper[Any]] = Channel(timeout=timeout) - - resp = self.generate( - model=model, - prompt=prompt, - system=system, - messages=messages, - tools=tools, - return_tool_requests=return_tool_requests, - tool_choice=tool_choice, - config=config, - max_turns=max_turns, - context=context, - output_format=output_format, - output_content_type=output_content_type, - output_instructions=output_instructions, - output_constrained=output_constrained, - output=output, - docs=docs, - use=use, - on_chunk=lambda c: stream.send(c), - ) - stream.set_close_future(asyncio.create_task(resp)) - - return stream, stream.closed - - async def retrieve( - self, - retriever: str | RetrieverRef | None = None, - query: str | DocumentData | None = None, - options: dict[str, object] | None = None, - ) -> RetrieverResponse: - """Retrieves documents based on query. - - Args: - retriever: Optional retriever name or reference to use. - query: Text query or a DocumentData containing query text. - options: Optional retriever-specific options. - - Returns: - The generated response with documents. - """ - retriever_name: str - retriever_config: dict[str, object] = {} - - if isinstance(retriever, RetrieverRef): - retriever_name = retriever.name - retriever_config = retriever.config or {} - if retriever.version: - retriever_config['version'] = retriever.version - elif isinstance(retriever, str): - retriever_name = retriever - else: - raise ValueError('Retriever must be specified as a string name or a RetrieverRef.') - - if isinstance(query, str): - query = Document.from_text(query) - - request_options = {**(retriever_config or {}), **(options or {})} - - retrieve_action = await self.registry.resolve_retriever(retriever_name) - if retrieve_action is None: - raise ValueError(f'Retriever "{retriever_name}" not found') - - if query is None: - raise ValueError('Query must be specified for retrieval.') - - return ( - await retrieve_action.arun( - RetrieverRequest( - query=query, - options=request_options if request_options else None, - ) - ) - ).response - - async def index( - self, - indexer: str | IndexerRef | None = None, - documents: list[Document] | None = None, - options: dict[str, object] | None = None, - ) -> None: - """Indexes documents. - - Args: - indexer: Optional indexer name or reference to use. - documents: Documents to index. - options: Optional indexer-specific options. - """ - indexer_name: str - indexer_config: dict[str, object] = {} - - if isinstance(indexer, IndexerRef): - indexer_name = indexer.name - indexer_config = indexer.config or {} - if indexer.version: - indexer_config['version'] = indexer.version - elif isinstance(indexer, str): - indexer_name = indexer - else: - raise ValueError('Indexer must be specified as a string name or an IndexerRef.') - - req_options = {**(indexer_config or {}), **(options or {})} - - index_action = await self.registry.resolve_action(cast(ActionKind, ActionKind.INDEXER), indexer_name) - if index_action is None: - raise ValueError(f'Indexer "{indexer_name}" not found') - - if documents is None: - raise ValueError('Documents must be specified for indexing.') - - _ = await index_action.arun( - IndexerRequest( - # Document subclasses DocumentData, so this is type-safe at runtime. - # list is invariant so list[Document] isn't assignable to list[DocumentData] - documents=cast(list[DocumentData], documents), - options=req_options if req_options else None, - ) - ) - - async def embed( - self, - embedder: str | EmbedderRef | None = None, - content: str | Document | DocumentData | None = None, - metadata: dict[str, object] | None = None, - options: dict[str, object] | None = None, - ) -> list[Embedding]: - """Embeds a single document or string. - - Generates vector embeddings for a single piece of content using the - specified embedder. This is the primary method for embedding individual - items. - - When using an EmbedderRef, the config and version from the ref are - extracted and merged with any provided options. The merge order is: - {version, ...config, ...options} (options take precedence). - - Args: - embedder: Embedder name (e.g., 'googleai/text-embedding-004') or - an EmbedderRef with configuration. - content: A single string, Document, or DocumentData to embed. - metadata: Optional metadata to apply to the document. Only used - when content is a string. - options: Optional embedder-specific options (e.g., task_type). - - Returns: - A list containing the Embedding for the input content. - - Raises: - ValueError: If embedder is not specified or not found. - ValueError: If content is not specified. - - Example - Basic string embedding: - >>> embeddings = await ai.embed(embedder='googleai/text-embedding-004', content='Hello, world!') - >>> print(len(embeddings[0].embedding)) # Vector dimensions - - Example - With metadata: - >>> embeddings = await ai.embed( - ... embedder='googleai/text-embedding-004', - ... content='Product description', - ... metadata={'category': 'electronics'}, - ... ) - - Example - With embedder options: - >>> embeddings = await ai.embed( - ... embedder='googleai/text-embedding-004', - ... content='Search query', - ... options={'task_type': 'RETRIEVAL_QUERY'}, - ... ) - - Example - Using EmbedderRef: - >>> ref = create_embedder_ref('googleai/text-embedding-004', config={'task_type': 'CLUSTERING'}) - >>> embeddings = await ai.embed(embedder=ref, content='Text') - """ - embedder_name = self._resolve_embedder_name(embedder) - embedder_config: dict[str, object] = {} - - # Extract config and version from EmbedderRef (not done for embed_many per JS behavior) - if isinstance(embedder, EmbedderRef): - embedder_config = embedder.config or {} - if embedder.version: - embedder_config['version'] = embedder.version # Handle version from ref - - # Merge options passed to embed() with config from EmbedderRef - final_options = {**(embedder_config or {}), **(options or {})} - - embed_action = await self.registry.resolve_embedder(embedder_name) - if embed_action is None: - raise ValueError(f'Embedder "{embedder_name}" not found') - - if content is None: - raise ValueError('Content must be specified for embedding.') - - documents = [Document.from_text(content, metadata)] if isinstance(content, str) else [content] - - # Document subclasses DocumentData, so this is type-safe at runtime. - # list is invariant so list[Document] isn't assignable to list[DocumentData] - response = ( - await embed_action.arun( - EmbedRequest( - input=documents, # pyright: ignore[reportArgumentType] - options=final_options, - ) - ) - ).response - return response.embeddings - - async def embed_many( - self, - embedder: str | EmbedderRef | None = None, - content: list[str] | list[Document] | list[DocumentData] | None = None, - metadata: dict[str, object] | None = None, - options: dict[str, object] | None = None, - ) -> list[Embedding]: - """Embeds multiple documents or strings in a single batch call. - - Generates vector embeddings for multiple pieces of content in one API - call. This is more efficient than calling embed() multiple times when - you have a batch of items to embed. - - Important: Unlike embed(), this method does NOT extract config/version - from EmbedderRef. It only uses the ref to resolve the embedder name - and passes options directly. This matches the JS canonical behavior. - - Args: - embedder: Embedder name (e.g., 'googleai/text-embedding-004') or - an EmbedderRef. - content: List of strings, Documents, or DocumentData to embed. - metadata: Optional metadata to apply to all items. Only used when - content items are strings. - options: Optional embedder-specific options. - - Returns: - List of Embedding objects, one per input item (same order). - - Raises: - ValueError: If embedder is not specified or not found. - ValueError: If content is not specified. - - Example - Basic batch embedding: - >>> embeddings = await ai.embed_many( - ... embedder='googleai/text-embedding-004', - ... content=['Doc 1', 'Doc 2', 'Doc 3'], - ... ) - >>> for i, emb in enumerate(embeddings): - ... print(f'Doc {i}: {len(emb.embedding)} dims') - - Example - With shared metadata: - >>> embeddings = await ai.embed_many( - ... embedder='googleai/text-embedding-004', - ... content=['text1', 'text2'], - ... metadata={'batch_id': 'batch-001'}, - ... ) - - Example - With options (EmbedderRef config is NOT extracted): - >>> embeddings = await ai.embed_many( - ... embedder='googleai/text-embedding-004', - ... content=documents, - ... options={'task_type': 'RETRIEVAL_DOCUMENT'}, - ... ) - """ - if content is None: - raise ValueError('Content must be specified for embedding.') - - # Convert strings to Documents if needed - documents: list[Document | DocumentData] = [ - Document.from_text(item, metadata) if isinstance(item, str) else item for item in content - ] - - # Resolve embedder name (JS embedMany does not extract config/version from ref) - embedder_name = self._resolve_embedder_name(embedder) - - embed_action = await self.registry.resolve_embedder(embedder_name) - if embed_action is None: - raise ValueError(f'Embedder "{embedder_name}" not found') - - response = (await embed_action.arun(EmbedRequest(input=documents, options=options))).response - return response.embeddings - - async def evaluate( - self, - evaluator: str | EvaluatorRef | None = None, - dataset: list[BaseDataPoint] | None = None, - options: dict[str, object] | None = None, - eval_run_id: str | None = None, - ) -> EvalResponse: - """Evaluates a dataset using an evaluator. - - Args: - evaluator: Name or reference of the evaluator to use. - dataset: Dataset to evaluate. - options: Evaluation options. - eval_run_id: Optional ID for the evaluation run. - - Returns: - The evaluation results. - """ - evaluator_name: str = '' - evaluator_config: dict[str, object] = {} - - if isinstance(evaluator, EvaluatorRef): - evaluator_name = evaluator.name - evaluator_config = evaluator.config_schema or {} - elif isinstance(evaluator, str): - evaluator_name = evaluator - else: - raise ValueError('Evaluator must be specified as a string name or an EvaluatorRef.') - - final_options = {**(evaluator_config or {}), **(options or {})} - - eval_action = await self.registry.resolve_evaluator(evaluator_name) - if eval_action is None: - raise ValueError(f'Evaluator "{evaluator_name}" not found') - - if not eval_run_id: - eval_run_id = str(uuid.uuid4()) - - if dataset is None: - raise ValueError('Dataset must be specified for evaluation.') - - return ( - await eval_action.arun( - EvalRequest( - dataset=dataset, - options=final_options, - eval_run_id=eval_run_id, - ) - ) - ).response - - @staticmethod - def current_context() -> dict[str, Any] | None: - """Retrieves the current execution context for the running action. - - This allows tools and other actions to access context data (like auth - or metadata) passed through the execution chain via ContextVars. - This provides parity with the JavaScript SDK's context handling. - - Returns: - The current context dictionary, or None if not running in an action. - """ - return ActionRunContext._current_context() # pyright: ignore[reportPrivateUsage] - - def dynamic_tool( - self, - name: str, - fn: Callable[..., object], - description: str | None = None, - metadata: dict[str, object] | None = None, - ) -> Action: - """Creates an unregistered tool action. - - This is useful for creating tools that are passed directly to generate() - without being registered in the global registry. Dynamic tools behave exactly - like registered tools but offer more flexibility for runtime-defined logic. - - Args: - name: The unique name of the tool. - fn: The function that implements the tool logic. - description: Optional human-readable description of what the tool does. - metadata: Optional dictionary of metadata about the tool. - - Returns: - An Action instance of kind TOOL, configured for dynamic execution. - """ - tool_meta: dict[str, object] = metadata.copy() if metadata else {} - tool_meta['type'] = 'tool' - tool_meta['dynamic'] = True - return Action( - kind=ActionKind.TOOL, - name=name, - fn=fn, - description=description, - metadata=tool_meta, - ) - - async def flush_tracing(self) -> None: - """Flushes all registered trace processors. - - This ensures all pending spans are exported before the application - shuts down, preventing loss of telemetry data. - """ - provider = trace_api.get_tracer_provider() - if isinstance(provider, TracerProvider): - await ensure_async(provider.force_flush)() - - async def run( - self, - name: str, - func_or_input: object, - maybe_fn: Callable[..., T | Awaitable[T]] | None = None, - metadata: dict[str, Any] | None = None, - ) -> T: - """Runs a function as a discrete step within a trace. - - This method is used to create sub-spans (steps) within a flow or other action. - Each run step is recorded separately in the trace, making it easier to - debug and monitor the internal execution of complex flows. - - It supports two call signatures: - 1. `run(name, fn)`: Runs the provided function. - 2. `run(name, input, fn)`: Passes the input to the function and records it. - - Args: - name: The descriptive name of the span/step. - func_or_input: Either the function to execute, or input data to pass - to `maybe_fn`. - maybe_fn: An optional function to execute if `func_or_input` is - provided as input data. - metadata: Optional metadata to associate with the generated trace span. - - Returns: - The result of the function execution. - """ - fn: Callable[..., T | Awaitable[T]] - input_data: Any = None - has_input = False - - if maybe_fn: - fn = maybe_fn - input_data = func_or_input - has_input = True - elif callable(func_or_input): - fn = cast(Callable[..., T | Awaitable[T]], func_or_input) - else: - raise ValueError('A function must be provided to run.') - - span_metadata = SpanMetadata(name=name, metadata=metadata) - with run_in_new_span(span_metadata, labels={'genkit:type': 'flowStep'}) as span: - try: - if has_input: - span.set_input(input_data) - result = await ensure_async(fn)(input_data) - else: - result = await ensure_async(fn)() - - span.set_output(result) - return result - except Exception: - # We catch all exceptions here to ensure they are captured by - # the trace span context manager before being re-raised. - # The GenkitSpan wrapper (run_in_new_span) handles recording - # the exception details. - raise - - async def check_operation(self, operation: Operation) -> Operation: - """Checks the status of a long-running background operation. - - This method matches JS checkOperation from js/ai/src/check-operation.ts. - - It looks up the background action by the operation's action key and - calls its check method to get updated status. - - Args: - operation: The Operation object to check. Must have an action - field specifying which background model created it. - - Returns: - An updated Operation object with the current status. - - Raises: - ValueError: If the operation is missing original request information - or if the background action cannot be resolved. - - Example: - >>> # Start a background operation - >>> response = await ai.generate(model=veo_model, prompt='A cat') - >>> operation = response.operation - >>> # Poll until done - >>> while not operation.done: - ... await asyncio.sleep(5) - ... operation = await ai.check_operation(operation) - >>> # Use the result - >>> print(operation.output) - """ - return await check_operation_impl(self.registry, operation) - - async def cancel_operation(self, operation: Operation) -> Operation: - """Cancels a long-running background operation. - - This method attempts to cancel an in-progress operation. Not all - background models support cancellation. - - If cancellation is not supported, returns the operation unchanged - (matching JS behavior). - - Args: - operation: The Operation object to cancel. Must have an action - field specifying which background model created it. - - Returns: - An updated Operation object reflecting the cancellation attempt. - - Raises: - ValueError: If the operation is missing original request information - or if the background action cannot be resolved. - - Example: - >>> # Start a background operation - >>> response = await ai.generate(model=veo_model, prompt='A cat') - >>> operation = response.operation - >>> # Cancel it - >>> operation = await ai.cancel_operation(operation) - """ - if not operation.action: - raise ValueError('Provided operation is missing original request information') - - background_action = await lookup_background_action(self.registry, operation.action) - if background_action is None: - raise ValueError(f'Failed to resolve background action from original request: {operation.action}') - - return await background_action.cancel(operation) - - async def generate_operation( - self, - model: str | None = None, - prompt: str | Part | list[Part] | None = None, - system: str | Part | list[Part] | None = None, - messages: list[Message] | None = None, - tools: list[str] | None = None, - return_tool_requests: bool | None = None, - tool_choice: ToolChoice | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - max_turns: int | None = None, - context: dict[str, object] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_constrained: bool | None = None, - output: OutputConfig | OutputConfigDict | Output[Any] | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | None = None, - on_chunk: ModelStreamingCallback | None = None, - ) -> Operation: - """Generate content using a long-running model and return an Operation. - - This method is for models that support long-running operations (like - video generation with Veo). It returns an Operation that can be polled - with check_operation() until completion. - - Note: This is a beta feature. Only models that support long-running - operations (model.supports.long_running = True) can be used with this - method. - - The Operation Flow - ================== - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ generate_operation() Flow β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Resolve Model β”‚ β”‚ - β”‚ β”‚ Check supports β”‚ β”‚ - β”‚ β”‚ long_running β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ β”‚ β”‚ - β”‚ β–Ό β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Error β”‚ β”‚ Call generate β”‚ β”‚ - β”‚ β”‚ (no β”‚ β”‚ Get operation β”‚ β”‚ - β”‚ β”‚ LRO) β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Operation β”‚ β”‚ - β”‚ β”‚ done=False β”‚ ──► poll with check_operation() β”‚ - β”‚ β”‚ id=... β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - Args: - model: The model to use for generation (must support long_running). - prompt: The prompt text or parts. - system: System message for the model. - messages: Conversation history. - tools: Tool names available to the model. - return_tool_requests: Whether to return tool requests. - tool_choice: How the model should choose tools. - config: Generation configuration. - max_turns: Maximum conversation turns. - context: Additional context data. - output_format: Output format (e.g., 'json'). - output_content_type: Output content type. - output_instructions: Output formatting instructions. - output_constrained: Whether to constrain output to schema. - output: Typed output configuration. - use: Middleware to apply. - docs: Documents for grounding. - on_chunk: Callback for streaming chunks. - - Returns: - An Operation object for tracking the long-running generation. - - Raises: - GenkitError: If the model doesn't support long-running operations. - GenkitError: If the model didn't return an operation. - - Example: - >>> # Generate video with Veo (long-running operation) - >>> operation = await ai.generate_operation( - ... model='googleai/veo-2.0-generate-001', - ... prompt='A banana riding a bicycle.', - ... ) - >>> # Poll until done - >>> while not operation.done: - ... await asyncio.sleep(5) - ... operation = await ai.check_operation(operation) - >>> # Access result - >>> print(operation.output) - """ - # Resolve the model and check for long_running support - resolved_model = model or self.registry.default_model - if not resolved_model: - raise GenkitError( - status='INVALID_ARGUMENT', - message='No model specified for generate_operation.', - ) - - model_action = await self.registry.resolve_action(ActionKind.MODEL, resolved_model) - if not model_action: - raise GenkitError( - status='NOT_FOUND', - message=f"Model '{resolved_model}' not found.", - ) - - # Check if model supports long-running operations - model_info = model_action.metadata.get('model') if model_action.metadata else None - supports_long_running = False - if model_info: - # model_info can be ModelInfo or dict - if hasattr(model_info, 'supports'): - supports_attr = getattr(model_info, 'supports', None) - if supports_attr: - supports_long_running = getattr(supports_attr, 'long_running', False) - elif isinstance(model_info, dict): - # Cast to dict[str, Any] for type checker - model_info_dict: dict[str, Any] = model_info # type: ignore[assignment] - supports = model_info_dict.get('supports') - if isinstance(supports, dict): - supports_long_running = bool(supports.get('longRunning', False)) - - if not supports_long_running: - raise GenkitError( - status='INVALID_ARGUMENT', - message=f"Model '{model_action.name}' does not support long running operations.", - ) - - # Call generate - response = await self.generate( - model=model, - prompt=prompt, - system=system, - messages=messages, - tools=tools, - return_tool_requests=return_tool_requests, - tool_choice=tool_choice, - config=config, - max_turns=max_turns, - context=context, - output_format=output_format, - output_content_type=output_content_type, - output_instructions=output_instructions, - output_constrained=output_constrained, - output=output, - use=use, - docs=docs, - on_chunk=on_chunk, - ) - - # Extract operation from response - if not hasattr(response, 'operation') or not response.operation: - raise GenkitError( - status='FAILED_PRECONDITION', - message=f"Model '{model_action.name}' did not return an operation.", - ) - - return response.operation diff --git a/py/packages/genkit/src/genkit/ai/_base_async.py b/py/packages/genkit/src/genkit/ai/_base_async.py deleted file mode 100644 index 5914b5d292..0000000000 --- a/py/packages/genkit/src/genkit/ai/_base_async.py +++ /dev/null @@ -1,235 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Asynchronous server gateway interface implementation for Genkit.""" - -import asyncio -import signal -import socket -import threading -from collections.abc import Coroutine -from typing import Any, TypeVar - -import anyio -import uvicorn - -from genkit.aio.loop import run_loop -from genkit.blocks.formats import built_in_formats -from genkit.blocks.generate import define_generate_action -from genkit.core.environment import is_dev_environment -from genkit.core.logging import get_logger -from genkit.core.plugin import Plugin -from genkit.core.reflection import create_reflection_asgi_app -from genkit.core.registry import Registry - -from ._registry import GenkitRegistry -from ._runtime import RuntimeManager -from ._server import ServerSpec - -logger = get_logger(__name__) - -T = TypeVar('T') - - -class _ReflectionServer(uvicorn.Server): - """A uvicorn.Server subclass that signals startup completion via a threading.Event.""" - - def __init__(self, config: uvicorn.Config, ready: threading.Event) -> None: - """Initialize the server with a ready event to set on startup.""" - super().__init__(config) - self._ready = ready - - async def startup(self, sockets: list | None = None) -> None: - """Override to set the ready event once uvicorn finishes binding.""" - try: - await super().startup(sockets=sockets) - finally: - self._ready.set() - - -class GenkitBase(GenkitRegistry): - """Base class with shared infra for Genkit instances (sync and async).""" - - def __init__( - self, - plugins: list[Plugin] | None = None, - model: str | None = None, - reflection_server_spec: ServerSpec | None = None, - ) -> None: - """Initialize a new Genkit instance. - - Args: - plugins: List of plugins to initialize. - model: Model name to use. - reflection_server_spec: Server spec for the reflection - server. If not provided in dev mode, a default will be used. - """ - super().__init__() - self._reflection_server_spec: ServerSpec | None = reflection_server_spec - self._reflection_ready = threading.Event() - self._initialize_registry(model, plugins) - # Ensure the default generate action is registered for async usage. - define_generate_action(self.registry) - # In dev mode, start the reflection server immediately in a background - # daemon thread so it's available regardless of which web framework (or - # none) the user chooses. - if is_dev_environment(): - self._start_reflection_background() - - def _start_reflection_background(self) -> None: - """Start the Dev UI reflection server in a background daemon thread. - - The thread owns its own asyncio event loop so it never conflicts with - the main thread's loop (whether that's uvicorn, FastAPI, or none). - Sets ``self._reflection_ready`` once the server is listening. - """ - - def _thread_main() -> None: - async def _run() -> None: - sockets: list[socket.socket] | None = None - spec = self._reflection_server_spec - if spec is None: - # Bind to port 0 to let the OS choose an available port and - # pass the socket to uvicorn to avoid a check-then-bind race. - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.bind(('127.0.0.1', 0)) - sock.listen(2048) - host, port = sock.getsockname() - spec = ServerSpec(scheme='http', host=host, port=port) - self._reflection_server_spec = spec - sockets = [sock] - - server = _make_reflection_server(self.registry, spec, ready=self._reflection_ready) - async with RuntimeManager(spec, lazy_write=True) as runtime_manager: - server_task = asyncio.create_task(server.serve(sockets=sockets)) - - # _ReflectionServer.startup() sets _reflection_ready once uvicorn binds. - # Use asyncio.to_thread so we don't block the event loop. - await asyncio.to_thread(self._reflection_ready.wait) - - if server.should_exit: - logger.warning(f'Reflection server at {spec.url} failed to start.') - return - - runtime_manager.write_runtime_file() - await logger.ainfo(f'Genkit Dev UI reflection server running at {spec.url}') - - # Keep running until the process exits (daemon thread). - await server_task - - asyncio.run(_run()) - - t = threading.Thread(target=_thread_main, daemon=True, name='genkit-reflection-server') - t.start() - - def _initialize_registry(self, model: str | None, plugins: list[Plugin] | None) -> None: - """Initialize the registry for the Genkit instance. - - Args: - model: Model name to use. - plugins: List of plugins to initialize. - - Raises: - ValueError: If an invalid plugin is provided. - - Returns: - None - """ - self.registry.default_model = model - for fmt in built_in_formats: - self.define_format(fmt) - - if not plugins: - logger.warning('No plugins provided to Genkit') - else: - for plugin in plugins: - if isinstance(plugin, Plugin): # pyright: ignore[reportUnnecessaryIsInstance] - self.registry.register_plugin(plugin) - else: - raise ValueError(f'Invalid {plugin=} provided to Genkit: must be of type `genkit.ai.Plugin`') - - def run_main(self, coro: Coroutine[Any, Any, T]) -> T | None: - """Run the user's main coroutine. - - In development mode (`GENKIT_ENV=dev`), this runs the user's coroutine - then blocks until Ctrl+C or SIGTERM, keeping the background reflection - server (started in ``__init__``) alive for the Dev UI. - - In production mode, this simply runs the user's coroutine to completion - using ``uvloop.run()`` for performance if available, otherwise - ``asyncio.run()``. - - Args: - coro: The main coroutine provided by the user. - - Returns: - The result of the user's coroutine, or None on graceful shutdown. - """ - if not is_dev_environment(): - logger.info('Running in production mode.') - return run_loop(coro) - - logger.info('Running in development mode.') - - async def dev_runner() -> T | None: - user_result: T | None = None - try: - user_result = await coro - logger.debug('User coroutine completed successfully.') - except Exception: - logger.exception('User coroutine failed') - - # Block until Ctrl+C (SIGINT handled by anyio) or SIGTERM, keeping - # the daemon reflection thread alive. - logger.info('Script done β€” Dev UI running. Press Ctrl+C to stop.') - try: - async with anyio.create_task_group() as tg: - - async def _handle_sigterm(tg_: anyio.abc.TaskGroup) -> None: # type: ignore[name-defined] - with anyio.open_signal_receiver(signal.SIGTERM) as sigs: - async for _ in sigs: - tg_.cancel_scope.cancel() - return - - tg.start_soon(_handle_sigterm, tg) - await anyio.sleep_forever() - except anyio.get_cancelled_exc_class(): - pass - - logger.info('Dev UI server stopped.') - return user_result - - return anyio.run(dev_runner) - - -def _make_reflection_server(registry: Registry, spec: ServerSpec, ready: threading.Event) -> _ReflectionServer: - """Make a reflection server for the given registry and spec. - - This is a helper function to make it easier to test the reflection server - in isolation. - - Args: - registry: The registry to use for the reflection server. - spec: The spec to use for the reflection server. - ready: threading.Event to set once uvicorn finishes binding. - - Returns: - A uvicorn server instance. - """ - app = create_reflection_asgi_app(registry=registry) - # pyrefly: ignore[bad-argument-type] - Starlette app is valid ASGI app for uvicorn - config = uvicorn.Config(app, host=spec.host, port=spec.port, loop='asyncio') - return _ReflectionServer(config, ready=ready) diff --git a/py/packages/genkit/src/genkit/ai/_registry.py b/py/packages/genkit/src/genkit/ai/_registry.py deleted file mode 100644 index 96571a7b17..0000000000 --- a/py/packages/genkit/src/genkit/ai/_registry.py +++ /dev/null @@ -1,1679 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Genkit maintains a registry of all actions. - -An **action** is a remote callable function that uses typed-JSON RPC over HTTP -to allow the framework and users to define custom AI functionality. There are -several kinds of action defined by [ActionKind][genkit.core.action.ActionKind]: - -| Kind | Description | -|---------------|-------------| -| `'chat-llm'` | Chat LLM | -| `'custom'` | Custom | -| `'embedder'` | Embedder | -| `'evaluator'` | Evaluator | -| `'flow'` | Flow | -| `'indexer'` | Indexer | -| `'model'` | Model | -| `'prompt'` | Prompt | -| `'resource'` | Resource | -| `'retriever'` | Retriever | -| `'text-llm'` | Text LLM | -| `'tool'` | Tool | -| `'util'` | Utility | -""" - -import asyncio -import inspect -import traceback -import uuid -from collections.abc import AsyncIterator, Awaitable, Callable -from functools import wraps -from typing import Any, Generic, ParamSpec, cast, overload - -from pydantic import BaseModel -from typing_extensions import Never, TypeVar - -from genkit.aio import ensure_async -from genkit.blocks.background_model import ( - BackgroundAction, - CancelModelOpFn, - CheckModelOpFn, - StartModelOpFn, - define_background_model as define_background_model_block, -) -from genkit.blocks.dap import ( - DapConfig, - DapFn, - DynamicActionProvider, - define_dynamic_action_provider as define_dap_block, -) -from genkit.blocks.document import Document -from genkit.blocks.embedding import EmbedderFn, EmbedderOptions -from genkit.blocks.evaluator import BatchEvaluatorFn, EvaluatorFn -from genkit.blocks.formats.types import FormatDef -from genkit.blocks.interfaces import Input, Output -from genkit.blocks.model import ModelFn, ModelMiddleware -from genkit.blocks.prompt import ( - ExecutablePrompt, - define_helper, - define_partial, - define_prompt, - define_schema, -) -from genkit.blocks.reranker import ( - RankedDocument, - RerankerFn, - RerankerOptions, - RerankerRef, - define_reranker as define_reranker_block, - rerank as rerank_block, -) -from genkit.blocks.resource import ( - ResourceFn, - ResourceOptions, - define_resource as define_resource_block, -) -from genkit.blocks.retriever import IndexerFn, RetrieverFn -from genkit.blocks.tools import ToolRunContext -from genkit.codec import dump_dict -from genkit.core.action import Action, ActionResponse, ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.logging import get_logger -from genkit.core.registry import Registry -from genkit.core.schema import to_json_schema -from genkit.core.tracing import run_in_new_span -from genkit.core.typing import ( - DocumentData, - DocumentPart, - EvalFnResponse, - EvalRequest, - EvalResponse, - EvalStatusEnum, - GenerationCommonConfig, - Message, - ModelInfo, - Part, - RetrieverResponse, - Score, - SpanMetadata, - ToolChoice, -) - -EVALUATOR_METADATA_KEY_DISPLAY_NAME = 'evaluatorDisplayName' -EVALUATOR_METADATA_KEY_DEFINITION = 'evaluatorDefinition' -EVALUATOR_METADATA_KEY_IS_BILLED = 'evaluatorIsBilled' - -logger = get_logger(__name__) - -# TypeVars for generic input/output typing -InputT = TypeVar('InputT') -OutputT = TypeVar('OutputT') -P = ParamSpec('P') -R = TypeVar('R') -T = TypeVar('T') -CallT = TypeVar('CallT') -ChunkT = TypeVar('ChunkT', default=Never) - - -def get_func_description(func: Callable[..., Any], description: str | None = None) -> str: - """Get the description of a function. - - Args: - func: The function to get the description of. - description: The description to use if the function docstring is - empty. - """ - if description is not None: - return description - if func.__doc__ is not None: - return func.__doc__ - return '' - - -R = TypeVar('R') - - -class SimpleRetrieverOptions(BaseModel, Generic[R]): - """Configuration options for `define_simple_retriever`. - - This class defines how items returned by a simple retriever handler are - mapped into Genkit `DocumentData` objects. - - Attributes: - name: The unique name of the retriever. - content: Specifies how to extract content from the returned items. - Can be a string key (for dict items) or a callable that transforms the item. - metadata: Specifies how to extract metadata from the returned items. - Can be a list of keys (for dict items) or a callable that transforms the item. - config_schema: Optional Pydantic schema or JSON schema for retriever configuration. - """ - - name: str - content: str | Callable[[R], str | list[DocumentPart]] | None = None - metadata: list[str] | Callable[[R], dict[str, Any]] | None = None - config_schema: type[BaseModel] | dict[str, Any] | None = None - - -def _item_to_document(item: R, options: SimpleRetrieverOptions[R]) -> DocumentData: - """Internal helper to convert a raw item to a Genkit DocumentData.""" - if isinstance(item, (Document, DocumentData)): - return item - - if isinstance(item, str): - return Document.from_text(item) - - if callable(options.content): - transformed = options.content(item) - if isinstance(transformed, str): - return Document.from_text(transformed) - else: - # transformed is list[DocumentPart] - return DocumentData(content=transformed) - - if isinstance(options.content, str) and isinstance(item, dict): - item_dict = cast(dict[str, object], item) - return Document.from_text(str(item_dict[options.content])) - - if options.content is None and isinstance(item, str): - return Document.from_text(item) - - raise ValueError(f'Cannot convert item to document without content option. Item: {item}') - - -def _item_to_metadata(item: R, options: SimpleRetrieverOptions[R]) -> dict[str, Any] | None: - """Internal helper to extract metadata from a raw item for a Document.""" - if isinstance(item, str): - return None - - if isinstance(options.metadata, list) and isinstance(item, dict): - item_dict = cast(dict[str, object], item) - result: dict[str, Any] = {} - for key in options.metadata: - str_key = str(key) - value = item_dict.get(str_key) - if value is not None: - result[str_key] = value - return result - - if callable(options.metadata): - return options.metadata(item) - - if options.metadata is None and isinstance(item, dict): - out = cast(dict[str, Any], item.copy()) - if isinstance(options.content, str) and options.content in out: - del out[options.content] - return out - - return None - - -class GenkitRegistry: - """User-facing API for interacting with Genkit registry.""" - - def __init__(self) -> None: - """Initialize the Genkit registry.""" - self.registry: Registry = Registry() - - @overload - # pyrefly: ignore[inconsistent-overload] - Overloads differentiate async vs sync returns - def flow( - self, name: str | None = None, description: str | None = None - ) -> Callable[[Callable[P, Awaitable[T]]], 'FlowWrapper[P, Awaitable[T], T]']: ... - - @overload - # pyrefly: ignore[inconsistent-overload] - Overloads differentiate async vs sync returns - # Overloads appear to overlap because T could be Awaitable[T], but at runtime we - # distinguish async vs sync functions correctly. - def flow( # pyright: ignore[reportOverlappingOverload] - self, name: str | None = None, description: str | None = None - ) -> Callable[[Callable[P, T]], 'FlowWrapper[P, T, T]']: ... - - def flow( # pyright: ignore[reportInconsistentOverload] - self, name: str | None = None, description: str | None = None - ) -> Callable[[Callable[P, Awaitable[T]] | Callable[P, T]], 'FlowWrapper[P, Awaitable[T] | T, T]']: - """Decorator to register a function as a flow. - - Args: - name: Optional name for the flow. If not provided, uses the - function name. - description: Optional description for the flow. If not provided, - uses the function docstring. - - Returns: - A decorator function that registers the flow. - """ - - def wrapper(func: Callable[P, Awaitable[T]] | Callable[P, T]) -> 'FlowWrapper[P, Awaitable[T] | T, T]': - """Register the decorated function as a flow. - - Args: - func: The function to register as a flow. - - Returns: - The wrapped function that executes the flow. - """ - flow_name = name if name is not None else getattr(func, '__name__', 'unnamed_flow') - flow_description = get_func_description(func, description) - action = self.registry.register_action( - name=flow_name, - kind=cast(ActionKind, ActionKind.FLOW), - # pyrefly: ignore[bad-argument-type] - func union type is valid for register_action - fn=func, - description=flow_description, - span_metadata={'genkit:metadata:flow:name': flow_name}, - ) - - # pyrefly: ignore[bad-argument-type] - func is valid for wraps despite union type - @wraps(func) - async def async_wrapper(*args: P.args, **_kwargs: P.kwargs) -> T: - """Asynchronous wrapper for the flow function. - - Args: - *args: Positional arguments to pass to the flow function. - **_kwargs: Keyword arguments (unused, for signature compatibility). - - Returns: - The response from the flow function. - """ - # Flows accept at most one input argument - input_arg = cast(T | None, args[0] if args else None) - return (await action.arun(input_arg)).response - - # pyrefly: ignore[bad-argument-type] - func is valid for wraps despite union type - @wraps(func) - def sync_wrapper(*args: P.args, **_kwargs: P.kwargs) -> T: - """Synchronous wrapper for the flow function. - - Args: - *args: Positional arguments to pass to the flow function. - **_kwargs: Keyword arguments (unused, for signature compatibility). - - Returns: - The response from the flow function. - """ - # Flows accept at most one input argument - input_arg = cast(T | None, args[0] if args else None) - return action.run(input_arg).response - - wrapped_fn = cast( - Callable[P, Awaitable[T]] | Callable[P, T], async_wrapper if action.is_async else sync_wrapper - ) - flow = FlowWrapper( - fn=cast(Callable[P, Awaitable[T] | T], wrapped_fn), - action=cast(Action[Any, T, Never], action), - ) - return flow - - return wrapper - - def define_helper(self, name: str, fn: Callable[..., Any]) -> None: - """Define a Handlebars helper function in the registry. - - Args: - name: The name of the helper function. - fn: The helper function to register. - """ - define_helper(self.registry, name, fn) - - def define_partial(self, name: str, source: str) -> None: - """Define a Handlebars partial template in the registry. - - Partials are reusable template fragments that can be included - in other prompts using {{>partialName}} syntax. - - Args: - name: The name of the partial. - source: The template source code for the partial. - """ - define_partial(self.registry, name, source) - - def define_schema(self, name: str, schema: type[BaseModel]) -> type[BaseModel]: - """Register a Pydantic schema for use in prompts. - - Schemas registered with this method can be referenced by name in - .prompt files using the `output.schema` field. - - Args: - name: The name to register the schema under. - schema: The Pydantic model class to register. - - Returns: - The schema that was registered (for convenience). - - Example: - ```python - RecipeSchema = ai.define_schema('Recipe', Recipe) - ``` - - Then in a .prompt file: - ```yaml - output: - schema: Recipe - ``` - """ - define_schema(self.registry, name, schema) - return schema - - def define_json_schema(self, name: str, json_schema: dict[str, object]) -> dict[str, object]: - """Register a JSON schema for use in prompts. - - This method registers a raw JSON Schema (as a dictionary) rather than - a Pydantic model class. Use this when you have a JSON Schema from an - external source or need more control over the schema definition. - - Schema Types Comparison - ======================= - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Schema Registration Methods β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ define_schema() β”‚ define_json_schema() β”‚ - β”‚ ──────────────────────────┼───────────────────────────────────│ - β”‚ Input: Pydantic class β”‚ Input: JSON Schema dict β”‚ - β”‚ Type-safe β”‚ Dynamic/external schemas β”‚ - β”‚ Auto-converts to JSON β”‚ Direct JSON Schema control β”‚ - β”‚ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - Args: - name: The name to register the schema under. - json_schema: The JSON Schema dictionary to register. - - Returns: - The JSON schema that was registered (for convenience). - - Example: - ```python - # Register a JSON Schema directly - recipe_schema = ai.define_json_schema( - 'Recipe', - { - 'type': 'object', - 'properties': { - 'title': {'type': 'string'}, - 'ingredients': {'type': 'array', 'items': {'type': 'string'}}, - 'instructions': {'type': 'string'}, - }, - 'required': ['title', 'ingredients', 'instructions'], - }, - ) - ``` - - Then in a .prompt file: - ```yaml - output: - schema: Recipe - ``` - - See Also: - - define_schema: For registering Pydantic models - - JSON Schema spec: https://json-schema.org/ - """ - self.registry.register_schema(name, json_schema) - return json_schema - - def define_dynamic_action_provider( - self, - config: DapConfig | str, - fn: DapFn, - ) -> DynamicActionProvider: - """Define and register a Dynamic Action Provider (DAP). - - A DAP is a factory that can dynamically provide actions at runtime, - enabling integration with external systems like MCP (Model Context - Protocol) servers, plugin marketplaces, or other dynamic action sources. - - Dynamic Action Provider Overview - ================================ - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ How DAPs Work β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ 1. Register DAP with Genkit β”‚ - β”‚ 2. When resolving an unknown action, Genkit queries DAPs β”‚ - β”‚ 3. DAP fetches actions from external source (cached) β”‚ - β”‚ 4. Actions are returned and can be used like static actions β”‚ - β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Genkit β”‚ ──► β”‚ DAP β”‚ ──► β”‚ External β”‚ β”‚ - β”‚ β”‚ Registry β”‚ β”‚ Cache β”‚ β”‚ System β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ (MCP, etc.) β”‚ β”‚ - β”‚ β–² β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ Actions β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - Args: - config: DAP configuration (DapConfig) or just a name string. - - name: Unique identifier for this DAP - - description: What this DAP provides - - cache_config: Caching behavior (ttl_millis) - - metadata: Additional metadata - fn: Async function that returns actions organized by type. - Should return a dict like: {'tool': [action1, action2], ...} - - Returns: - The registered DynamicActionProvider. - - Example: - ```python - from genkit.ai import Genkit - from genkit.blocks.dap import DapConfig, DapCacheConfig - - ai = Genkit() - - - # Simple DAP - just a name - async def get_tools(): - return { - 'tool': [ - ai.dynamic_tool(name='tool1', fn=lambda x: x), - ] - } - - - dap = ai.define_dynamic_action_provider('my-tools', get_tools) - - # DAP with custom caching - dap = ai.define_dynamic_action_provider( - config=DapConfig( - name='mcp-tools', - description='Tools from MCP server', - cache_config=DapCacheConfig(ttl_millis=10000), - ), - fn=get_tools, - ) - - # Invalidate cache when needed - dap.invalidate_cache() - - # Get a specific action - action = await dap.get_action('tool', 'tool1') - ``` - - Use Cases: - - MCP Integration: Connect to Model Context Protocol servers - - Plugin Systems: Load actions from external plugins - - Multi-tenant: Provide tenant-specific actions - - Feature Flags: Enable/disable actions at runtime - - See Also: - - genkit.plugins.mcp: MCP plugin using DAPs - - JS implementation: js/core/src/dynamic-action-provider.ts - """ - return define_dap_block(self.registry, config, fn) - - def tool( - self, name: str | None = None, description: str | None = None - ) -> Callable[[Callable[P, T]], Callable[P, T]]: - """Decorator to register a function as a tool. - - Args: - name: Optional name for the flow. If not provided, uses the function - name. - description: Description for the tool to be passed to the model; - if not provided, uses the function docstring. - - Returns: - A decorator function that registers the tool. - """ - - def wrapper(func: Callable[P, T]) -> Callable[P, T]: - """Register the decorated function as a tool. - - Args: - func: The function to register as a tool. - - Returns: - The wrapped function that executes the tool. - """ - tool_name = name if name is not None else getattr(func, '__name__', 'unnamed_tool') - tool_description = get_func_description(func, description) - - input_spec = inspect.getfullargspec(func) - - func_any = cast(Callable[..., Any], func) - - def tool_fn_wrapper(*args: Any) -> Any: # noqa: ANN401 - # Dynamic dispatch based on function signature - pyright can't verify ParamSpec here - match len(input_spec.args): - case 0: - return func_any() - case 1: - return func_any(args[0]) - case 2: - return func_any(args[0], ToolRunContext(cast(ActionRunContext, args[1]))) - case _: - raise ValueError('tool must have 0-2 args...') - - action = self.registry.register_action( - name=tool_name, - kind=cast(ActionKind, ActionKind.TOOL), - description=tool_description, - fn=tool_fn_wrapper, - metadata_fn=func, - ) - - @wraps(func) - async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any: # noqa: ANN401 - """Asynchronous wrapper for the tool function. - - Args: - *args: Positional arguments to pass to the tool function. - **kwargs: Keyword arguments to pass to the tool function. - - Returns: - The response from the tool function. - """ - action_any = cast(Any, action) - return (await action_any.arun(*args, **kwargs)).response - - @wraps(func) - def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any: # noqa: ANN401 - """Synchronous wrapper for the tool function. - - Args: - *args: Positional arguments to pass to the tool function. - **kwargs: Keyword arguments to pass to the tool function. - - Returns: - The response from the tool function. - """ - action_any = cast(Any, action) - return action_any.run(*args, **kwargs).response - - return cast(Callable[P, T], async_wrapper if action.is_async else sync_wrapper) - - return wrapper - - def define_retriever( - self, - name: str, - fn: RetrieverFn[Any], - config_schema: type[BaseModel] | dict[str, object] | None = None, - metadata: dict[str, object] | None = None, - description: str | None = None, - ) -> Action: - """Define a retriever action. - - Args: - name: Name of the retriever. - fn: Function implementing the retriever behavior. - config_schema: Optional schema for retriever configuration. - metadata: Optional metadata for the retriever. - description: Optional description for the retriever. - """ - retriever_meta: dict[str, object] = dict(metadata) if metadata else {} - retriever_info: dict[str, object] - existing_retriever = retriever_meta.get('retriever') - if isinstance(existing_retriever, dict): - retriever_info = {str(key): value for key, value in existing_retriever.items()} - else: - retriever_info = {} - retriever_meta['retriever'] = retriever_info - label_value = retriever_info.get('label') - if not isinstance(label_value, str) or not label_value: - retriever_info['label'] = name - if config_schema: - retriever_info['customOptions'] = to_json_schema(config_schema) - - retriever_description = get_func_description(fn, description) - return self.registry.register_action( - name=name, - kind=cast(ActionKind, ActionKind.RETRIEVER), - fn=fn, - metadata=retriever_meta, - description=retriever_description, - ) - - def define_simple_retriever( - self, - options: SimpleRetrieverOptions[R] | str, - handler: Callable[[DocumentData, Any], list[R] | Awaitable[list[R]]], - description: str | None = None, - ) -> Action: - """Define a simple retriever action. - - A simple retriever makes it easy to map existing data into documents - that can be used for prompt augmentation. - - Args: - options: Configuration options for the retriever, or just the name. - handler: A function that queries a datastore and returns items - from which to extract documents. - description: Optional description for the retriever. - - Returns: - The registered Action for the retriever. - """ - if isinstance(options, str): - options = SimpleRetrieverOptions(name=options) - - async def retriever_fn(query: Document, options_obj: Any) -> RetrieverResponse: # noqa: ANN401 - items = await ensure_async(handler)(query, options_obj) - docs = [] - for item in items: - doc = _item_to_document(item, options) - if not isinstance(item, str): - doc.metadata = _item_to_metadata(item, options) - docs.append(doc) - return RetrieverResponse(documents=docs) - - return self.define_retriever( - name=options.name, - fn=retriever_fn, - config_schema=options.config_schema, - description=description, - ) - - def define_indexer( - self, - name: str, - fn: IndexerFn[Any], - config_schema: type[BaseModel] | dict[str, object] | None = None, - metadata: dict[str, object] | None = None, - description: str | None = None, - ) -> Action: - """Define an indexer action. - - Args: - name: Name of the indexer. - fn: Function implementing the indexer behavior. - config_schema: Optional schema for indexer configuration. - metadata: Optional metadata for the indexer. - description: Optional description for the indexer. - """ - indexer_meta: dict[str, object] = dict(metadata) if metadata else {} - indexer_info: dict[str, object] - existing_indexer = indexer_meta.get('indexer') - if isinstance(existing_indexer, dict): - indexer_info = {str(key): value for key, value in existing_indexer.items()} - else: - indexer_info = {} - indexer_meta['indexer'] = indexer_info - label_value = indexer_info.get('label') - if not isinstance(label_value, str) or not label_value: - indexer_info['label'] = name - if config_schema: - indexer_info['customOptions'] = to_json_schema(config_schema) - - indexer_description = get_func_description(fn, description) - return self.registry.register_action( - name=name, - kind=cast(ActionKind, ActionKind.INDEXER), - fn=fn, - metadata=indexer_meta, - description=indexer_description, - ) - - def define_reranker( - self, - name: str, - fn: RerankerFn[Any], - config_schema: type[BaseModel] | dict[str, object] | None = None, - metadata: dict[str, object] | None = None, - description: str | None = None, - ) -> Action: - """Define a reranker action. - - Rerankers reorder documents based on their relevance to a query. - They are commonly used in RAG pipelines to improve retrieval quality. - - Args: - name: Name of the reranker. - fn: Function implementing the reranker behavior. Should accept - (query_doc, documents, options) and return RerankerResponse. - config_schema: Optional schema for reranker configuration. - metadata: Optional metadata for the reranker. - description: Optional description for the reranker. - - Returns: - The registered Action for the reranker. - - Example: - >>> async def my_reranker(query, docs, options): - ... # Score documents based on relevance to query - ... scored = [(doc, compute_score(query, doc)) for doc in docs] - ... scored.sort(key=lambda x: x[1], reverse=True) - ... return RerankerResponse(documents=[...]) - >>> ai.define_reranker('my-reranker', my_reranker) - """ - # Extract label and config from metadata - reranker_label: str = name - reranker_config_schema: dict[str, object] | None = None - - # Check if metadata has reranker info - if metadata and 'reranker' in metadata: - existing = metadata['reranker'] - if isinstance(existing, dict): - existing_dict = cast(dict[str, object], existing) - label_val = existing_dict.get('label') - if isinstance(label_val, str) and label_val: - reranker_label = label_val - opts_val = existing_dict.get('customOptions') - if isinstance(opts_val, dict): - reranker_config_schema = cast(dict[str, object], opts_val) - - # Override with config_schema if provided - if config_schema: - reranker_config_schema = to_json_schema(config_schema) - - reranker_description = get_func_description(fn, description) - return define_reranker_block( - self.registry, - name=name, - fn=fn, - options=RerankerOptions( - config_schema=reranker_config_schema, - label=reranker_label, - ), - description=reranker_description, - ) - - async def rerank( - self, - reranker: str | Action | RerankerRef, - query: str | DocumentData, - documents: list[DocumentData], - options: object | None = None, - ) -> list[RankedDocument]: - """Rerank documents based on their relevance to a query. - - This method takes a query and a list of documents, and returns the - documents reordered by relevance as determined by the specified reranker. - - Args: - reranker: The reranker to use - can be a name string, Action, or RerankerRef. - query: The query to rank documents against - can be a string or DocumentData. - documents: The list of documents to rerank. - options: Optional configuration options for this rerank call. - - Returns: - A list of RankedDocument objects sorted by relevance score. - - Raises: - ValueError: If the reranker cannot be resolved. - - Example: - >>> ranked_docs = await ai.rerank( - ... reranker='my-reranker', - ... query='What is machine learning?', - ... documents=[doc1, doc2, doc3], - ... ) - >>> for doc in ranked_docs: - ... print(f'Score: {doc.score}, Text: {doc.text()}') - """ - return await rerank_block( - self.registry, - { - 'reranker': reranker, - 'query': query, - 'documents': documents, - 'options': options, - }, - ) - - def define_evaluator( - self, - name: str, - display_name: str, - definition: str, - fn: EvaluatorFn[Any], - is_billed: bool = False, - config_schema: type[BaseModel] | dict[str, object] | None = None, - metadata: dict[str, object] | None = None, - description: str | None = None, - ) -> Action: - """Define an evaluator action. - - This action runs the callback function on the every sample of - the input dataset. - - Args: - name: Name of the evaluator. - fn: Function implementing the evaluator behavior. - display_name: User-visible display name - definition: User-visible evaluator definition - is_billed: Whether the evaluator performs any billed actions - (paid APIs, LLMs etc.) - config_schema: Optional schema for evaluator configuration. - metadata: Optional metadata for the evaluator. - description: Optional description for the evaluator. - """ - evaluator_meta: dict[str, object] = dict(metadata) if metadata else {} - evaluator_info: dict[str, object] - existing_evaluator = evaluator_meta.get('evaluator') - if isinstance(existing_evaluator, dict): - evaluator_info = {str(key): value for key, value in existing_evaluator.items()} - else: - evaluator_info = {} - evaluator_meta['evaluator'] = evaluator_info - evaluator_info[EVALUATOR_METADATA_KEY_DEFINITION] = definition - evaluator_info[EVALUATOR_METADATA_KEY_DISPLAY_NAME] = display_name - evaluator_info[EVALUATOR_METADATA_KEY_IS_BILLED] = is_billed - label_value = evaluator_info.get('label') - if not isinstance(label_value, str) or not label_value: - evaluator_info['label'] = name - if config_schema: - evaluator_info['customOptions'] = to_json_schema(config_schema) - - evaluator_description = get_func_description(fn, description) - - async def eval_stepper_fn(req: EvalRequest) -> EvalResponse: - eval_responses: list[EvalFnResponse] = [] - for index in range(len(req.dataset)): - datapoint = req.dataset[index] - if datapoint.test_case_id is None: - datapoint.test_case_id = str(uuid.uuid4()) - span_metadata = SpanMetadata( - name=f'Test Case {datapoint.test_case_id}', - metadata={'evaluator:evalRunId': req.eval_run_id}, - ) - try: - # Try to run with tracing, but fallback if tracing infrastructure fails - # (e.g., in environments with NonRecordingSpans like pre-commit) - try: - with run_in_new_span(span_metadata, labels={'genkit:type': 'evaluator'}) as span: - span_id = span.span_id - trace_id = span.trace_id - try: - span.set_input(datapoint) - test_case_output = await fn(datapoint, req.options) - test_case_output.span_id = span_id - test_case_output.trace_id = trace_id - span.set_output(test_case_output) - eval_responses.append(test_case_output) - except Exception as e: - logger.debug(f'eval_stepper_fn error: {e!s}') - logger.debug(traceback.format_exc()) - evaluation = Score( - error=f'Evaluation of test case {datapoint.test_case_id} failed: \n{e!s}', - status=cast(EvalStatusEnum, EvalStatusEnum.FAIL), - ) - eval_responses.append( - # The ty type checker only recognizes aliases, so we use them - # to pass both ty check and runtime validation. - EvalFnResponse( - span_id=span_id, - trace_id=trace_id, - test_case_id=datapoint.test_case_id, - evaluation=evaluation, - ) - ) - # Raise to mark span as failed - raise e - except (AttributeError, UnboundLocalError): - # Fallback: run without span - try: - test_case_output = await fn(datapoint, req.options) - eval_responses.append(test_case_output) - except Exception as e: - logger.debug(f'eval_stepper_fn error: {e!s}') - logger.debug(traceback.format_exc()) - evaluation = Score( - error=f'Evaluation of test case {datapoint.test_case_id} failed: \n{e!s}', - status=cast(EvalStatusEnum, EvalStatusEnum.FAIL), - ) - eval_responses.append( - EvalFnResponse( - test_case_id=datapoint.test_case_id, - evaluation=evaluation, - ) - ) - except Exception: # noqa: S112 - intentionally continue processing other datapoints - # Continue to process other points - continue - return EvalResponse(eval_responses) - - return self.registry.register_action( - name=name, - kind=cast(ActionKind, ActionKind.EVALUATOR), - fn=eval_stepper_fn, - metadata=evaluator_meta, - description=evaluator_description, - ) - - def define_batch_evaluator( - self, - name: str, - display_name: str, - definition: str, - fn: BatchEvaluatorFn[Any], - is_billed: bool = False, - config_schema: type[BaseModel] | dict[str, object] | None = None, - metadata: dict[str, object] | None = None, - description: str | None = None, - ) -> Action: - """Define a batch evaluator action. - - This action runs the callback function on the entire dataset. - - Args: - name: Name of the evaluator. - fn: Function implementing the evaluator behavior. - display_name: User-visible display name - definition: User-visible evaluator definition - is_billed: Whether the evaluator performs any billed actions - (paid APIs, LLMs etc.) - config_schema: Optional schema for evaluator configuration. - metadata: Optional metadata for the evaluator. - description: Optional description for the evaluator. - """ - evaluator_meta: dict[str, object] = metadata.copy() if metadata else {} - if 'evaluator' not in evaluator_meta: - evaluator_meta['evaluator'] = {} - # Cast to dict for nested operations - pyrefly doesn't narrow nested dict types - evaluator_dict = cast(dict[str, object], evaluator_meta['evaluator']) - evaluator_dict[EVALUATOR_METADATA_KEY_DEFINITION] = definition - evaluator_dict[EVALUATOR_METADATA_KEY_DISPLAY_NAME] = display_name - evaluator_dict[EVALUATOR_METADATA_KEY_IS_BILLED] = is_billed - if 'label' not in evaluator_dict or not evaluator_dict['label']: - evaluator_dict['label'] = name - if config_schema: - evaluator_dict['customOptions'] = to_json_schema(config_schema) - - evaluator_description = get_func_description(fn, description) - return self.registry.register_action( - name=name, - kind=cast(ActionKind, ActionKind.EVALUATOR), - fn=fn, - metadata=evaluator_meta, - description=evaluator_description, - ) - - def define_model( - self, - name: str, - fn: ModelFn, - config_schema: type[BaseModel] | dict[str, object] | None = None, - metadata: dict[str, object] | None = None, - info: ModelInfo | None = None, - description: str | None = None, - ) -> Action: - """Define a custom model action. - - Args: - name: Name of the model. - fn: Function implementing the model behavior. - config_schema: Optional schema for model configuration. - metadata: Optional metadata for the model. - info: Optional ModelInfo for the model. - description: Optional description for the model. - """ - # Build model options dict - model_options: dict[str, object] = {} - - # Start with info if provided - if info: - model_info_dict = dump_dict(info) - if isinstance(model_info_dict, dict): - for key, value in model_info_dict.items(): - if isinstance(key, str): - model_options[key] = value - - # Check if metadata has model info - if metadata and 'model' in metadata: - existing = metadata['model'] - if isinstance(existing, dict): - existing_dict = cast(dict[str, object], existing) - for key, value in existing_dict.items(): - if isinstance(key, str) and key not in model_options: - model_options[key] = value - - # Default label to name if not set - if 'label' not in model_options or not model_options['label']: - model_options['label'] = name - - # Add config schema if provided - if config_schema: - model_options['customOptions'] = to_json_schema(config_schema) - - # Build the final metadata dict - model_meta: dict[str, object] = metadata.copy() if metadata else {} - model_meta['model'] = model_options - - model_description = get_func_description(fn, description) - return self.registry.register_action( - name=name, - kind=cast(ActionKind, ActionKind.MODEL), - fn=fn, - metadata=model_meta, - description=model_description, - ) - - def define_background_model( - self, - name: str, - start: StartModelOpFn, - check: CheckModelOpFn, - cancel: CancelModelOpFn | None = None, - label: str | None = None, - info: ModelInfo | None = None, - config_schema: type[BaseModel] | dict[str, object] | None = None, - metadata: dict[str, object] | None = None, - description: str | None = None, - ) -> BackgroundAction: - """Define a background model for long-running AI operations. - - Background models are used for tasks like video generation (Veo) or - large image generation that may take seconds or minutes to complete. - Unlike regular models that return results immediately, background models - return an Operation that can be polled for completion. - - This matches JS defineBackgroundModel from js/ai/src/model.ts. - - Args: - name: Unique name for this background model. - start: Async function to start the background operation. - Takes (GenerateRequest, ActionRunContext) -> Operation. - check: Async function to check operation status. - Takes (Operation) -> Operation. - cancel: Optional async function to cancel operations. - Takes (Operation) -> Operation. - label: Human-readable label (defaults to name). - info: Model capability information (ModelInfo). - config_schema: Schema for model configuration options. - metadata: Additional metadata for the model. - description: Description for the model action. - - Returns: - A BackgroundAction that can be used to start/check/cancel operations. - - Example: - >>> async def start_video(req: GenerateRequest, ctx) -> Operation: - ... job_id = await video_api.submit(req.messages[0].content[0].text) - ... return Operation(id=job_id, done=False) - >>> async def check_video(op: Operation) -> Operation: - ... status = await video_api.get_status(op.id) - ... if status.complete: - ... return Operation(id=op.id, done=True, output=...) - ... return Operation(id=op.id, done=False) - >>> action = ai.define_background_model( - ... name='video-gen', - ... start=start_video, - ... check=check_video, - ... ) - """ - return define_background_model_block( - registry=self.registry, - name=name, - start=start, - check=check, - cancel=cancel, - label=label, - info=info, - config_schema=config_schema, - metadata=metadata, - description=description, - ) - - def define_embedder( - self, - name: str, - fn: EmbedderFn, - options: EmbedderOptions | None = None, - metadata: dict[str, object] | None = None, - description: str | None = None, - ) -> Action: - """Define a custom embedder action. - - Args: - name: Name of the model. - fn: Function implementing the embedder behavior. - options: Optional options for the embedder. - metadata: Optional metadata for the model. - description: Optional description for the embedder. - """ - embedder_meta: dict[str, object] = dict(metadata) if metadata else {} - embedder_info: dict[str, object] - existing_embedder = embedder_meta.get('embedder') - if isinstance(existing_embedder, dict): - embedder_info = {str(key): value for key, value in existing_embedder.items()} - else: - embedder_info = {} - embedder_meta['embedder'] = embedder_info - - if options: - if options.label: - embedder_info['label'] = options.label - if options.dimensions: - embedder_info['dimensions'] = options.dimensions - if options.supports: - embedder_info['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True) - if options.config_schema: - embedder_info['customOptions'] = to_json_schema(options.config_schema) - - embedder_description = get_func_description(fn, description) - return self.registry.register_action( - name=name, - kind=cast(ActionKind, ActionKind.EMBEDDER), - fn=fn, - metadata=embedder_meta, - description=embedder_description, - ) - - def define_format(self, format: FormatDef) -> None: - """Registers a custom format in the registry. - - Args: - format: The format to register. - """ - self.registry.register_value('format', format.name, format) - - # Overload 1: Both input and output typed -> ExecutablePrompt[InputT, OutputT] - @overload - def define_prompt( - self, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, object] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, object] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, object] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - *, - input: 'Input[InputT]', - output: 'Output[OutputT]', - ) -> 'ExecutablePrompt[InputT, OutputT]': ... - - # Overload 2: Only input typed -> ExecutablePrompt[InputT, Any] - @overload - def define_prompt( - self, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, object] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, object] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, object] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - *, - input: 'Input[InputT]', - output: None = None, - ) -> 'ExecutablePrompt[InputT, Any]': ... - - # Overload 3: Only output typed -> ExecutablePrompt[Any, OutputT] - @overload - def define_prompt( - self, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, object] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, object] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, object] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - input: None = None, - *, - output: 'Output[OutputT]', - ) -> 'ExecutablePrompt[Any, OutputT]': ... - - # Overload 4: Neither typed -> ExecutablePrompt[Any, Any] - @overload - def define_prompt( - self, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, object] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, object] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, object] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - input: None = None, - output: None = None, - ) -> 'ExecutablePrompt[Any, Any]': ... - - def define_prompt( - self, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, object] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, object] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, object] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, object] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - input: 'Input[Any] | None' = None, - output: 'Output[Any] | None' = None, - ) -> 'ExecutablePrompt[Any, Any]': - """Define a prompt. - - Args: - name: Optional name for the prompt. - variant: Optional variant name for the prompt. - model: Optional model name to use for the prompt. - config: Optional configuration for the model. - description: Optional description for the prompt. - input_schema: Optional schema for the input to the prompt. - system: Optional system message for the prompt. - prompt: Optional prompt for the model. - messages: Optional messages for the model. - output_format: Optional output format for the prompt. - output_content_type: Optional output content type for the prompt. - output_instructions: Optional output instructions for the prompt. - output_schema: Optional schema for the output from the prompt. - output_constrained: Optional flag indicating whether the output - should be constrained. - max_turns: Optional maximum number of turns for the prompt. - return_tool_requests: Optional flag indicating whether tool requests - should be returned. - metadata: Optional metadata for the prompt. - tools: Optional list of tools to use for the prompt. - tool_choice: Optional tool choice for the prompt. - use: Optional list of model middlewares to use for the prompt. - docs: Optional list of documents or a callable to be used for grounding. - input: Typed input configuration using Input[T]. When provided, - the prompt's input parameter is type-checked. - output: Typed output configuration using Output[T]. When provided, - the response output is typed. - - Example: - ```python - from genkit import Input, Output - from pydantic import BaseModel - - - class RecipeInput(BaseModel): - dish: str - - - class Recipe(BaseModel): - name: str - ingredients: list[str] - - - # With typed input AND output - recipe_prompt = ai.define_prompt( - name='recipe', - prompt='Create a recipe for {dish}', - input=Input(schema=RecipeInput), - output=Output(schema=Recipe), - ) - - # Input is type-checked! - response = await recipe_prompt(RecipeInput(dish='pizza')) - response.output.name # βœ“ Typed as str - ``` - """ - if input is not None and output is not None: - return define_prompt( - self.registry, - name=name, - variant=variant, - model=model, - config=config, - description=description, - input_schema=input_schema, - system=system, - prompt=prompt, - messages=messages, - output_format=output_format, - output_content_type=output_content_type, - output_instructions=output_instructions, - output_schema=output_schema, - output_constrained=output_constrained, - max_turns=max_turns, - return_tool_requests=return_tool_requests, - metadata=metadata, - tools=tools, - tool_choice=tool_choice, - use=use, - docs=docs, - input=input, - output=output, - ) - if input is not None: - return define_prompt( - self.registry, - name=name, - variant=variant, - model=model, - config=config, - description=description, - input_schema=input_schema, - system=system, - prompt=prompt, - messages=messages, - output_format=output_format, - output_content_type=output_content_type, - output_instructions=output_instructions, - output_schema=output_schema, - output_constrained=output_constrained, - max_turns=max_turns, - return_tool_requests=return_tool_requests, - metadata=metadata, - tools=tools, - tool_choice=tool_choice, - use=use, - docs=docs, - input=input, - output=None, - ) - if output is not None: - return define_prompt( - self.registry, - name=name, - variant=variant, - model=model, - config=config, - description=description, - input_schema=input_schema, - system=system, - prompt=prompt, - messages=messages, - output_format=output_format, - output_content_type=output_content_type, - output_instructions=output_instructions, - output_schema=output_schema, - output_constrained=output_constrained, - max_turns=max_turns, - return_tool_requests=return_tool_requests, - metadata=metadata, - tools=tools, - tool_choice=tool_choice, - use=use, - docs=docs, - input=None, - output=output, - ) - return define_prompt( - self.registry, - name=name, - variant=variant, - model=model, - config=config, - description=description, - input_schema=input_schema, - system=system, - prompt=prompt, - messages=messages, - output_format=output_format, - output_content_type=output_content_type, - output_instructions=output_instructions, - output_schema=output_schema, - output_constrained=output_constrained, - max_turns=max_turns, - return_tool_requests=return_tool_requests, - metadata=metadata, - tools=tools, - tool_choice=tool_choice, - use=use, - docs=docs, - input=None, - output=None, - ) - - # Overload 1: Neither typed -> ExecutablePrompt[Any, Any] - @overload - def prompt( - self, - name: str, - variant: str | None = None, - *, - input: None = None, - output: None = None, - ) -> ExecutablePrompt[Any, Any]: ... - - # Overload 2: Only input typed - @overload - def prompt( - self, - name: str, - variant: str | None = None, - *, - input: Input[InputT], - output: None = None, - ) -> ExecutablePrompt[InputT, Any]: ... - - # Overload 3: Only output typed - @overload - def prompt( - self, - name: str, - variant: str | None = None, - *, - input: None = None, - output: Output[OutputT], - ) -> ExecutablePrompt[Any, OutputT]: ... - - # Overload 4: Both input and output typed - @overload - def prompt( - self, - name: str, - variant: str | None = None, - *, - input: Input[InputT], - output: Output[OutputT], - ) -> ExecutablePrompt[InputT, OutputT]: ... - - def prompt( - self, - name: str, - variant: str | None = None, - *, - input: Input[InputT] | None = None, - output: Output[OutputT] | None = None, - ) -> ExecutablePrompt[InputT, OutputT] | ExecutablePrompt[Any, Any]: - """Look up a prompt by name and optional variant. - - This matches the JavaScript prompt() function behavior. - - Can look up prompts that were: - 1. Defined programmatically using define_prompt() - 2. Loaded from .prompt files using load_prompt_folder() - - Args: - name: The name of the prompt. - variant: Optional variant name. - input: Optional typed input configuration. When provided, the - prompt's input parameter will be type-checked. - output: Optional typed output configuration. When provided, - response.output will be statically typed. - - Returns: - An ExecutablePrompt instance. - - Example: - ```python - # Without type hints (output is Any) - prompt = ai.prompt('greet') - - # With typed output (response.output is MySchema) - prompt = ai.prompt('greet', output=Output(schema=MySchema)) - response = await prompt(input={'name': 'World'}) - response.output # Statically typed as MySchema - ``` - """ - # Extract schema types if provided - input_schema = input.schema if input else None - output_schema = output.schema if output else None - - return ExecutablePrompt( - registry=self.registry, - _name=name, - variant=variant, - input_schema=input_schema, - output_schema=output_schema, - ) - - def define_resource( - self, - opts: 'ResourceOptions | None' = None, - fn: 'ResourceFn | None' = None, - *, - name: str | None = None, - uri: str | None = None, - template: str | None = None, - description: str | None = None, - metadata: dict[str, object] | None = None, - ) -> Action: - """Define a resource action. - - Args: - opts: Options defining the resource (e.g. uri, template, name). - fn: Function implementing the resource behavior. - name: Optional name for the resource. - uri: Optional URI for the resource. - template: Optional URI template for the resource. - description: Optional description for the resource. - metadata: Optional metadata for the resource. - - Returns: - The registered Action for the resource. - """ - if fn is None: - raise ValueError('A function `fn` must be provided to define a resource.') - if opts is None: - opts = {} - if name: - opts['name'] = name - if uri: - opts['uri'] = uri - if template: - opts['template'] = template - if description: - opts['description'] = description - if metadata: - opts['metadata'] = metadata - - return define_resource_block(self.registry, opts, fn) - - -class FlowWrapper(Generic[P, CallT, T, ChunkT]): - """A wrapper for flow functions to add `stream` method. - - This class wraps a flow function and provides a `stream` method for - asynchronous execution. - """ - - def __init__(self, fn: Callable[P, CallT], action: Action[Any, T, ChunkT]) -> None: - """Initialize the FlowWrapper. - - Args: - fn: The function to wrap. - action: The action to wrap. - """ - self._fn: Callable[P, CallT] = fn - self._action: Action[Any, T, ChunkT] = action - - def __call__(self, *args: P.args, **kwds: P.kwargs) -> CallT: - """Call the wrapped function. - - Args: - *args: Positional arguments to pass to the function. - **kwds: Keyword arguments to pass to the function. - - Returns: - The result of the function call. - """ - return self._fn(*args, **kwds) - - def stream( - self, - input: object = None, - context: dict[str, object] | None = None, - telemetry_labels: dict[str, object] | None = None, - timeout: float | None = None, - ) -> tuple[AsyncIterator[ChunkT], asyncio.Future[ActionResponse[T]]]: - """Run the flow and return an async iterator of the results. - - Args: - input: The input to the action. - context: The context to pass to the action. - telemetry_labels: The telemetry labels to pass to the action. - timeout: The timeout for the streaming action. - - Returns: - A tuple containing: - - An AsyncIterator of the chunks from the action. - - An asyncio.Future that resolves to the final result of the action. - """ - return self._action.stream(input=input, context=context, telemetry_labels=telemetry_labels, timeout=timeout) diff --git a/py/packages/genkit/src/genkit/ai/_server.py b/py/packages/genkit/src/genkit/ai/_server.py deleted file mode 100644 index 78a10364a2..0000000000 --- a/py/packages/genkit/src/genkit/ai/_server.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Functionality used by the Genkit veneer to start multiple servers. - -The following servers may be started depending upon the host environment: - -- Reflection API server. -- Flows server. - -The reflection API server is started only in dev mode, which is enabled by the -setting the environment variable `GENKIT_ENV` to `dev`. By default, the -reflection API server binds and listens on (localhost, 3100 or the next available -port). The flows server is the production servers that exposes flows and actions -over HTTP. -""" - -from dataclasses import dataclass - -from genkit.core.logging import get_logger - -logger = get_logger(__name__) - - -@dataclass -class ServerSpec: - """ServerSpec encapsulates the scheme, host and port information. - - This class defines the server binding and listening configuration. - """ - - port: int - scheme: str = 'http' - host: str = 'localhost' - - @property - def url(self) -> str: - """URL evaluates to the host base URL given the server specs.""" - return f'{self.scheme}://{self.host}:{self.port}' diff --git a/py/packages/genkit/src/genkit/aio/__init__.py b/py/packages/genkit/src/genkit/aio/__init__.py deleted file mode 100644 index 225dd7dcd2..0000000000 --- a/py/packages/genkit/src/genkit/aio/__init__.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Asynchronous utilities for the Genkit framework. - -This module provides async utilities used internally by Genkit for handling -concurrent operations, streaming responses, and async/sync interoperability. - -Overview: - Genkit is async-first, leveraging Python's asyncio for concurrent operations. - This module provides utilities that simplify async patterns used throughout - the framework, particularly for streaming model responses. - -Key Components: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Component β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Channel β”‚ AsyncIterator for streaming chunks with a final β”‚ - β”‚ β”‚ value (used for streaming model responses) β”‚ - β”‚ ensure_async() β”‚ Wrap sync/async functions to ensure async callable β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Using Channel for streaming: - - ```python - from genkit.aio import Channel - - # Create a channel for streaming chunks with a final result - channel: Channel[str, int] = Channel() - - - async def producer(): - for chunk in ['Hello', ' ', 'World']: - channel.send(chunk) - channel.close(final_value=len('Hello World')) - - - # Consume chunks - async for chunk in channel: - print(chunk, end='') - - # Get final value - result = await channel.result() - ``` - - Using ensure_async: - - ```python - from genkit.aio import ensure_async - - - def sync_fn(x: int) -> int: - return x * 2 - - - async_fn = ensure_async(sync_fn) - result = await async_fn(5) # Returns 10 - ``` - -See Also: - - asyncio documentation: https://docs.python.org/3/library/asyncio.html -""" - -from ._util import ensure_async -from .channel import Channel - -__all__ = [ - 'Channel', - 'ensure_async', -] diff --git a/py/packages/genkit/src/genkit/aio/_util.py b/py/packages/genkit/src/genkit/aio/_util.py deleted file mode 100644 index 09adadd267..0000000000 --- a/py/packages/genkit/src/genkit/aio/_util.py +++ /dev/null @@ -1,58 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AIO util module for defining and managing AIO utilities.""" - -import inspect -from collections.abc import Awaitable, Callable -from typing import Any - - -def ensure_async(fn: Callable[..., Any] | Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]: - """Ensure the function is async. - - This function handles three cases: - 1. `fn` is already an async function -> return as-is - 2. `fn` is a sync function that returns a regular value -> wrap in async - 3. `fn` is a sync function (e.g., lambda) that returns a coroutine -> await it - - Args: - fn: The function to ensure is async. - - Returns: - The async function. - """ - is_async = inspect.iscoroutinefunction(fn) - if is_async: - return fn - - async def async_wrapper(*args: object, **kwargs: object) -> Any: # noqa: ANN401 - """Wrap the function in an async function. - - Args: - *args: The arguments to the function. - **kwargs: The keyword arguments to the function. - - Returns: - The result of the function. - """ - result = fn(*args, **kwargs) - # Handle case where a sync function (e.g., lambda) returns a coroutine - if inspect.iscoroutine(result): - return await result - return result - - return async_wrapper diff --git a/py/packages/genkit/src/genkit/aio/channel.py b/py/packages/genkit/src/genkit/aio/channel.py deleted file mode 100644 index 0d7a62aff7..0000000000 --- a/py/packages/genkit/src/genkit/aio/channel.py +++ /dev/null @@ -1,221 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Asyncio helpers for asynchronous communication and data flow control.""" - -from __future__ import annotations - -import asyncio -from collections.abc import AsyncIterator -from typing import Any, Generic - -from typing_extensions import TypeVar - -from ._compat import wait_for - -T = TypeVar('T') # Type of items in the channel -R = TypeVar('R', default=Any) # Type of the close future result (defaults to Any) - - -class Channel(Generic[T, R]): - """An asynchronous channel for sending and receiving values. - - This class provides an asynchronous queue-like interface, allowing values to - be sent and received between different parts of an asynchronous program. It - supports both sending values and closing the channel, which will signal to - any receivers that no more values will be sent. - - The Channel class implements the async iterator protocol, allowing it to be - used in async for loops. - - Typical usage: - ```python - channel: Channel[int, int] = Channel(timeout=0.1) - - # Send values to the channel - channel.send(1) - channel.send(2) - - # Receive values from the channel in another task - async for value in channel: - print(value) # Will print 1, then 2 - ``` - - """ - - def __init__(self, timeout: float | int | None = None) -> None: - """Initializes a new Channel. - - The channel is initialized with an internal queue to store values, a - future to signal when the channel is closed, and an internal close - future that will be set when the channel should be closed. - - Args: - timeout: The timeout in seconds for the __anext__ method. If None, - waits indefinitely (default). - - Raises: - ValueError: If the timeout is negative. - """ - if timeout is not None and timeout < 0: - raise ValueError('Timeout must be non-negative') - - self.queue: asyncio.Queue[T] = asyncio.Queue() - self.closed: asyncio.Future[R] = asyncio.Future() - self._close_future: asyncio.Future[R] | None = None - self._timeout: float | int | None = timeout - - def __aiter__(self) -> AsyncIterator[T]: - """Returns the asynchronous iterator for the channel. - - Returns: - AsyncIterator[T]: The channel object itself, which implements the - `__anext__` method required for async iteration. - """ - return self - - async def __anext__(self) -> T: - """Retrieves the next value from the channel. - - If the queue is not empty, the value is returned immediately. - Otherwise, it waits until a value is available or the channel is closed. - - Implements the `__anext__` method required for async iteration. - - Raises: - StopAsyncIteration: If the channel is closed and no more values are - available, signaling the end of the iteration. - TimeoutError: If the timeout is exceeded while waiting for a - value and a timeout has been specified. - - Returns: - T: The next value from the channel. - """ - # If the queue has values, return the next value immediately. - # Otherwise, wait for a value to be available or the channel to close or - # a timeout to occur. - if not self.queue.empty(): - return self.queue.get_nowait() - - # Create the task to retrieve the next value - pop_task = asyncio.ensure_future(self._pop()) - if not self._close_future: - # Wait for the pop task with a timeout, raise TimeoutError if a - # timeout is specified and is exceeded and automatically cancel the - # pending task. - return await wait_for(pop_task, timeout=self._timeout) - - # Wait for either the pop task or the close future to complete. A - # timeout is added to prevent indefinite blocking, unless - # specifically set to None. - # NOTE: asyncio.wait does not cancel tasks on timeout by default. - finished, _pending = await asyncio.wait( - [pop_task, self._close_future], - return_when=asyncio.FIRST_COMPLETED, - timeout=self._timeout, - ) - - # If timeout occurred (nothing finished), cancel pop task and raise - # Note: Don't cancel _close_future as it's owned by external code - if not finished: - _ = pop_task.cancel() - raise TimeoutError('Channel timeout exceeded') - - # If the pop task completed, return its result. - if pop_task in finished: - return pop_task.result() - - # If the close future completed, raise StopAsyncIteration. - if self._close_future in finished: - # Cancel pop task if we're done, avoid warnings. - _ = pop_task.cancel() - raise StopAsyncIteration - - # Wait for the pop task with a timeout, raise TimeoutError if a timeout - # is specified and is exceeded and automatically cancel the pending - # task. - return await wait_for(pop_task, timeout=self._timeout) - - def send(self, value: T) -> None: - """Sends a value into the channel. - - The value is added to the internal queue for consumers to retrieve. - This is a non-blocking operation. - - Args: - value: The value to send through the channel. - - Raises: - asyncio.QueueFull: If the channel's internal queue is full. - - Returns: - None. - """ - self.queue.put_nowait(value) - - def set_close_future(self, future: asyncio.Future[R]) -> None: - """Sets a future that, when completed, will close the channel. - - When the provided future completes, the channel will be marked as - closed, signaling to consumers that no more values will be sent. - - Args: - future: The future to monitor for channel closure. When this future - completes, the channel will be closed. - - Raises: - ValueError: If the provided future is None. - """ - if future is None: # pyright: ignore[reportUnnecessaryComparison] - raise ValueError('Cannot set a None future') # pyright: ignore[reportUnreachable] - - def _handle_done(v: asyncio.Future[R]) -> None: - """Handle future completion, propagating results or errors to self.closed. - - This callback ensures proper propagation of the future's final state - (success, exception, or cancellation) to the channel's closed future, - allowing consumers to properly handle completion or errors. - """ - # Propagate cancellation to notify consumers that the operation was cancelled - if v.cancelled(): - _ = self.closed.cancel() - elif (exc := v.exception()) is not None: - self.closed.set_exception(exc) - else: - self.closed.set_result(v.result()) - - self._close_future = asyncio.ensure_future(future) - if self._close_future is not None: # pyright: ignore[reportUnnecessaryComparison] - self._close_future.add_done_callback(_handle_done) - - async def _pop(self) -> T: - """Asynchronously retrieves a value from the internal queue. - - This is an internal method used by the async iterator implementation to - wait until a value is available in the queue. - - Raises: - StopAsyncIteration: If a None value is retrieved, indicating the - channel is closed. - - Returns: - T: The retrieved value from the channel. - """ - r = await self.queue.get() - self.queue.task_done() - if not r: - raise StopAsyncIteration - return r diff --git a/py/packages/genkit/src/genkit/aio/loop.py b/py/packages/genkit/src/genkit/aio/loop.py deleted file mode 100644 index 56d8fb5bda..0000000000 --- a/py/packages/genkit/src/genkit/aio/loop.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Asyncio loop utilities.""" - -import asyncio -import threading -from asyncio import AbstractEventLoop -from collections.abc import AsyncIterable, Callable, Coroutine, Iterable -from typing import TypeVar - -from genkit.core.logging import get_logger - -logger = get_logger(__name__) - -T = TypeVar('T') - - -def create_loop() -> AbstractEventLoop: - """Creates a new asyncio event loop or returns the current one. - - This function attempts to get the current event loop. If no current loop - exists (e.g., in a new thread), it creates and returns a new event loop. - - Returns: - An asyncio event loop instance. - """ - try: - return asyncio.get_event_loop() - except Exception: - return asyncio.new_event_loop() - - -def run_async(loop: asyncio.AbstractEventLoop, fn: Callable[[], Coroutine[object, object, T]]) -> T | None: - """Runs an async callable on the given event loop and blocks until completion. - - If the loop is already running (e.g., called from within an async context), - it schedules the callable using `asyncio.run_coroutine_threadsafe` and uses - a threading lock to block until the callable finishes. - - If the loop is not running, it uses `loop.run_until_complete`. - - Args: - loop: The asyncio event loop to run the callable on. - fn: The async callable (e.g., a coroutine function) to execute. - - Returns: - The result returned by the callable `fn`. - - Raises: - Any exception raised by the callable `fn`. - """ - if loop.is_running(): - output: T | None = None - error: Exception | None = None - lock = threading.Lock() - _ = lock.acquire() - - async def run_fn() -> T | None: - nonlocal lock - nonlocal output - nonlocal error - try: - output = await fn() - return output - except Exception as e: - error = e - finally: - lock.release() - return None - - _ = asyncio.run_coroutine_threadsafe(run_fn(), loop=loop) - - def wait_for_done() -> None: - nonlocal lock - _ = lock.acquire() - - thread = threading.Thread(target=wait_for_done) - thread.start() - thread.join() - - if error: - raise error # pyright: ignore[reportUnreachable] - - return output - else: - return loop.run_until_complete(fn()) - - -def iter_over_async(ait: AsyncIterable[T], loop: asyncio.AbstractEventLoop) -> Iterable[T]: - """Synchronously iterates over an AsyncIterable using a specified event loop. - - This function bridges asynchronous iteration with synchronous code by - running the `__anext__` calls of the async iterator on the provided event - loop and yielding the results synchronously. - - Args: - ait: The asynchronous iterable to iterate over. - loop: The asyncio event loop to use for running `__anext__`. - - Yields: - Items from the asynchronous iterable. - """ - ait_iter = ait.__aiter__() - - async def get_next() -> tuple[bool, T | None]: - try: - obj = await ait_iter.__anext__() - return False, obj - except StopAsyncIteration: - return True, None - - while True: - done, obj = loop.run_until_complete(get_next()) - if done: - break - assert obj is not None # Type narrowing: obj is T when done=False - yield obj - - -def run_loop(coro: Coroutine[object, object, T], *, debug: bool | None = None) -> T: - """Runs a coroutine using uvloop if available. - - Otherwise uses plain `asyncio.run`. - - Args: - coro: The asynchronous coroutine to run. - debug: If True, run in debug mode. - """ - try: - # Lazy import: uvloop is optional and only loaded if available - import uvloop # noqa: PLC0415 - - logger.debug('βœ… Using uvloop (recommended)') - return uvloop.run(coro, debug=debug) - except ImportError as e: - logger.debug( - '❓ Using asyncio (install uvloop for better performance)', - error=e, - ) - return asyncio.run(coro, debug=debug) diff --git a/py/packages/genkit/src/genkit/blocks/__init__.py b/py/packages/genkit/src/genkit/blocks/__init__.py deleted file mode 100644 index bcc72a370b..0000000000 --- a/py/packages/genkit/src/genkit/blocks/__init__.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Building blocks for the Genkit framework. - -This package provides the core building blocks for AI applications, including -models, prompts, embeddings, retrievers, and tools. These blocks are composed -to create intelligent applications. - -Overview: - The blocks package contains the fundamental components used to build - Genkit applications. Each block type represents a specific AI capability - that can be composed together. - -Key Modules: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Module β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ model β”‚ Model registration and invocation (ai.generate) β”‚ - β”‚ prompt β”‚ Prompt templates and ExecutablePrompt β”‚ - β”‚ embedding β”‚ Text embeddings (ai.embed) β”‚ - β”‚ retriever β”‚ Document retrieval for RAG (ai.retrieve) β”‚ - β”‚ document β”‚ Document model for content + metadata β”‚ - β”‚ tools β”‚ Tool context and interrupt handling β”‚ - β”‚ evaluator β”‚ Evaluation functions for quality assessment β”‚ - β”‚ reranker β”‚ Document reranking for improved retrieval β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Usage: - Most users interact with blocks via the ``Genkit`` class methods: - - ```python - from genkit import Genkit - - ai = Genkit(...) - - # Model block (via ai.generate) - response = await ai.generate(prompt='Hello!') - - # Prompt block (via ai.prompt) - prompt = ai.prompt('greet', source='...', model='...') - - # Embedding block (via ai.embed) - embeddings = await ai.embed(content='text') - - # Retriever block (via ai.retrieve) - docs = await ai.retrieve(retriever='my_retriever', query='...') - ``` - -See Also: - - genkit.ai: Main Genkit class - - genkit.types: Type definitions -""" - - -def package_name() -> str: - """Get the fully qualified package name. - - Returns: - The string 'genkit.blocks', which is the fully qualified package name. - """ - return 'genkit.blocks' - - -__all__ = ['package_name'] diff --git a/py/packages/genkit/src/genkit/blocks/dap.py b/py/packages/genkit/src/genkit/blocks/dap.py deleted file mode 100644 index 23d6ac638a..0000000000 --- a/py/packages/genkit/src/genkit/blocks/dap.py +++ /dev/null @@ -1,516 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Dynamic Action Provider (DAP) support for Genkit. - -Dynamic Action Providers allow external systems to supply actions at runtime, -enabling integration with dynamic tooling systems like MCP (Model Context -Protocol) servers, plugin marketplaces, or other dynamic action sources. - -Overview -======== - -A Dynamic Action Provider is a registered action that acts as a factory for -other actions. When Genkit needs to resolve an action that isn't statically -registered, it queries relevant DAPs to see if they can provide it. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Dynamic Action Provider Flow β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Request β”‚ β”‚ DAP β”‚ β”‚ External β”‚ β”‚ - β”‚ β”‚ Action β”‚ ───► β”‚ Cache β”‚ ───► β”‚ System β”‚ β”‚ - β”‚ β”‚ Resolution β”‚ β”‚ (TTL) β”‚ β”‚ (e.g. MCP) β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ β–Ό β–Ό β”‚ - β”‚ β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ └───────────── β”‚ Return β”‚ ◄─── β”‚ Actions β”‚ β”‚ - β”‚ β”‚ Action β”‚ β”‚ Created β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Concepts -============ - -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ Concept β”‚ Description β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ DAP β”‚ Dynamic Action Provider - a factory for actions β”‚ -β”‚ DapFn β”‚ Async function that returns actions by type β”‚ -β”‚ DapConfig β”‚ Configuration including name, description, caching β”‚ -β”‚ DapValue β”‚ Dictionary mapping action types to action lists β”‚ -β”‚ Cache TTL β”‚ Time-to-live for cached actions (default: 3000ms) β”‚ -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Use Cases -========= - -1. **MCP Integration**: Connect to MCP servers that provide tools/resources -2. **Plugin Marketplaces**: Load tools from external plugin systems -3. **Multi-tenant Systems**: Provide tenant-specific actions dynamically -4. **Feature Flags**: Enable/disable actions based on runtime configuration - -Example: - ```python - from genkit.ai import Genkit - from genkit.blocks.dap import define_dynamic_action_provider - - ai = Genkit() - - - # Define a DAP that provides tools from an external source - async def get_mcp_tools(): - # Connect to MCP server and get available tools - tools = await mcp_client.list_tools() - return { - 'tool': [ - ai.dynamic_tool( - name=t.name, - description=t.description, - fn=lambda input: mcp_client.call_tool(t.name, input), - ) - for t in tools - ] - } - - - dap = define_dynamic_action_provider( - ai.registry, - config={'name': 'mcp-tools', 'cache_config': {'ttl_millis': 5000}}, - fn=get_mcp_tools, - ) - ``` - -Caveats: - - DAPs are queried during action resolution, so they should be fast - - Cache TTL balances freshness with performance - choose wisely - - Actions from DAPs have the DAP name prefixed for identification - -See Also: - - MCP Plugin: genkit.plugins.mcp for Model Context Protocol integration - - JS Implementation: js/core/src/dynamic-action-provider.ts -""" - -import asyncio -import time -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, field -from typing import Any - -from genkit.core.action import Action -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry - -ActionMetadataLike = Mapping[str, object] -"""Type alias for action metadata - any string-keyed mapping. - -This type represents objects that behave like action metadata dictionaries, -providing key-based access to action properties like 'name', 'description', -'type', etc. - -Examples of compatible types: - - dict[str, object] - - dict[str, Any] - - ActionMetadata (from genkit.core.action) - - Any other Mapping[str, object] -""" - -DapValue = dict[str, list[Action[Any, Any]]] -"""Dictionary mapping action type names to lists of actions.""" - -DapFn = Callable[[], Awaitable[DapValue]] -"""Async function that returns actions organized by type.""" - -DapMetadata = dict[str, list[ActionMetadataLike]] -"""Dictionary mapping action type names to lists of action metadata.""" - - -@dataclass -class DapCacheConfig: - """Configuration for DAP caching behavior. - - Attributes: - ttl_millis: Time-to-live for cache in milliseconds. - - Negative: No caching (always fetch fresh) - - Zero/None: Default (3000 milliseconds) - - Positive: Cache validity duration in milliseconds - """ - - ttl_millis: int | None = None - - -@dataclass -class DapConfig: - """Configuration for a Dynamic Action Provider. - - Attributes: - name: Unique name for this DAP (used as prefix for actions). - description: Human-readable description of what this DAP provides. - cache_config: Optional caching configuration. - metadata: Additional metadata to attach to the DAP action. - """ - - name: str - description: str | None = None - cache_config: DapCacheConfig | None = None - metadata: dict[str, Any] = field(default_factory=dict) - - -class SimpleCache: - """Thread-safe cache for DAP values with TTL expiration. - - This cache ensures that concurrent requests for the same data share - a single fetch operation, preventing thundering herd problems. - """ - - def __init__( - self, - dap: 'DynamicActionProvider', - config: DapConfig, - dap_fn: DapFn, - ) -> None: - """Initialize the cache. - - Args: - dap: The parent DAP action. - config: DAP configuration including TTL. - dap_fn: Function to fetch actions from the external source. - """ - self._dap = dap - self._dap_fn = dap_fn - self._value: DapValue | None = None - self._expires_at: float | None = None - self._fetch_task: asyncio.Task[DapValue] | None = None - - # Determine TTL (default 3000ms) - ttl = config.cache_config.ttl_millis if config.cache_config else None - self._ttl_millis = 3000 if ttl is None or ttl == 0 else ttl - - async def get_or_fetch(self, skip_trace: bool = False) -> DapValue: - """Get cached value or fetch fresh data if stale. - - This method handles concurrent requests by sharing a single - fetch operation across all waiters. - - Args: - skip_trace: If True, don't run the DAP action (used by DevUI - to avoid excessive trace entries). - - Returns: - The DAP value containing actions by type. - """ - # Check if cache is still valid - is_stale = ( - self._value is None - or self._expires_at is None - or self._ttl_millis < 0 - or time.time() * 1000 > self._expires_at - ) - - if not is_stale and self._value is not None: - return self._value - - # If there's already a fetch in progress, wait for it - if self._fetch_task is not None: - return await self._fetch_task - - # Start a new fetch - self._fetch_task = asyncio.create_task(self._do_fetch(skip_trace)) - try: - return await self._fetch_task - finally: - self._fetch_task = None - - async def _do_fetch(self, skip_trace: bool) -> DapValue: - """Perform the actual fetch operation. - - Args: - skip_trace: If True, skip running the DAP action. - - Returns: - Fresh DAP value. - """ - try: - self._value = await self._dap_fn() - self._expires_at = time.time() * 1000 + self._ttl_millis - - # Run the DAP action for tracing (unless skipped) - if not skip_trace: - metadata = transform_dap_value(self._value) - await self._dap.action.arun(metadata) - - return self._value - except Exception: - self.invalidate() - raise - - def invalidate(self) -> None: - """Invalidate the cache, forcing a fresh fetch on next access.""" - self._value = None - self._expires_at = None - - -def transform_dap_value(value: DapValue) -> DapMetadata: - """Transform DAP value to metadata format for logging. - - Args: - value: DAP value with actions. - - Returns: - DAP metadata with action metadata. - """ - metadata: DapMetadata = {} - for action_type, actions in value.items(): - action_metadata_list: list[ActionMetadataLike] = [] - for action in actions: - # Action.metadata is dict[str, object] which satisfies ActionMetadataLike - meta: ActionMetadataLike = action.metadata if action.metadata else {} - action_metadata_list.append(meta) - metadata[action_type] = action_metadata_list - return metadata - - -class DynamicActionProvider: - """A Dynamic Action Provider that lazily resolves actions. - - This class wraps a DAP function and provides methods to query - for actions by type and name, with caching for performance. - """ - - def __init__( - self, - action: Action[Any, Any], - config: DapConfig, - dap_fn: DapFn, - ) -> None: - """Initialize the DAP. - - Args: - action: The underlying DAP action. - config: DAP configuration. - dap_fn: Function to fetch actions. - """ - self.action = action - self.config = config - self._cache = SimpleCache(self, config, dap_fn) - - def invalidate_cache(self) -> None: - """Invalidate the cache, forcing a fresh fetch on next access.""" - self._cache.invalidate() - - async def get_action( - self, - action_type: str, - action_name: str, - ) -> Action[Any, Any] | None: - """Get a specific action by type and name. - - Args: - action_type: The type of action (e.g., 'tool', 'model'). - action_name: The name of the action. - - Returns: - The action if found, None otherwise. - """ - result = await self._cache.get_or_fetch() - actions = result.get(action_type, []) - for action in actions: - if action.name == action_name: - return action - return None - - async def list_action_metadata( - self, - action_type: str, - action_name: str, - ) -> list[ActionMetadataLike]: - """List metadata for actions matching type and name pattern. - - Args: - action_type: The type of action. - action_name: Name or pattern to match: - - '*': Match all actions of this type - - 'prefix*': Match actions starting with prefix - - 'exact': Match only this exact name - - Returns: - List of matching action metadata. - """ - result = await self._cache.get_or_fetch() - actions = result.get(action_type, []) - if not actions: - return [] - - metadata_list: list[ActionMetadataLike] = [] - for action in actions: - meta: ActionMetadataLike = action.metadata if action.metadata else {} - metadata_list.append(meta) - - # Match all - if action_name == '*': - return metadata_list - - # Prefix match - if action_name.endswith('*'): - prefix = action_name[:-1] - return [m for m in metadata_list if str(m.get('name', '')).startswith(prefix)] - - # Exact match - return [m for m in metadata_list if m.get('name') == action_name] - - async def get_action_metadata_record( - self, - dap_prefix: str, - ) -> dict[str, ActionMetadataLike]: - """Get all actions as a metadata record for reflection API. - - This is used by the DevUI to list available actions. - - Args: - dap_prefix: Prefix to add to action keys. - - Returns: - Dictionary mapping action keys to metadata. - """ - dap_actions: dict[str, ActionMetadataLike] = {} - - # Skip trace to avoid excessive DevUI trace entries - result = await self._cache.get_or_fetch(skip_trace=True) - - for action_type, actions in result.items(): - for action in actions: - if not action.name: - raise ValueError(f'Invalid metadata when listing dynamic actions from {dap_prefix} - name required') - key = f'{dap_prefix}:{action_type}/{action.name}' - dap_actions[key] = action.metadata if action.metadata else {} - - return dap_actions - - -def is_dynamic_action_provider(obj: object) -> bool: - """Check if an object is a Dynamic Action Provider. - - Args: - obj: Object to check. - - Returns: - True if the object is a DAP. - """ - if isinstance(obj, DynamicActionProvider): - return True - if hasattr(obj, 'metadata'): - metadata = getattr(obj, 'metadata', None) - if isinstance(metadata, dict): - return metadata.get('type') == 'dynamic-action-provider' - return False - - -def define_dynamic_action_provider( - registry: Registry, - config: DapConfig | str, - fn: DapFn, -) -> DynamicActionProvider: - """Define and register a Dynamic Action Provider. - - A DAP is a factory that can dynamically provide actions at runtime. - This is useful for integrating with external systems like MCP servers - or plugin marketplaces. - - Args: - registry: The registry to register the DAP with. - config: DAP configuration or just a name string. - fn: Async function that returns actions organized by type. - - Returns: - The registered DynamicActionProvider. - - Example: - ```python - # Simple DAP that provides tools - async def get_tools(): - return { - 'tool': [ - ai.dynamic_tool(name='tool1', fn=...), - ai.dynamic_tool(name='tool2', fn=...), - ] - } - - - dap = define_dynamic_action_provider( - registry, - config='my-tools', - fn=get_tools, - ) - - # DAP with custom caching - dap = define_dynamic_action_provider( - registry, - config=DapConfig( - name='mcp-tools', - description='Tools from MCP server', - cache_config=DapCacheConfig(ttl_millis=10000), - ), - fn=get_tools, - ) - ``` - - See Also: - - JS implementation: js/core/src/dynamic-action-provider.ts - """ - # Normalize config - cfg = DapConfig(name=config) if isinstance(config, str) else config - - # Create metadata with DAP type marker - action_metadata = { - **cfg.metadata, - 'type': 'dynamic-action-provider', - } - - # Define the underlying action - # The action itself just returns its input (for logging purposes) - async def dap_action(input: DapMetadata) -> DapMetadata: - return input - - action = registry.register_action( - name=cfg.name, - kind=ActionKind.DYNAMIC_ACTION_PROVIDER, - description=cfg.description, - fn=dap_action, - metadata=action_metadata, - ) - - # Wrap in DynamicActionProvider - dap = DynamicActionProvider(action, cfg, fn) - - return dap - - -__all__ = [ - 'ActionMetadataLike', - 'DapCacheConfig', - 'DapConfig', - 'DapFn', - 'DapMetadata', - 'DapValue', - 'DynamicActionProvider', - 'SimpleCache', - 'define_dynamic_action_provider', - 'is_dynamic_action_provider', - 'transform_dap_value', -] diff --git a/py/packages/genkit/src/genkit/blocks/document.py b/py/packages/genkit/src/genkit/blocks/document.py deleted file mode 100644 index 0900ed86c6..0000000000 --- a/py/packages/genkit/src/genkit/blocks/document.py +++ /dev/null @@ -1,252 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Document model for the Genkit framework. - -This module provides types for creating and managing multi-modal documents in -Genkit. -""" - -from __future__ import annotations - -import warnings -from copy import deepcopy -from typing import Any - -from genkit.core.typing import ( - DocumentData, - DocumentPart, - Embedding, - Media, - MediaPart, - TextPart, -) - -TEXT_DATA_TYPE: str = 'text' - - -class Document(DocumentData): - """Represents document content along with its metadata. - - This object can be embedded, indexed or retrieved. Each document can contain - multiple parts (for example text and an image). - """ - - def __init__( - self, - content: list[DocumentPart], - metadata: dict[str, Any] | None = None, - ) -> None: - """Initializes a Document object. - - Performs a deep copy of the provided content and metadata to prevent - unintended modifications to the original objects. - - Args: - content: A list of DocumentPart objects representing the document's content. - metadata: An optional dictionary containing metadata about the document. - """ - doc_content = deepcopy(content) - doc_metadata = deepcopy(metadata) - super().__init__(content=doc_content, metadata=doc_metadata) - - @staticmethod - def from_document_data(document_data: DocumentData) -> Document: - """Constructs a Document instance from a DocumentData object. - - Args: - document_data: The DocumentData object containing content and metadata. - - Returns: - A new Document instance initialized with the provided data. - """ - return Document(content=document_data.content, metadata=document_data.metadata) - - @staticmethod - def from_text(text: str, metadata: dict[str, Any] | None = None) -> Document: - """Constructs a Document instance from a single text string. - - Args: - text: The text content for the document. - metadata: Optional metadata for the document. - - Returns: - A new Document instance containing a single text part. - """ - # NOTE: DocumentPart is a RootModel requiring root=TextPart(...) syntax. - return Document(content=[DocumentPart(root=TextPart(text=text))], metadata=metadata) - - @staticmethod - def from_media( - url: str, - content_type: str | None = None, - metadata: dict[str, Any] | None = None, - ) -> Document: - """Constructs a Document instance from a single media URL. - - Args: - url: The URL of the media content. - content_type: Optional MIME type of the media. - metadata: Optional metadata for the document. - - Returns: - A new Document instance containing a single media part. - """ - return Document( - # NOTE: DocumentPart is a RootModel requiring root=MediaPart(...) syntax. - # Using contentType alias for ty type checker compatibility. - content=[DocumentPart(root=MediaPart(media=Media(url=url, content_type=content_type)))], - metadata=metadata, - ) - - @staticmethod - def from_data( - data: str, - data_type: str | None = None, - metadata: dict[str, Any] | None = None, - ) -> Document: - """Constructs a Document instance from a single data string, determining type. - - If `data_type` is 'text', creates a text document. Otherwise, assumes - `data` is a URL and creates a media document. - - Args: - data: The data string (either text content or a media URL). - data_type: The type of the data ('text' or a media content type). - metadata: Optional metadata for the document. - - Returns: - A new Document instance. - """ - if data_type == TEXT_DATA_TYPE: - return Document.from_text(data, metadata) - return Document.from_media(data, data_type, metadata) - - def text(self) -> str: - """Concatenates all text parts of the document's content. - - Returns: - A single string containing the text from all text parts, joined - without delimiters. - """ - texts = [] - for p in self.content: - # Handle both TextPart objects and potential dict representations - # p.root is the underlying TextPart or MediaPart - part = p.root if hasattr(p, 'root') else p - text_val = getattr(part, 'text', None) - if isinstance(text_val, str): - texts.append(text_val) - return ''.join(texts) - - def media(self) -> list[Media]: - """Retrieves all media parts from the document's content. - - Returns: - A list of Media objects contained within the document. - """ - return [ - part.root.media for part in self.content if isinstance(part.root, MediaPart) and part.root.media is not None - ] - - def data(self) -> str: - """Gets the primary data content of the document. - - Returns the concatenated text if available, otherwise the URL of the - first media part. Returns an empty string if the document has neither - text nor media. - - Returns: - The primary data string (text or media URL). - """ - if self.text(): - return self.text() - - if self.media(): - return self.media()[0].url - - return '' - - def data_type(self) -> str | None: - """Gets the data type corresponding to the primary data content. - - Returns 'text' if the primary data is text. If the primary data is media, - returns the content type of the first media part. Returns None if the - document has no primary data or the media content type is not set. - - Returns: - The data type string ('text' or a MIME type) or None. - """ - if self.text(): - return TEXT_DATA_TYPE - - if self.media() and self.media()[0].content_type: - return self.media()[0].content_type - - return None - - def get_embedding_documents(self, embeddings: list[Embedding]) -> list[Document]: - """Creates multiple Document instances from a single document and its embeddings. - - Since embedders can return multiple embeddings for one input document, - but storage often requires a 1:1 document-to-embedding relationship, - this method duplicates the original document for each embedding. - Embedding metadata is added to the corresponding document's metadata - under the 'embedMetadata' key. - - Args: - embeddings: A list of Embedding objects generated for this document. - - Returns: - A list of Document objects, one for each provided embedding. - """ - documents = [] - for embedding in embeddings: - content = deepcopy(self.content) - metadata = deepcopy(self.metadata) - if embedding.metadata: - if not metadata: - metadata = {} - metadata['embedMetadata'] = embedding.metadata - documents.append(Document(content=content, metadata=metadata)) - _ = check_unique_documents(documents) - return documents - - -def check_unique_documents(documents: list[Document]) -> bool: - """Checks if a list of documents contains duplicates based on their JSON representation. - - Prints a warning if duplicates are found, as this can cause issues with - vector storage systems that key documents by their content hash. - Duplicate documents might arise if embeddings are generated without unique - metadata for different parts or aspects of the same original document. - - Args: - documents: A list of Document objects to check. - - Returns: - True if all documents are unique, False otherwise (primarily for testing). - """ - seen = set() - for doc in documents: - if doc.model_dump_json() in seen: - warnings.warn( - 'Embedding documents are not unique. Are you missing embed metadata?', - stacklevel=2, - ) - return False - seen.add(doc.model_dump_json()) - return True diff --git a/py/packages/genkit/src/genkit/blocks/embedding.py b/py/packages/genkit/src/genkit/blocks/embedding.py deleted file mode 100644 index 9ca93989bf..0000000000 --- a/py/packages/genkit/src/genkit/blocks/embedding.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Embedding types and utilities for Genkit. - -This module provides types and utilities for working with text embeddings -in Genkit. Embeddings are numerical vector representations of text that -capture semantic meaning, enabling similarity search, clustering, and -other AI applications. - -Terminology: - - +-------------------+------------------------------------------------------+ - | Term | Description | - +===================+======================================================+ - | Embedding | A numerical vector (list of floats) representing | - | | semantic meaning. Similar texts produce similar | - | | vectors. Contains 'embedding' field and metadata. | - +-------------------+------------------------------------------------------+ - | Embedder | A model/service that converts text to embeddings. | - | | Examples: 'googleai/text-embedding-004'. Registered | - | | as actions, invoked via embed() and embed_many(). | - +-------------------+------------------------------------------------------+ - | EmbedderRef | Reference bundling embedder name with optional | - | | config and version. Useful for reusing settings. | - +-------------------+------------------------------------------------------+ - | Document | Structured content to embed. Create via | - | | Document.from_text(). Can include metadata. | - +-------------------+------------------------------------------------------+ - | Dimensions | Size of embedding vector (e.g., 768 or 1536 floats). | - | | Higher dimensions = more nuance but more storage. | - +-------------------+------------------------------------------------------+ - -Key Components: - - +-------------------+------------------------------------------------------+ - | Component | Description | - +===================+======================================================+ - | EmbedderRef | Reference to an embedder with optional configuration | - +-------------------+------------------------------------------------------+ - | EmbedderOptions | Configuration options for defining embedders | - +-------------------+------------------------------------------------------+ - | EmbedderSupports | Declares what input types an embedder supports | - +-------------------+------------------------------------------------------+ - | Embedder | Runtime wrapper around an embedder action | - +-------------------+------------------------------------------------------+ - | create_embedder | Factory function for creating embedder references | - | _ref | | - +-------------------+------------------------------------------------------+ - -Usage with Genkit: - The primary way to use embeddings is through the Genkit class methods: - - - ai.embed(): Embed a single piece of content - - ai.embed_many(): Embed multiple pieces of content in batch - -Example - Single embedding: - >>> embeddings = await ai.embed(embedder='googleai/text-embedding-004', content='Hello, world!') - >>> vector = embeddings[0].embedding - -Example - Batch embedding: - >>> embeddings = await ai.embed_many(embedder='googleai/text-embedding-004', content=['Doc 1', 'Doc 2', 'Doc 3']) - -Example - Using EmbedderRef with configuration: - >>> ref = create_embedder_ref('googleai/text-embedding-004', config={'task_type': 'CLUSTERING'}, version='v1') - >>> embeddings = await ai.embed(embedder=ref, content='My text') - -Note on embed() vs embed_many(): - - embed() extracts config/version from EmbedderRef and merges with options - - embed_many() does NOT extract config from EmbedderRef - pass options directly - -See Also: - - genkit.ai.Genkit.embed: Single content embedding method - - genkit.ai.Genkit.embed_many: Batch embedding method - - genkit.core.typing.Embedding: The embedding result type -""" - -from collections.abc import Callable -from typing import Any, ClassVar, cast - -from pydantic import BaseModel, ConfigDict -from pydantic.alias_generators import to_camel - -from genkit.blocks.document import Document -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.schema import to_json_schema -from genkit.core.typing import DocumentData, EmbedRequest, EmbedResponse - - -class EmbedderSupports(BaseModel): - """Embedder capability support.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - input: list[str] | None = None - multilingual: bool | None = None - - -class EmbedderOptions(BaseModel): - """Configuration options for an embedder.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, alias_generator=to_camel) - - config_schema: dict[str, Any] | None = None - label: str | None = None - supports: EmbedderSupports | None = None - dimensions: int | None = None - - -class EmbedderRef(BaseModel): - """Reference to an embedder with configuration.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - name: str - config: Any | None = None - version: str | None = None - - -class Embedder: - """Runtime embedder wrapper around an embedder Action.""" - - def __init__(self, name: str, action: Action) -> None: - """Initialize the Embedder. - - Args: - name: The name of the embedder. - action: The underlying action to execute. - """ - self.name: str = name - self._action: Action = action - - async def embed( - self, - documents: list[Document], - options: dict[str, Any] | None = None, - ) -> EmbedResponse: - """Embed a list of documents. - - Args: - documents: The documents to embed. - options: Optional configuration for the embedding request. - - Returns: - The generated embedding response. - """ - # NOTE: Document subclasses DocumentData, so this is type-safe at runtime. - # NOTE: Document subclasses DocumentData, so this is type-safe at runtime. - return ( - await self._action.arun(EmbedRequest(input=cast(list['DocumentData'], documents), options=options)) - ).response - - -EmbedderFn = Callable[[EmbedRequest], EmbedResponse] - - -def embedder_action_metadata( - name: str, - options: EmbedderOptions | None = None, -) -> ActionMetadata: - """Creates metadata for an embedder action. - - Args: - name: The name of the embedder. - options: Configuration options for the embedder. - - Returns: - The action metadata for the embedder. - """ - options = options if options is not None else EmbedderOptions() - embedder_metadata_dict: dict[str, object] = {'embedder': {}} - embedder_info = cast(dict[str, object], embedder_metadata_dict['embedder']) - - if options.label: - embedder_info['label'] = options.label - - embedder_info['dimensions'] = options.dimensions - - if options.supports: - embedder_info['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True) - - embedder_info['customOptions'] = options.config_schema if options.config_schema else None - - return ActionMetadata( - kind=cast(ActionKind, ActionKind.EMBEDDER), - name=name, - input_json_schema=to_json_schema(EmbedRequest), - output_json_schema=to_json_schema(EmbedResponse), - metadata=embedder_metadata_dict, - ) - - -def create_embedder_ref(name: str, config: dict[str, Any] | None = None, version: str | None = None) -> EmbedderRef: - """Creates an EmbedderRef instance.""" - return EmbedderRef(name=name, config=config, version=version) diff --git a/py/packages/genkit/src/genkit/blocks/evaluator.py b/py/packages/genkit/src/genkit/blocks/evaluator.py deleted file mode 100644 index bd6cdb2da6..0000000000 --- a/py/packages/genkit/src/genkit/blocks/evaluator.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Evaluator type definitions for the Genkit framework. - -This module defines the type interfaces for evaluators in the Genkit framework. -Evaluators are used for assessint the quality of output of a Genkit flow or -model. -""" - -from collections.abc import Callable, Coroutine -from typing import Any, ClassVar, TypeVar, cast - -from pydantic import BaseModel, ConfigDict -from pydantic.alias_generators import to_camel - -from genkit.core.action import ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.schema import to_json_schema -from genkit.core.typing import ( - BaseDataPoint, - EvalFnResponse, - EvalRequest, -) - -T = TypeVar('T') - -# User-provided evaluator function that evaluates a single datapoint. -# Must be async (coroutine function). -EvaluatorFn = Callable[[BaseDataPoint, T], Coroutine[Any, Any, EvalFnResponse]] - -# User-provided batch evaluator function that evaluates an EvaluationRequest -BatchEvaluatorFn = Callable[[EvalRequest, T], Coroutine[Any, Any, list[EvalFnResponse]]] - - -class EvaluatorRef(BaseModel): - """Reference to an evaluator.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, alias_generator=to_camel) - - name: str - config_schema: dict[str, object] | None = None - - -def evaluator_ref(name: str, config_schema: dict[str, object] | None = None) -> EvaluatorRef: - """Create a reference to an evaluator. - - Args: - name: Name of the evaluator. - config_schema: Optional schema for evaluator configuration. - - Returns: - An EvaluatorRef instance. - """ - return EvaluatorRef(name=name, config_schema=config_schema) - - -def evaluator_action_metadata( - name: str, - config_schema: type | dict[str, Any] | None = None, -) -> ActionMetadata: - """Generates an ActionMetadata for evaluators. - - Args: - name: Name of the evaluator. - config_schema: Optional schema for evaluator configuration. - - Returns: - An ActionMetadata instance for the evaluator. - """ - return ActionMetadata( - kind=cast(ActionKind, ActionKind.EVALUATOR), - name=name, - input_json_schema=to_json_schema(EvalRequest), - output_json_schema=to_json_schema(list[EvalFnResponse]), - metadata={'evaluator': {'customOptions': to_json_schema(config_schema) if config_schema else None}}, - ) diff --git a/py/packages/genkit/src/genkit/blocks/formats/__init__.py b/py/packages/genkit/src/genkit/blocks/formats/__init__.py deleted file mode 100644 index 5f922ecccc..0000000000 --- a/py/packages/genkit/src/genkit/blocks/formats/__init__.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Genkit format package. Provides implementation for various formats like json, jsonl, etc. - -Genkit includes built-in formats to handle common output patterns. These formats can be -specified in `GenerateActionOptions` or `OutputConfig`. - -| Format | Description | Constraints | -| :--- | :--- | :--- | -| `json` | JSON object. Default when a schema is provided. | Enforced by `constrained: true`. | -| `text` | Returns the raw text output from the model. | No constraints. | -| `array` | JSON array of items. Useful for generating lists. | Enforced by `constrained: true`. | -| `enum` | Parses output as a single enum value. | Enforced by `constrained: true`. | -| `jsonl` | Newline-delimited JSON. Useful for streaming lists. | No constraints. | - -Usage Example: - - # JSON Format (default with schema) - ai.generate( - output=OutputConfig( - schema={'type': 'object', 'properties': {'foo': {'type': 'string'}}} - ) - ) - - # Array Format - ai.generate( - output=OutputConfig( - format='array', - schema={'type': 'array', 'items': {'type': 'string'}} - ) - ) - - # Enum Format - ai.generate( - output=OutputConfig( - format='enum', - schema={'type': 'string', 'enum': ['cat', 'dog']} - ) - ) -""" - -from genkit.blocks.formats.array import ArrayFormat -from genkit.blocks.formats.enum import EnumFormat -from genkit.blocks.formats.json import JsonFormat -from genkit.blocks.formats.jsonl import JsonlFormat -from genkit.blocks.formats.text import TextFormat -from genkit.blocks.formats.types import FormatDef, Formatter, FormatterConfig - - -def package_name() -> str: - """Get the fully qualified package name.""" - return 'genkit.blocks.formats' - - -built_in_formats = [ - ArrayFormat(), - EnumFormat(), - JsonFormat(), - JsonlFormat(), - TextFormat(), -] - - -__all__ = [ - 'ArrayFormat', - 'EnumFormat', - 'FormatDef', - 'Formatter', - 'FormatterConfig', - 'JsonFormat', - 'JsonlFormat', - 'TextFormat', - 'package_name', -] diff --git a/py/packages/genkit/src/genkit/blocks/interfaces.py b/py/packages/genkit/src/genkit/blocks/interfaces.py deleted file mode 100644 index 2c24d49e1e..0000000000 --- a/py/packages/genkit/src/genkit/blocks/interfaces.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Shared interfaces and typing helpers across blocks.""" - -from __future__ import annotations - -from typing import Generic, TypedDict, TypeVar - -InputT = TypeVar('InputT') -OutputT = TypeVar('OutputT') - - -class OutputConfigDict(TypedDict, total=False): - """TypedDict for output configuration when passed as a dict.""" - - format: str | None - content_type: str | None - instructions: bool | str | None - schema: type | dict[str, object] | None - constrained: bool | None - - -class Input(Generic[InputT]): - """Typed input configuration that preserves schema type information. - - This class provides a type-safe way to configure input schemas for prompts. - When you pass a Pydantic model as the schema, the prompt's input parameter - will be properly typed. - """ - - def __init__(self, schema: type[InputT]) -> None: - """Initialize typed input configuration. - - Args: - schema: The type/class for structured input. - """ - self.schema: type[InputT] = schema - - -class Output(Generic[OutputT]): - """Typed output configuration that preserves schema type information. - - This class provides a type-safe way to configure output options for generate(). - When you pass a Pydantic model or other type as the schema, the return type - of generate() will be properly typed. - """ - - def __init__( - self, - schema: type[OutputT], - format: str = 'json', - content_type: str | None = None, - instructions: bool | str | None = None, - constrained: bool | None = None, - ) -> None: - """Initialize typed output configuration. - - Args: - schema: The type/class for structured output. - format: Output format name. Defaults to 'json'. - content_type: Optional MIME content type. - instructions: Optional formatting instructions. - constrained: Whether to constrain output to schema. - """ - self.schema: type[OutputT] = schema - self.format: str = format - self.content_type: str | None = content_type - self.instructions: bool | str | None = instructions - self.constrained: bool | None = constrained - - def to_dict(self) -> OutputConfigDict: - """Convert to OutputConfigDict for internal use.""" - result: OutputConfigDict = {'schema': self.schema, 'format': self.format} - if self.content_type is not None: - result['content_type'] = self.content_type - if self.instructions is not None: - result['instructions'] = self.instructions - if self.constrained is not None: - result['constrained'] = self.constrained - return result diff --git a/py/packages/genkit/src/genkit/blocks/model.py b/py/packages/genkit/src/genkit/blocks/model.py deleted file mode 100644 index 00bea2f2e9..0000000000 --- a/py/packages/genkit/src/genkit/blocks/model.py +++ /dev/null @@ -1,544 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Model type definitions for the Genkit framework. - -This module defines the type interfaces for AI models in the Genkit framework. -These types ensure consistent interaction with different AI models and provide -type safety when working with model inputs and outputs. - -Example: - def my_model(request: GenerateRequest) -> GenerateResponse: - # Model implementation - return GenerateResponse(...) - - model_fn: ModelFn = my_model -""" - -from __future__ import annotations - -from collections.abc import Awaitable, Callable, Sequence -from functools import cached_property -from typing import Any, Generic, cast - -from pydantic import BaseModel, Field, PrivateAttr -from typing_extensions import TypeVar - -from genkit.core.action import ActionMetadata, ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.extract import extract_json -from genkit.core.schema import to_json_schema -from genkit.core.typing import ( - Candidate, - DocumentPart, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationUsage, - Media, - MediaModel, - Message, - ModelInfo, - Part, - Text, - ToolRequestPart, -) - -# type ModelFn = Callable[[GenerateRequest], GenerateResponse] -ModelFn = Callable[[GenerateRequest, ActionRunContext], GenerateResponse] - -# type ModelMiddlewareNext = Callable[[GenerateRequest, ActionRunContext], Awaitable[GenerateResponse]] -ModelMiddlewareNext = Callable[[GenerateRequest, ActionRunContext], Awaitable[GenerateResponse]] -# type ModelMiddleware = Callable[ -# [GenerateRequest, ActionRunContext, ModelMiddlewareNext], -# Awaitable[GenerateResponse], -# ] -ModelMiddleware = Callable[ - [GenerateRequest, ActionRunContext, ModelMiddlewareNext], - Awaitable[GenerateResponse], -] - -# TypeVar for generic output type in GenerateResponseWrapper -OutputT = TypeVar('OutputT', default=object) - - -class ModelReference(BaseModel): - """Reference to a model with configuration.""" - - name: str - config_schema: object | None = None - info: ModelInfo | None = None - version: str | None = None - config: dict[str, object] | None = None - - -class MessageWrapper(Message): - """A wrapper around the base Message type providing utility methods. - - This class extends the standard `Message` by adding convenient cached properties - like `text` (for concatenated text content) and `tool_requests`. - It stores the original message in the `_original_message` attribute. - """ - - def __init__( - self, - message: Message, - ) -> None: - """Initializes the MessageWrapper. - - Args: - message: The original Message object to wrap. - """ - super().__init__( - role=message.role, - content=message.content, - metadata=message.metadata, - ) - self._original_message: Message = message - - @cached_property - def text(self) -> str: - """Returns all text parts of the current chunk joined into a single string. - - Returns: - str: The combined text content from the current chunk. - """ - return text_from_message(self) - - @cached_property - def tool_requests(self) -> list[ToolRequestPart]: - """Returns all tool request parts of the response as a list. - - Returns: - list[ToolRequestPart]: list of tool requests present in this response. - """ - return [p.root for p in self.content if isinstance(p.root, ToolRequestPart)] - - @cached_property - def interrupts(self) -> list[ToolRequestPart]: - """Returns all interrupted tool request parts of the message as a list. - - Returns: - list[ToolRequestPart]: list of interrupted tool requests. - """ - return [p for p in self.tool_requests if p.metadata and p.metadata.root.get('interrupt')] - - -class GenerateResponseWrapper(GenerateResponse, Generic[OutputT]): - """A wrapper around GenerateResponse providing utility methods. - - Extends the base `GenerateResponse` with cached properties (`text`, `output`, - `messages`, `tool_requests`) and methods for validation (`assert_valid`, - `assert_valid_schema`). It also handles optional message/chunk parsing. - - When used with `Output[T]`, the `output` property is typed as `T`. - """ - - # _message_parser is a private attribute that Pydantic will ignore - _message_parser: Callable[[MessageWrapper], object] | None = PrivateAttr(None) - # _schema_type stores the Pydantic class for runtime validation - _schema_type: type[BaseModel] | None = PrivateAttr(None) - # Override the parent's message field with our wrapper type (intentional Liskov violation) - # pyrefly: ignore[bad-override] - Intentional covariant override for wrapper functionality - message: MessageWrapper | None = None # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__( - self, - response: GenerateResponse, - request: GenerateRequest, - message_parser: Callable[[MessageWrapper], object] | None = None, - schema_type: type[BaseModel] | None = None, - ) -> None: - """Initializes a GenerateResponseWrapper instance. - - Args: - response: The original GenerateResponse object. - request: The GenerateRequest object associated with the response. - message_parser: An optional function to parse the output from the message. - schema_type: Optional Pydantic model class for runtime validation. - """ - # Wrap the message if it's not already a MessageWrapper - wrapped_message: MessageWrapper | None = None - if response.message is not None: - wrapped_message = ( - MessageWrapper(response.message) - if not isinstance(response.message, MessageWrapper) - else response.message - ) - - super().__init__( - message=wrapped_message, - finish_reason=response.finish_reason, - finish_message=response.finish_message, - latency_ms=response.latency_ms, - usage=response.usage if response.usage is not None else GenerationUsage(), - custom=response.custom if response.custom is not None else {}, - request=request, - candidates=response.candidates, - operation=response.operation, - ) - # Set subclass-specific fields after parent initialization - self._message_parser = message_parser - self._schema_type = schema_type - - def assert_valid(self) -> None: - """Validates the basic structure of the response. - - Note: This method is currently a placeholder (TODO). - - Raises: - AssertionError: If the response structure is considered invalid. - """ - # TODO(#4343): implement - pass - - def assert_valid_schema(self) -> None: - """Validates that the response message conforms to any specified output schema. - - Note: This method is currently a placeholder (TODO). - - Raises: - AssertionError: If the response message does not conform to the schema. - """ - # TODO(#4343): implement - pass - - @cached_property - def text(self) -> str: - """Returns all text parts of the response joined into a single string. - - Returns: - str: The combined text content from the response. - """ - if self.message is None: - return '' - return self.message.text - - @cached_property - def output(self) -> OutputT: - """Parses out JSON data from the text parts of the response. - - When used with `Output[T]`, returns the parsed output typed as `T`. - If a schema_type was provided and the parsed output is a dict, - validates and returns a proper Pydantic model instance. - - Returns: - The parsed JSON data from the response, typed according to the schema. - """ - if self._message_parser and self.message is not None: - parsed = self._message_parser(self.message) - else: - parsed = extract_json(self.text) - - # If we have a schema type and the parsed output is a dict, validate and - # return a proper Pydantic instance. Skip if parsed is already the correct - # type or if it's not a dict (e.g., custom formats may return strings). - if self._schema_type is not None and parsed is not None and isinstance(parsed, dict): - return cast(OutputT, self._schema_type.model_validate(parsed)) - - return cast(OutputT, parsed) - - @cached_property - def messages(self) -> list[Message]: - """Returns all messages of the response, including request messages as a list. - - Returns: - list[Message]: list of messages. - """ - if self.message is None: - return list(self.request.messages) if self.request else [] - return [ - *(self.request.messages if self.request else []), - self.message._original_message, # pyright: ignore[reportPrivateUsage] - ] - - @cached_property - def tool_requests(self) -> list[ToolRequestPart]: - """Returns all tool request parts of the response as a list. - - Returns: - list[ToolRequestPart]: list of tool requests present in this response. - """ - if self.message is None: - return [] - return self.message.tool_requests - - @cached_property - def interrupts(self) -> list[ToolRequestPart]: - """Returns all interrupted tool request parts of the response as a list. - - Returns: - list[ToolRequestPart]: list of interrupted tool requests. - """ - if self.message is None: - return [] - return self.message.interrupts - - -class GenerateResponseChunkWrapper(GenerateResponseChunk): - """A wrapper around GenerateResponseChunk providing utility methods. - - Extends the base `GenerateResponseChunk` with cached properties for accessing - the text content of the current chunk (`text`), the accumulated text from all - previous chunks including the current one (`accumulated_text`), and parsed - output from the accumulated text (`output`). It also stores previous chunks. - """ - - # Field(exclude=True) means these fields are not included in serialization - previous_chunks: list[GenerateResponseChunk] = Field(default_factory=list, exclude=True) - chunk_parser: Callable[[GenerateResponseChunkWrapper], object] | None = Field(None, exclude=True) - - def __init__( - self, - chunk: GenerateResponseChunk, - previous_chunks: list[GenerateResponseChunk], - index: int, - chunk_parser: Callable[[GenerateResponseChunkWrapper], object] | None = None, - ) -> None: - """Initializes the GenerateResponseChunkWrapper. - - Args: - chunk: The raw GenerateResponseChunk to wrap. - previous_chunks: A list of preceding chunks in the stream. - index: The index of this chunk in the sequence of messages/chunks. - chunk_parser: An optional function to parse the output from the chunk. - """ - super().__init__( - role=chunk.role, - index=index, - content=chunk.content, - custom=chunk.custom, - aggregated=chunk.aggregated, - ) - # Set subclass-specific fields after parent initialization - self.previous_chunks = previous_chunks - self.chunk_parser = chunk_parser - - @cached_property - def text(self) -> str: - """Returns all text parts of the current chunk joined into a single string. - - Returns: - str: The combined text content from the current chunk. - """ - parts: list[str] = [] - for p in self.content: - text_val = p.root.text - if text_val is not None: - # Handle Text RootModel (access .root) or plain str - if isinstance(text_val, Text): - parts.append(str(text_val.root) if text_val.root is not None else '') - else: - parts.append(str(text_val)) - return ''.join(parts) - - @cached_property - def accumulated_text(self) -> str: - """Returns all text parts from previous chunks plus the latest chunk. - - Returns: - str: The combined text content from all chunks seen so far. - """ - parts: list[str] = [] - if self.previous_chunks: - for chunk in self.previous_chunks: - for p in chunk.content: - text_val = p.root.text - if text_val: - # Handle Text RootModel (access .root) or plain str - if isinstance(text_val, Text): - parts.append(str(text_val.root) if text_val.root is not None else '') - else: - parts.append(str(text_val)) - return ''.join(parts) + self.text - - @cached_property - def output(self) -> object: - """Parses out JSON data from the accumulated text parts of the response. - - Returns: - Any: The parsed JSON data from the accumulated chunks. - """ - if self.chunk_parser: - return self.chunk_parser(self) - return extract_json(self.accumulated_text) - - -class PartCounts(BaseModel): - """Stores counts of different types of media parts. - - Attributes: - characters: Total number of characters in text parts. - images: Total number of image parts. - videos: Total number of video parts. - audio: Total number of audio parts. - """ - - characters: int = 0 - images: int = 0 - videos: int = 0 - audio: int = 0 - - -def text_from_message(msg: Message) -> str: - """Extracts and concatenates text content from all parts of a Message. - - Args: - msg: The Message object. - - Returns: - A single string containing all text found in the message parts. - """ - return text_from_content(msg.content) - - -def text_from_content(content: Sequence[Part | DocumentPart]) -> str: - """Extracts and concatenates text content from a list of Parts or DocumentParts. - - Args: - content: A sequence of Part or DocumentPart objects. - - Returns: - A single string containing all text found in the parts. - """ - return ''.join(str(p.root.text) for p in content if hasattr(p.root, 'text') and p.root.text is not None) - - -def get_basic_usage_stats(input_: list[Message], response: Message | list[Candidate]) -> GenerationUsage: - """Calculates basic usage statistics based on input and output messages/candidates. - - Counts characters, images, videos, and audio files for both input and output. - - Args: - input_: A list of input Message objects. - response: Either a single output Message object or a list of Candidate objects. - - Returns: - A GenerationUsage object populated with the calculated counts. - """ - request_parts = [] - - for msg in input_: - request_parts.extend(msg.content) - - response_parts = [] - if isinstance(response, list): - for candidate in response: - response_parts.extend(candidate.message.content) - else: - response_parts = response.content - - input_counts = get_part_counts(parts=request_parts) - output_counts = get_part_counts(parts=response_parts) - - return GenerationUsage( - input_characters=input_counts.characters, - input_images=input_counts.images, - input_videos=input_counts.videos, - input_audio_files=input_counts.audio, - output_characters=output_counts.characters, - output_images=output_counts.images, - output_videos=output_counts.videos, - output_audio_files=output_counts.audio, - ) - - -def get_part_counts(parts: list[Part]) -> PartCounts: - """Counts the occurrences of different media types within a list of Parts. - - Iterates through the parts, summing character lengths and counting image, - video, and audio parts based on content type or data URL prefix. - - Args: - parts: A list of Part objects to analyze. - - Returns: - A PartCounts object containing the aggregated counts. - """ - part_counts = PartCounts() - - for part in parts: - text_val = part.root.text - if text_val: - # Handle Text RootModel (access .root) or plain str - if isinstance(text_val, Text): - part_counts.characters += len(str(text_val.root)) if text_val.root else 0 - else: - part_counts.characters += len(str(text_val)) - - media = part.root.media - - if media: - # Handle Media BaseModel vs MediaModel RootModel - if isinstance(media, Media): - content_type = media.content_type or '' - url = media.url or '' - elif isinstance(media, MediaModel) and hasattr(media.root, 'content_type'): - content_type = getattr(media.root, 'content_type', '') or '' - url = getattr(media.root, 'url', '') or '' - else: - content_type = '' - url = '' - is_image = content_type.startswith('image') or url.startswith('data:image') - is_video = content_type.startswith('video') or url.startswith('data:video') - is_audio = content_type.startswith('audio') or url.startswith('data:audio') - - part_counts.images += 1 if is_image else 0 - part_counts.videos += 1 if is_video else 0 - part_counts.audio += 1 if is_audio else 0 - - return part_counts - - -def model_action_metadata( - name: str, - info: dict[str, object] | None = None, - config_schema: type | dict[str, Any] | None = None, -) -> ActionMetadata: - """Generates an ActionMetadata for models.""" - info = info if info is not None else {} - return ActionMetadata( - kind=cast(ActionKind, ActionKind.MODEL), - name=name, - input_json_schema=to_json_schema(GenerateRequest), - output_json_schema=to_json_schema(GenerateResponse), - metadata={'model': {**info, 'customOptions': to_json_schema(config_schema) if config_schema else None}}, - ) - - -def model_ref( - name: str, - namespace: str | None = None, - info: ModelInfo | None = None, - version: str | None = None, - config: dict[str, object] | None = None, -) -> ModelReference: - """The factory function equivalent to export function modelRef(...). - - Args: - name: The model name. - namespace: Optional namespace to prefix the name. - info: Optional model info. - version: Optional model version. - config: Optional model configuration. - - Returns: - A ModelReference instance. - """ - # Logic: if (options.namespace && !name.startsWith(options.namespace + '/')) - final_name = f'{namespace}/{name}' if namespace and not name.startswith(f'{namespace}/') else name - - return ModelReference(name=final_name, info=info, version=version, config=config) diff --git a/py/packages/genkit/src/genkit/blocks/prompt.py b/py/packages/genkit/src/genkit/blocks/prompt.py deleted file mode 100644 index 838580c73c..0000000000 --- a/py/packages/genkit/src/genkit/blocks/prompt.py +++ /dev/null @@ -1,2446 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Prompt management and templating for the Genkit framework. - -This module provides the ExecutablePrompt class and related types for managing -AI prompts in Genkit. It enables defining reusable prompts with templates, -input/output schemas, tool configurations, and more. - -Overview: - Prompts in Genkit are reusable templates that can be executed against AI - models. They encapsulate the prompt text, system instructions, model - configuration, and other generation options. The ExecutablePrompt class - provides a callable interface matching the JavaScript SDK. - -Key Concepts: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Term β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ ExecutablePrompt β”‚ A prompt that can be called like a function β”‚ - β”‚ PromptGenerateOpts β”‚ Options to override prompt defaults at runtime β”‚ - β”‚ GenerateStreamResp β”‚ Response object with stream and response props β”‚ - β”‚ .prompt files β”‚ Dotprompt template files (YAML frontmatter + HBS) β”‚ - β”‚ Variant β”‚ Alternative version of a prompt (e.g., casual) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Prompt Execution Flow: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Define β”‚ ──► β”‚ Render β”‚ ──► β”‚ Generate β”‚ - β”‚ Prompt β”‚ β”‚ Template β”‚ β”‚ Response β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ β”‚ β”‚ - β”‚ ai.define_prompt() β”‚ prompt.render() β”‚ model.generate() - β”‚ or .prompt file β”‚ β”‚ - β–Ό β–Ό β–Ό - ExecutablePrompt GenerateActionOptions GenerateResponse - -Key Operations: - - Define prompts programmatically with `ai.define_prompt()` - - Load prompts from .prompt files with `load_prompt_folder()` - - Look up prompts by name with `ai.prompt()` - - Execute prompts with `await prompt(input)` or `prompt.stream(input)` - - Render prompts without executing with `await prompt.render(input)` - - Override options at runtime with `opts` parameter - -Example: - Basic usage with programmatic prompt: - - ```python - from genkit.ai import Genkit - - ai = Genkit(model='googleai/gemini-2.0-flash') - - # Define a prompt - recipe_prompt = ai.define_prompt( - name='recipe', - system='You are a helpful chef.', - prompt='Create a recipe for {{food}}.', - config={'temperature': 0.7}, - ) - - # Execute the prompt - response = await recipe_prompt({'food': 'pizza'}) - print(response.text) - - # Override options at runtime - response = await recipe_prompt( - {'food': 'salad'}, - opts={ - 'config': {'temperature': 0.5}, # Merged with prompt config - 'model': 'googleai/gemini-1.5-pro', # Override model - }, - ) - - # Stream the response - result = recipe_prompt.stream({'food': 'soup'}) - async for chunk in result.stream: - print(chunk.text, end='') - final = await result.response - ``` - - Using .prompt files (Dotprompt): - - ``` - # prompts/recipe.prompt - --- - model: googleai/gemini-2.0-flash - config: - temperature: 0.7 - input: - schema: - food: string - --- - You are a helpful chef. - Create a recipe for {{food}}. - ``` - - ```python - # Load and use the prompt - recipe = ai.prompt('recipe') - response = await recipe({'food': 'curry'}) - ``` - -Caveats: - - Config values are MERGED (not replaced) when using opts.config - - The `system` and `prompt` fields cannot be overridden via opts - - Message resolvers receive opts.messages as `history`, not appended - - Python uses snake_case (e.g., `as_tool()`) vs JS camelCase (`asTool()`) - -See Also: - - JavaScript implementation: js/ai/src/prompt.ts - - Dotprompt documentation: https://genkit.dev/docs/dotprompt -""" - -import asyncio -import os -import weakref -from collections.abc import AsyncIterable, Awaitable, Callable -from pathlib import Path -from typing import Any, ClassVar, Generic, TypedDict, TypeVar, cast, overload - -from dotpromptz.typing import ( - DataArgument, - PromptFunction, - PromptInputConfig, - PromptMetadata, -) -from pydantic import BaseModel, ConfigDict - -from genkit.aio import Channel, ensure_async -from genkit.blocks.generate import ( - StreamingCallback as ModelStreamingCallback, - generate_action, - to_tool_definition, -) -from genkit.blocks.interfaces import Input, Output -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - GenerateResponseWrapper, - ModelMiddleware, -) -from genkit.core.action import Action, ActionRunContext, create_action_key -from genkit.core.action.types import ActionKind -from genkit.core.error import GenkitError -from genkit.core.logging import get_logger -from genkit.core.registry import Registry -from genkit.core.schema import to_json_schema -from genkit.core.typing import ( - DocumentData, - GenerateActionOptions, - GenerateActionOutputConfig, - GenerateRequest, - GenerationCommonConfig, - Message, - OutputConfig, - Part, - Resume, - Role, - TextPart, - ToolChoice, - ToolRequestPart, - ToolResponsePart, -) - -logger = get_logger(__name__) - -# TypeVars for generic input/output typing -InputT = TypeVar('InputT') -OutputT = TypeVar('OutputT') - - -class OutputOptions(TypedDict, total=False): - """Output configuration options for prompt generation. - - This matches the JavaScript OutputOptions interface and allows overriding - output configuration when executing a prompt. - - Overview: - OutputOptions controls how the model's response should be formatted. - You can request JSON output, constrain it to a schema, or provide - custom formatting instructions. - - Attributes: - format: Output format. Common values: 'json', 'text'. Defaults to - model's native format if not specified. - content_type: MIME content type for the output (e.g., 'application/json'). - instructions: Instructions for formatting output. Can be: - - True: Use default formatting instructions - - False: Disable formatting instructions - - str: Custom instructions to append to prompt - schema: Output schema for structured output. Can be: - - A Python type/class (Pydantic model or dataclass) - - A dict representing JSON Schema - - A string referencing a registered schema name - json_schema: Direct JSON Schema definition for the output. - constrained: Whether to constrain model output to the schema. - When True, the model will be forced to output valid JSON - matching the schema (if supported by the model). - - Example: - ```python - from pydantic import BaseModel - - - class Recipe(BaseModel): - name: str - ingredients: list[str] - steps: list[str] - - - response = await prompt( - {'food': 'pizza'}, - opts={ - 'output': { - 'format': 'json', - 'schema': Recipe, - 'constrained': True, - } - }, - ) - recipe = response.output # Parsed Recipe object - ``` - """ - - format: str | None - content_type: str | None - instructions: bool | str | None - schema: type | dict[str, Any] | str | None - json_schema: dict[str, Any] | None - constrained: bool | None - - -class ResumeOptions(TypedDict, total=False): - """Options for resuming generation after an interrupt. - - This matches the JavaScript ResumeOptions interface and enables - human-in-the-loop workflows where tool execution can be paused - for user confirmation or input. - - Overview: - When a tool is defined as an "interrupt", the model's tool call - is not automatically executed. Instead, the response contains - the tool request, allowing your application to: - 1. Present the tool call to a user for approval - 2. Modify the tool arguments - 3. Provide a custom response without executing the tool - - ResumeOptions is used to continue generation after handling - the interrupt. - - Attributes: - respond: Tool response part(s) to respond to interrupt tool requests. - Each response must have a matching `name` (and `ref` if supplied) - for its corresponding tool request. - restart: Tool request part(s) to restart with additional metadata. - This re-executes the tool with `resumed` metadata passed to - the tool function. - metadata: Additional metadata to annotate the created tool message - under the "resume" key. - - Example: - ```python - # Define an interrupt tool for user confirmation - @ai.tool(name='book_flight', interrupt=True) - def book_flight(destination: str, date: str) -> str: - # This won't be called automatically - return f'Booked flight to {destination} on {date}' - - - # First generate - gets interrupted - response = await prompt({'request': 'Book a flight to Paris'}) - interrupt = response.interrupts[0] - - # Resume after user confirms - response = await prompt( - {'request': 'Book a flight to Paris'}, - opts={ - 'messages': response.messages, - 'resume': { - 'respond': book_flight.respond(interrupt, 'Confirmed: Flight booked to Paris'), - }, - }, - ) - ``` - - See Also: - - Interrupts documentation: https://genkit.dev/docs/tool-calling#pause-agentic-loops-with-interrupts - """ - - respond: ToolResponsePart | list[ToolResponsePart] | None - restart: ToolRequestPart | list[ToolRequestPart] | None - metadata: dict[str, Any] | None - - -class PromptGenerateOptions(TypedDict, total=False): - """Options for generating with a prompt at runtime. - - This matches the JavaScript PromptGenerateOptions type (GenerateOptions - minus 'prompt' and 'system' fields, which are defined by the prompt). - - Overview: - PromptGenerateOptions allows overriding a prompt's default configuration - when executing it. Options are passed as the second argument to - ExecutablePrompt.__call__(), render(), and stream(). - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Merge Behavior β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ config β”‚ MERGED: {...promptConfig, ...optsConfig} β”‚ - β”‚ metadata β”‚ MERGED: {...promptMetadata, ...optsMetadata} β”‚ - β”‚ All other fields β”‚ REPLACED: opts value overrides prompt value β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - Attributes: - model: Override the model to use for generation. Can be a model name - string like 'googleai/gemini-2.0-flash'. - config: Model configuration options (temperature, topK, etc.). - These are MERGED with the prompt's config, not replaced. - messages: Conversation history for multi-turn prompts. Behavior: - - If prompt has messages template: passed as 'history' to resolver - - If prompt has no messages: used directly as the messages - docs: Additional documents for RAG/grounding context. - tools: Override the tools available for this generation. - resources: Dynamic resources (MCP resources) to make available. - tool_choice: Tool selection strategy: - - 'auto': Model decides when to use tools (default) - - 'required': Model must use at least one tool - - 'none': Model cannot use tools - output: Override output configuration (format, schema, etc.). - resume: Options for resuming after an interrupt (human-in-the-loop). - return_tool_requests: If True, return tool calls without auto-executing. - Useful for custom tool handling or inspection. - max_turns: Maximum tool call iterations (default: 5). Limits - back-and-forth between model and tools. - on_chunk: Callback function called with each response chunk during - streaming. Signature: `(chunk: GenerateResponseChunkWrapper) -> None` - use: Middleware to apply to this generation request. - context: Additional context data passed to tools and sub-actions. - Useful for passing auth info, request metadata, etc. - step_name: Custom name for this generate call in trace views. - metadata: Additional metadata for the generation request. - - Example: - ```python - # Basic override - response = await prompt({'topic': 'AI'}, opts={'config': {'temperature': 0.9}}) - - # Multi-turn conversation - response = await prompt( - {'question': 'What about safety?'}, - opts={ - 'messages': previous_response.messages, # Continue conversation - 'config': {'temperature': 0.5}, - }, - ) - - - # Streaming with callback - def on_chunk(chunk): - print(chunk.text, end='', flush=True) - - - response = await prompt({'topic': 'Space'}, opts={'on_chunk': on_chunk}) - - # Override model and tools - response = await prompt( - {'task': 'analyze'}, - opts={ - 'model': 'googleai/gemini-1.5-pro', - 'tools': ['search', 'calculator'], - 'tool_choice': 'auto', - }, - ) - ``` - - Caveats: - - Cannot override 'prompt' or 'system' (defined by the prompt itself) - - Config is MERGED, not replaced - to clear a config value, set it explicitly - - Message handling depends on whether prompt defines a messages template - """ - - model: str | None - config: dict[str, Any] | GenerationCommonConfig | None - messages: list[Message] | None - docs: list[DocumentData] | None - tools: list[str] | None - resources: list[str] | None - tool_choice: ToolChoice | None - output: OutputOptions | None - resume: ResumeOptions | None - return_tool_requests: bool | None - max_turns: int | None - on_chunk: ModelStreamingCallback | None - use: list[ModelMiddleware] | None - context: dict[str, Any] | None - step_name: str | None - metadata: dict[str, Any] | None - - -class GenerateStreamResponse(Generic[OutputT]): - r"""Response from a streaming prompt execution. - - This class provides a consistent interface matching the JavaScript - GenerateStreamResponse, with both stream and response properties - accessible simultaneously. - - When the prompt has a typed output schema, `response` returns - `GenerateResponseWrapper[OutputT]` with typed `.output` property. - - Overview: - When you call `prompt.stream()`, you get a GenerateStreamResponse - that allows you to: - 1. Iterate over response chunks as they arrive (via `stream`) - 2. Await the complete response when streaming finishes (via `response`) - - This enables real-time UIs that show text as it's generated while - still having access to the complete response for logging, analysis, - or error handling. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Stream vs Response β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ stream β”‚ AsyncIterable - yields chunks as generated β”‚ - β”‚ response β”‚ Awaitable - resolves to complete response β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - Attributes: - stream: Async iterable of response chunks. Each chunk contains partial - text and metadata as it's generated by the model. - response: Awaitable that resolves to the complete GenerateResponseWrapper - once streaming is finished. Contains the full text, usage stats, - finish reason, and any tool calls. - - Example: - Basic streaming to console: - - ```python - result = prompt.stream({'topic': 'AI'}) - - # Stream chunks to console in real-time - async for chunk in result.stream: - print(chunk.text, end='', flush=True) - - # Get complete response for logging - final = await result.response - print(f'\\nFinish reason: {final.finish_reason}') - print(f'Token usage: {final.usage}') - ``` - - Streaming to a web response: - - ```python - async def generate_stream(request): - result = prompt.stream({'question': request.question}) - - async def event_stream(): - async for chunk in result.stream: - yield f'data: {chunk.text}\\n\\n' - yield 'data: [DONE]\\n\\n' - - return StreamingResponse(event_stream()) - ``` - - Getting response without consuming stream: - - ```python - result = prompt.stream({'topic': 'news'}) - - # You can await response directly - stream is consumed internally - final = await result.response - print(final.text) # Complete text - ``` - - Caveats: - - The stream can only be consumed once - - Awaiting `response` without consuming `stream` will still work - - If an error occurs during streaming, it's raised when awaiting `response` - - See Also: - - JavaScript GenerateStreamResponse: js/ai/src/generate.ts - """ - - def __init__( - self, - channel: Channel[GenerateResponseChunkWrapper, GenerateResponseWrapper[OutputT]], - response_future: asyncio.Future[GenerateResponseWrapper[OutputT]], - ) -> None: - """Initialize the stream response. - - Args: - channel: The channel providing response chunks. This is an async - iterable that yields GenerateResponseChunkWrapper objects. - response_future: Future that resolves to the complete response - when streaming is finished. - """ - self._channel: Channel[GenerateResponseChunkWrapper, GenerateResponseWrapper[OutputT]] = channel - self._response_future: asyncio.Future[GenerateResponseWrapper[OutputT]] = response_future - - @property - def stream(self) -> AsyncIterable[GenerateResponseChunkWrapper]: - """Get the async iterable of response chunks. - - Returns: - An async iterable that yields GenerateResponseChunkWrapper objects - as they are received from the model. Each chunk contains: - - text: The partial text generated so far - - index: The chunk index - - Additional metadata from the model - """ - return self._channel - - @property - def response(self) -> Awaitable[GenerateResponseWrapper[OutputT]]: - """Get the awaitable for the complete response. - - Returns: - An awaitable that resolves to a GenerateResponseWrapper containing: - - text: The complete generated text - - output: The typed output (when using Output[T]) - - messages: The full message history - - usage: Token usage statistics - - finish_reason: Why generation stopped (e.g., 'stop', 'length') - - Any tool calls or interrupts from the response - """ - return self._response_future - - -class PromptCache: - """Model for a prompt cache.""" - - user_prompt: PromptFunction[Any] | None = None - system: PromptFunction[Any] | None = None - messages: PromptFunction[Any] | None = None - - -class PromptConfig(BaseModel): - """Model for a prompt action.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) - - variant: str | None = None - model: str | None = None - config: dict[str, Any] | GenerationCommonConfig | None = None - description: str | None = None - input_schema: type | dict[str, Any] | str | None = None - system: str | Part | list[Part] | Callable[..., Any] | None = None - prompt: str | Part | list[Part] | Callable[..., Any] | None = None - messages: str | list[Message] | Callable[..., Any] | None = None - output_format: str | None = None - output_content_type: str | None = None - output_instructions: bool | str | None = None - output_schema: type | dict[str, Any] | str | None = None - output_constrained: bool | None = None - max_turns: int | None = None - return_tool_requests: bool | None = None - metadata: dict[str, Any] | None = None - tools: list[str] | None = None - tool_choice: ToolChoice | None = None - use: list[ModelMiddleware] | None = None - docs: list[DocumentData] | Callable[..., Any] | None = None - tool_responses: list[Part] | None = None - resources: list[str] | None = None - - -class ExecutablePrompt(Generic[InputT, OutputT]): - r"""A prompt that can be executed with a given input and configuration. - - This class matches the JavaScript ExecutablePrompt interface, providing - a callable object that generates AI responses from a prompt template. - - When defined with input/output schemas via `Input[I]` and `Output[O]`, - the prompt is typed as `ExecutablePrompt[I, O]`: - - Input is type-checked when calling the prompt - - Output is typed on `response.output` - - Overview: - ExecutablePrompt is the main way to work with prompts in Genkit. It - wraps a prompt definition (template, model, config, etc.) and provides - methods to execute it, stream responses, or render it without execution. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ ExecutablePrompt Methods β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ __call__(input,opts)β”‚ Execute prompt, return complete response β”‚ - β”‚ stream(input, opts) β”‚ Execute prompt, return streaming response β”‚ - β”‚ render(input, opts) β”‚ Render template without executing β”‚ - β”‚ as_tool() β”‚ Convert prompt to a tool action β”‚ - β”‚ ref β”‚ Property with prompt name and metadata β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - Attributes: - ref: A dict containing the prompt's name and metadata. - - Example: - Basic execution: - - ```python - # Get a prompt (from ai.define_prompt or ai.prompt) - recipe = ai.prompt('recipe') - - # Execute with input - response = await recipe({'food': 'pizza'}) - print(response.text) - - # Execute with options override - response = await recipe( - {'food': 'salad'}, - opts={ - 'config': {'temperature': 0.5}, - 'model': 'googleai/gemini-1.5-pro', - }, - ) - ``` - - Streaming: - - ```python - result = recipe.stream({'food': 'soup'}) - - async for chunk in result.stream: - print(chunk.text, end='') - - final = await result.response - print(f'\\nTokens used: {final.usage}') - ``` - - Rendering without execution: - - ```python - # Get the GenerateActionOptions without calling the model - options = await recipe.render({'food': 'curry'}) - print(options.messages) # See rendered messages - print(options.config) # See merged config - - # Manually execute - response = await ai.generate(options) - ``` - - Converting to a tool: - - ```python - # Use the prompt as a tool in another prompt - recipe_tool = await recipe.as_tool() - - response = await ai.generate( - prompt='Suggest a healthy meal', - tools=[recipe_tool], - ) - ``` - - Caveats: - - Config values passed via opts are MERGED with prompt config - - The prompt and system fields cannot be overridden at runtime - - Lazy resolution: prompts loaded from files are resolved on first use - - See Also: - - JavaScript ExecutablePrompt: js/ai/src/prompt.ts - - Dotprompt: https://genkit.dev/docs/dotprompt - """ - - def __init__( - self, - registry: Registry, - variant: str | None = None, - model: str | None = None, - config: dict[str, Any] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, Any] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, Any] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, Any] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - resources: list[str] | None = None, - _name: str | None = None, # prompt name for action lookup - _ns: str | None = None, # namespace for action lookup - _prompt_action: Action | None = None, # reference to PROMPT action - # TODO(#4344): - # docs: list[Document]): - ) -> None: - """Initializes an ExecutablePrompt instance. - - Args: - registry: The registry to use for resolving models and tools. - variant: The variant of the prompt. - model: The model to use for generation. - config: The generation configuration. - description: A description of the prompt. - input_schema: type | dict[str, Any] | str | None = None, - system: str | Part | list[Part] | Callable | None = None, - prompt: str | Part | list[Part] | Callable | None = None, - messages: str | list[Message] | Callable | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: Instructions for formatting the output. - output_schema: type | dict[str, Any] | str | None = None, - output_constrained: Whether the output should be constrained to the output schema. - max_turns: The maximum number of turns in a conversation. - return_tool_requests: Whether to return tool requests. - metadata: Metadata to associate with the prompt. - tools: A list of tool names to use with the prompt. - tool_choice: The tool choice strategy. - use: A list of model middlewares to apply. - docs: A list of documents to be used for grounding. - resources: A list of resource URIs to be used for grounding. - """ - self._registry: Registry = registry - self._variant: str | None = variant - self._model: str | None = model - self._config: dict[str, Any] | GenerationCommonConfig | None = config - self._description: str | None = description - self._input_schema: type | dict[str, Any] | str | None = input_schema - self._system: str | Part | list[Part] | Callable[..., Any] | None = system - self._prompt: str | Part | list[Part] | Callable[..., Any] | None = prompt - self._messages: str | list[Message] | Callable[..., Any] | None = messages - self._output_format: str | None = output_format - self._output_content_type: str | None = output_content_type - self._output_instructions: bool | str | None = output_instructions - self._output_schema: type | dict[str, Any] | str | None = output_schema - self._output_constrained: bool | None = output_constrained - self._max_turns: int | None = max_turns - self._return_tool_requests: bool | None = return_tool_requests - self._metadata: dict[str, Any] | None = metadata - self._tools: list[str] | None = tools - self._tool_choice: ToolChoice | None = tool_choice - self._use: list[ModelMiddleware] | None = use - self._docs: list[DocumentData] | Callable[..., Any] | None = docs - self._resources: list[str] | None = resources - self._cache_prompt: PromptCache = PromptCache() - self._name: str | None = _name # Store name/ns for action lookup (used by as_tool()) - self._ns: str | None = _ns - self._prompt_action: Action | None = _prompt_action - - @property - def ref(self) -> dict[str, Any]: - """Returns a reference object for this prompt. - - The reference object contains the prompt's name and metadata. - """ - return { - 'name': registry_definition_key(self._name, self._variant, self._ns) if self._name else None, - 'metadata': self._metadata, - } - - async def _ensure_resolved(self) -> None: - """Ensures the prompt is resolved from the registry if only a name was provided.""" - if self._prompt_action or not self._name: - return - - # Preserve Pydantic schema type if it was explicitly provided via ai.prompt(..., output=Output(schema=T)) - # The resolved prompt from .prompt file will have a dict schema, but we want to keep the Pydantic type - # for runtime validation to get proper typed output. - original_output_schema = self._output_schema - - resolved = await lookup_prompt(self._registry, self._name, self._variant) - self._model = resolved._model - self._config = resolved._config - self._description = resolved._description - self._input_schema = resolved._input_schema - self._system = resolved._system - self._prompt = resolved._prompt - self._messages = resolved._messages - self._output_format = resolved._output_format - self._output_content_type = resolved._output_content_type - self._output_instructions = resolved._output_instructions - # Keep original Pydantic type if provided, otherwise use resolved (dict) schema - if isinstance(original_output_schema, type) and issubclass(original_output_schema, BaseModel): - self._output_schema = original_output_schema - else: - self._output_schema = resolved._output_schema - self._output_constrained = resolved._output_constrained - self._max_turns = resolved._max_turns - self._return_tool_requests = resolved._return_tool_requests - self._metadata = resolved._metadata - self._tools = resolved._tools - self._tool_choice = resolved._tool_choice - self._use = resolved._use - self._docs = resolved._docs - self._resources = resolved._resources - self._prompt_action = resolved._prompt_action - - async def __call__( - self, - input: InputT | None = None, - opts: PromptGenerateOptions | None = None, - ) -> GenerateResponseWrapper[OutputT]: - """Executes the prompt with the given input and configuration. - - This method matches the JavaScript ExecutablePrompt callable interface, - accepting an optional `opts` parameter that can override the prompt's - default configuration. - - Args: - input: The input to the prompt template. When the prompt is defined - with `input=Input(schema=T)`, this should be an instance of T. - opts: Optional generation options to override prompt defaults. - Can include: model, config, messages, docs, tools, output, - tool_choice, return_tool_requests, max_turns, on_chunk, - use (middleware), context, resume, and metadata. - - Returns: - The generated response with typed output. - - Example: - ```python - # With typed input/output - class RecipeInput(BaseModel): - dish: str - - - prompt = ai.define_prompt( - name='recipe', - input=Input(schema=RecipeInput), - output=Output(schema=Recipe), - prompt='Create a recipe for {dish}', - ) - - response = await prompt(RecipeInput(dish='pizza')) - response.output.name # Typed! - ``` - """ - await self._ensure_resolved() - effective_opts: PromptGenerateOptions = opts if opts else {} - - # Extract streaming callback and middleware from opts - on_chunk = effective_opts.get('on_chunk') - middleware = effective_opts.get('use') or self._use - context = effective_opts.get('context') - - result = await generate_action( - self._registry, - await self.render(input=input, opts=effective_opts), - on_chunk=on_chunk, - middleware=middleware, - context=context if context else ActionRunContext._current_context(), # pyright: ignore[reportPrivateUsage] - ) - # Cast to preserve the generic type parameter - return cast(GenerateResponseWrapper[OutputT], result) - - def stream( - self, - input: InputT | None = None, - opts: PromptGenerateOptions | None = None, - *, - timeout: float | None = None, - ) -> GenerateStreamResponse[OutputT]: - r"""Streams the prompt execution with the given input and configuration. - - This method matches the JavaScript ExecutablePrompt.stream() interface, - returning a GenerateStreamResponse with both stream and response properties. - - Args: - input: The input to the prompt template. When the prompt is defined - with `input=Input(schema=T)`, this should be an instance of T. - opts: Optional generation options to override prompt defaults. - Can include: model, config, messages, docs, tools, output, - tool_choice, return_tool_requests, max_turns, use (middleware), - context, resume, and metadata. - timeout: Optional timeout in seconds for the streaming operation. - - Returns: - A GenerateStreamResponse with: - - stream: AsyncIterable of response chunks - - response: Awaitable that resolves to the typed complete response - - Example: - ```python - prompt = ai.define_prompt( - name='story', - input=Input(schema=StoryInput), - output=Output(schema=Story), - prompt='Write a story about {topic}', - ) - - # Stream the response - result = prompt.stream(StoryInput(topic='adventure')) - async for chunk in result.stream: - print(chunk.text, end='') - - # Get the final typed response - final = await result.response - print(f'Title: {final.output.title}') - ``` - """ - effective_opts: PromptGenerateOptions = opts if opts else {} - channel: Channel[GenerateResponseChunkWrapper, GenerateResponseWrapper[OutputT]] = Channel(timeout=timeout) - - # Create a copy of opts with the streaming callback - stream_opts: PromptGenerateOptions = {**effective_opts, 'on_chunk': lambda c: channel.send(c)} - - resp = self.__call__(input=input, opts=stream_opts) - response_future: asyncio.Future[GenerateResponseWrapper[OutputT]] = asyncio.create_task(resp) - channel.set_close_future(response_future) - - return GenerateStreamResponse[OutputT](channel=channel, response_future=response_future) - - async def render( - self, - input: InputT | dict[str, Any] | None = None, - opts: PromptGenerateOptions | None = None, - ) -> GenerateActionOptions: - """Renders the prompt template with the given input and options. - - This method matches the JavaScript ExecutablePrompt.render() interface, - accepting an optional `opts` parameter that can override the prompt's - default configuration. - - Args: - input: The input to the prompt template. Can be a typed input model - or a dict with template variables. - opts: Optional generation options to override prompt defaults. - Can include: model, config, messages, docs, tools, output, - tool_choice, return_tool_requests, max_turns, context, - resume, and metadata. - - Returns: - The rendered prompt as a GenerateActionOptions object, ready to - be passed to generate(). - - Example: - ```python - prompt = ai.define_prompt( - name='recipe', - input=Input(schema=RecipeInput), - prompt='Create a recipe for {dish}', - ) - - # Render without executing - options = await prompt.render(RecipeInput(dish='pizza')) - - # Then generate manually - response = await ai.generate(options) - ``` - """ - await self._ensure_resolved() - effective_opts: PromptGenerateOptions = opts if opts else {} - - # Extract context from opts - context = effective_opts.get('context') - - # Extract output options from opts - output_opts = effective_opts.get('output') or {} - - # Merge config: opts.config is MERGED with prompt config (JS behavior) - # This allows overriding specific config values while keeping others - opts_config_value = effective_opts.get('config') - if opts_config_value is not None: - prompt_config = self._config or {} - opts_config = opts_config_value or {} - # Convert Pydantic models to dicts for merging - if isinstance(prompt_config, BaseModel): - prompt_config = prompt_config.model_dump(exclude_none=True) - if isinstance(opts_config, BaseModel): - opts_config = opts_config.model_dump(exclude_none=True) - merged_config = ( - { - **(prompt_config if isinstance(prompt_config, dict) else {}), - **(opts_config if isinstance(opts_config, dict) else {}), - } - if prompt_config or opts_config - else None - ) - else: - merged_config = self._config - - # Merge model: opts.model overrides prompt model - merged_model = effective_opts.get('model') or self._model - - # Merge tools: opts.tools overrides prompt tools - merged_tools = effective_opts.get('tools') or self._tools - - # Merge output options: opts.output overrides prompt output settings - merged_output_format = output_opts.get('format') or self._output_format - merged_output_content_type = output_opts.get('content_type') or self._output_content_type - merged_output_schema = output_opts.get('schema') or output_opts.get('json_schema') or self._output_schema - merged_output_constrained = ( - output_opts.get('constrained') if output_opts.get('constrained') is not None else self._output_constrained - ) - merged_output_instructions = ( - output_opts.get('instructions') - if output_opts.get('instructions') is not None - else self._output_instructions - ) - - # Merge other options (opts values override prompt values) - merged_tool_choice = effective_opts.get('tool_choice') or self._tool_choice - merged_return_tool_requests = ( - effective_opts.get('return_tool_requests') - if effective_opts.get('return_tool_requests') is not None - else self._return_tool_requests - ) - merged_max_turns = ( - effective_opts.get('max_turns') if effective_opts.get('max_turns') is not None else self._max_turns - ) - merged_metadata = ( - {**(self._metadata or {}), **(effective_opts.get('metadata') or {})} - if effective_opts.get('metadata') - else self._metadata - ) - - # Build the merged PromptConfig - prompt_options = PromptConfig( - model=merged_model, - prompt=self._prompt, - system=self._system, - messages=self._messages, - tools=merged_tools, - return_tool_requests=merged_return_tool_requests, - tool_choice=merged_tool_choice, - config=merged_config, - max_turns=merged_max_turns, - output_format=merged_output_format, - output_content_type=merged_output_content_type, - output_instructions=merged_output_instructions, - output_schema=merged_output_schema, - output_constrained=merged_output_constrained, - input_schema=self._input_schema, - metadata=merged_metadata, - docs=self._docs, - resources=effective_opts.get('resources') or self._resources, - ) - - model = prompt_options.model or self._registry.default_model - if model is None: - raise GenkitError(status='INVALID_ARGUMENT', message='No model configured.') - - resolved_msgs: list[Message] = [] - # Convert input to dict for render functions - # If input is a Pydantic model, convert to dict; otherwise use as-is - render_input: dict[str, Any] - if input is None: - render_input = {} - elif isinstance(input, dict): - # Type narrow: input is dict here, assign to dict[str, Any] typed variable - render_input = {str(k): v for k, v in input.items()} - elif isinstance(input, BaseModel): - # Pydantic v2 model - render_input = input.model_dump() - elif hasattr(input, 'dict'): - # Pydantic v1 model - dict_func = getattr(input, 'dict', None) - render_input = cast(Callable[[], dict[str, Any]], dict_func)() - else: - # Fallback: cast to dict (should not happen with proper typing) - render_input = cast(dict[str, Any], input) - # Get opts.messages for history (matching JS behavior) - opts_messages = effective_opts.get('messages') - - # Render system prompt - if prompt_options.system: - result = await render_system_prompt( - self._registry, render_input, prompt_options, self._cache_prompt, context - ) - resolved_msgs.append(result) - - # Handle messages (matching JS behavior): - # - If prompt has messages template: render it (opts.messages passed as history to resolvers) - # - If prompt has no messages: use opts.messages directly - if prompt_options.messages: - # Prompt defines messages - render them (resolvers receive opts_messages as history) - resolved_msgs.extend( - await render_message_prompt( - self._registry, - render_input, - prompt_options, - self._cache_prompt, - context, - history=opts_messages, - ) - ) - elif opts_messages: - # Prompt has no messages template - use opts.messages directly - resolved_msgs.extend(opts_messages) - - # Render user prompt - if prompt_options.prompt: - result = await render_user_prompt(self._registry, render_input, prompt_options, self._cache_prompt, context) - resolved_msgs.append(result) - - # If schema is set but format is not explicitly set, default to 'json' format - if prompt_options.output_schema and not prompt_options.output_format: - output_format = 'json' - else: - output_format = prompt_options.output_format - - # Build output config - output = GenerateActionOutputConfig() - if output_format: - output.format = output_format - if prompt_options.output_content_type: - output.content_type = prompt_options.output_content_type - if prompt_options.output_instructions is not None: - output.instructions = prompt_options.output_instructions - _resolve_output_schema(self._registry, prompt_options.output_schema, output) - if prompt_options.output_constrained is not None: - output.constrained = prompt_options.output_constrained - - # Handle resume options - resume = None - resume_opts = effective_opts.get('resume') - if resume_opts: - respond = resume_opts.get('respond') - if respond: - resume = Resume(respond=respond) if isinstance(respond, list) else Resume(respond=[respond]) - - # Merge docs: opts.docs extends prompt docs - merged_docs = await render_docs(render_input, prompt_options, context) - opts_docs = effective_opts.get('docs') - if opts_docs: - merged_docs = [*merged_docs, *opts_docs] if merged_docs else list(opts_docs) - - return GenerateActionOptions( - model=model, - messages=resolved_msgs, - config=prompt_options.config, - tools=prompt_options.tools, - return_tool_requests=prompt_options.return_tool_requests, - tool_choice=prompt_options.tool_choice, - output=output, - max_turns=prompt_options.max_turns, - docs=merged_docs, - resume=resume, - ) - - async def as_tool(self) -> Action: - """Expose this prompt as a tool. - - Returns the PROMPT action, which can be used as a tool. - """ - await self._ensure_resolved() - # If we have a direct reference to the action, use it - if self._prompt_action is not None: - return self._prompt_action - - # Otherwise, try to look it up using name/variant/ns - if self._name is None: - raise GenkitError( - status='FAILED_PRECONDITION', - message=( - 'Prompt name not available. This prompt was not created via define_prompt_async() or load_prompt().' - ), - ) - - lookup_key = registry_lookup_key(self._name, self._variant, self._ns) - - action = await self._registry.resolve_action_by_key(lookup_key) - - if action is None or action.kind != ActionKind.PROMPT: - raise GenkitError( - status='NOT_FOUND', - message=f'PROMPT action not found for prompt "{self._name}"', - ) - - return action - - -# Overload 1: Both input and output typed -> ExecutablePrompt[InputT, OutputT] -@overload -def define_prompt( - registry: Registry, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, Any] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, Any] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, Any] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, Any] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - *, - input: 'Input[InputT]', - output: 'Output[OutputT]', -) -> 'ExecutablePrompt[InputT, OutputT]': ... - - -# Overload 2: Only input typed -> ExecutablePrompt[InputT, Any] -@overload -def define_prompt( - registry: Registry, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, Any] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, Any] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, Any] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, Any] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - *, - input: 'Input[InputT]', - output: None = None, -) -> 'ExecutablePrompt[InputT, Any]': ... - - -# Overload 3: Only output typed -> ExecutablePrompt[Any, OutputT] -@overload -def define_prompt( - registry: Registry, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, Any] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, Any] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, Any] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, Any] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - input: None = None, - *, - output: 'Output[OutputT]', -) -> 'ExecutablePrompt[Any, OutputT]': ... - - -# Overload 4: Neither typed -> ExecutablePrompt[Any, Any] -@overload -def define_prompt( - registry: Registry, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, Any] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, Any] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, Any] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, Any] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - input: None = None, - output: None = None, -) -> 'ExecutablePrompt[Any, Any]': ... - - -# Implementation -def define_prompt( - registry: Registry, - name: str | None = None, - variant: str | None = None, - model: str | None = None, - config: dict[str, Any] | GenerationCommonConfig | None = None, - description: str | None = None, - input_schema: type | dict[str, Any] | str | None = None, - system: str | Part | list[Part] | Callable[..., Any] | None = None, - prompt: str | Part | list[Part] | Callable[..., Any] | None = None, - messages: str | list[Message] | Callable[..., Any] | None = None, - output_format: str | None = None, - output_content_type: str | None = None, - output_instructions: bool | str | None = None, - output_schema: type | dict[str, Any] | str | None = None, - output_constrained: bool | None = None, - max_turns: int | None = None, - return_tool_requests: bool | None = None, - metadata: dict[str, Any] | None = None, - tools: list[str] | None = None, - tool_choice: ToolChoice | None = None, - use: list[ModelMiddleware] | None = None, - docs: list[DocumentData] | Callable[..., Any] | None = None, - input: 'Input[Any] | None' = None, - output: 'Output[Any] | None' = None, -) -> 'ExecutablePrompt[Any, Any]': - """Defines an executable prompt. - - Args: - registry: The registry to use for resolving models and tools. - name: The name of the prompt. - variant: The variant of the prompt. - model: The model to use for generation. - config: The generation configuration. - description: A description of the prompt. - input_schema: The input schema for the prompt. - system: The system message for the prompt. - prompt: The user prompt. - messages: A list of messages to include in the prompt. - output_format: The output format. - output_content_type: The output content type. - output_instructions: Instructions for formatting the output. - output_schema: The output schema (use `output` parameter for typed outputs). - output_constrained: Whether the output should be constrained to the output schema. - max_turns: The maximum number of turns in a conversation. - return_tool_requests: Whether to return tool requests. - metadata: Metadata to associate with the prompt. - tools: A list of tool names to use with the prompt. - tool_choice: The tool choice strategy. - use: A list of model middlewares to apply. - docs: A list of documents to be used for grounding. - input: Typed input configuration using Input[T]. When provided, the - prompt's input parameter is type-checked. - output: Typed output configuration using Output[T]. When provided, the - response output is typed. - - Returns: - An ExecutablePrompt instance. When both `input=Input(schema=I)` and - `output=Output(schema=O)` are provided, returns `ExecutablePrompt[I, O]` - with typed input and output. - - Example: - ```python - from genkit import Input, Output - from pydantic import BaseModel - - - class RecipeInput(BaseModel): - dish: str - - - class Recipe(BaseModel): - name: str - ingredients: list[str] - - - # With typed input AND output - recipe_prompt = define_prompt( - registry, - name='recipe', - prompt='Create a recipe for {dish}', - input=Input(schema=RecipeInput), - output=Output(schema=Recipe), - ) - - # Input is type-checked! - response = await recipe_prompt(RecipeInput(dish='pizza')) - response.output.name # βœ“ Typed as str - ``` - """ - # If Input[T] is provided, extract its schema - effective_input_schema = input_schema - if input is not None: - effective_input_schema = input.schema - - # If Output[T] is provided, extract its configuration - effective_output_schema = output_schema - effective_output_format = output_format - effective_output_content_type = output_content_type - effective_output_instructions = output_instructions - effective_output_constrained = output_constrained - - if output is not None: - effective_output_schema = output.schema - effective_output_format = output.format if output.format else output_format - if output.content_type is not None: - effective_output_content_type = output.content_type - if output.instructions is not None: - effective_output_instructions = output.instructions - if output.constrained is not None: - effective_output_constrained = output.constrained - - executable_prompt: ExecutablePrompt[Any, Any] = ExecutablePrompt( - registry, - variant=variant, - model=model, - config=config, - description=description, - input_schema=effective_input_schema, - system=system, - prompt=prompt, - messages=messages, - output_format=effective_output_format, - output_content_type=effective_output_content_type, - output_instructions=effective_output_instructions, - output_schema=effective_output_schema, - output_constrained=effective_output_constrained, - max_turns=max_turns, - return_tool_requests=return_tool_requests, - metadata=metadata, - tools=tools, - tool_choice=tool_choice, - use=use, - docs=docs, - _name=name, - ) - - if name: - # Register actions for this prompt - action_metadata: dict[str, object] = { - 'type': 'prompt', - 'source': 'programmatic', - 'prompt': { - 'name': name, - 'variant': variant or '', - }, - } - - async def prompt_action_fn(input: Any = None) -> GenerateRequest: # noqa: ANN401 - """PROMPT action function - renders prompt and returns GenerateRequest.""" - options = await executable_prompt.render(input=input) - return await to_generate_request(registry, options) - - async def executable_prompt_action_fn(input: Any = None) -> GenerateActionOptions: # noqa: ANN401 - """EXECUTABLE_PROMPT action function - renders prompt and returns GenerateActionOptions.""" - return await executable_prompt.render(input=input) - - action_name = registry_definition_key(name, variant) - prompt_action = registry.register_action( - kind=cast(ActionKind, ActionKind.PROMPT), - name=action_name, - fn=prompt_action_fn, - metadata=action_metadata, - ) - - executable_prompt_action = registry.register_action( - kind=cast(ActionKind, ActionKind.EXECUTABLE_PROMPT), - name=action_name, - fn=executable_prompt_action_fn, - metadata=action_metadata, - ) - - # Link them - executable_prompt._prompt_action = prompt_action # pyright: ignore[reportPrivateUsage] - # Dynamic attributes set at runtime - these are custom attrs added to Action objects - setattr(prompt_action, '_executable_prompt', weakref.ref(executable_prompt)) # noqa: B010 - setattr(executable_prompt_action, '_executable_prompt', weakref.ref(executable_prompt)) # noqa: B010 - - return executable_prompt - - -def _resolve_output_schema( - registry: Registry, - output_schema: type | dict[str, Any] | str | None, - output: GenerateActionOutputConfig, -) -> None: - """Resolve output schema and populate the output config. - - Handles three types of output_schema: - - str: Schema name - look up JSON schema and type from registry - - Pydantic type: Store both JSON schema and type for runtime validation - - dict: Raw JSON schema - convert directly - - Args: - registry: The registry to use for schema lookups. - output_schema: The schema to resolve (string name, Pydantic type, or dict). - output: The output config to populate with json_schema and schema_type. - """ - if output_schema is None: - return - - if isinstance(output_schema, str): - # Schema name - look up from registry - resolved_schema = registry.lookup_schema(output_schema) - if resolved_schema: - output.json_schema = resolved_schema - # Also look up the schema type for runtime validation - schema_type = registry.lookup_schema_type(output_schema) - if schema_type: - output.schema_type = schema_type - elif isinstance(output_schema, type) and issubclass(output_schema, BaseModel): - # Pydantic type - store both JSON schema and type - output.json_schema = to_json_schema(output_schema) - output.schema_type = output_schema - else: - # dict (raw JSON schema) - output.json_schema = to_json_schema(output_schema) - - -async def to_generate_action_options(registry: Registry, options: PromptConfig) -> GenerateActionOptions: - """Converts the given parameters to a GenerateActionOptions object. - - Args: - registry: The registry to use for resolving models and tools. - options: The prompt configuration. - - Returns: - A GenerateActionOptions object. - """ - model = options.model or registry.default_model - if model is None: - raise GenkitError(status='INVALID_ARGUMENT', message='No model configured.') - - cache = PromptCache() - resolved_msgs: list[Message] = [] - # Use empty dict instead of None for render functions - render_input: dict[str, Any] = {} - if options.system: - result = await render_system_prompt(registry, render_input, options, cache) - resolved_msgs.append(result) - if options.messages: - resolved_msgs.extend(await render_message_prompt(registry, render_input, options, cache)) - if options.prompt: - result = await render_user_prompt(registry, render_input, options, cache) - resolved_msgs.append(result) - - # If is schema is set but format is not explicitly set, default to - # `json` format. - output_format = 'json' if options.output_schema and not options.output_format else options.output_format - - output = GenerateActionOutputConfig() - if output_format: - output.format = output_format - if options.output_content_type: - output.content_type = options.output_content_type - if options.output_instructions is not None: - output.instructions = options.output_instructions - _resolve_output_schema(registry, options.output_schema, output) - if options.output_constrained is not None: - output.constrained = options.output_constrained - - resume = None - if options.tool_responses: - # Filter for only ToolResponsePart instances - tool_response_parts = [r.root for r in options.tool_responses if isinstance(r.root, ToolResponsePart)] - if tool_response_parts: - resume = Resume(respond=tool_response_parts) - - return GenerateActionOptions( - model=model, - messages=resolved_msgs, - config=options.config, - tools=options.tools, - return_tool_requests=options.return_tool_requests, - tool_choice=options.tool_choice, - output=output, - max_turns=options.max_turns, - docs=await render_docs(render_input, options), - resume=resume, - ) - - -async def to_generate_request(registry: Registry, options: GenerateActionOptions) -> GenerateRequest: - """Converts GenerateActionOptions to a GenerateRequest. - - This function resolves tool names into their respective tool definitions - by looking them up in the provided registry. it also validates that the - provided options contain at least one message. - - Args: - registry: The Registry instance used to look up tool actions. - options: The GenerateActionOptions containing the configuration, - messages, and tool references to be converted. - - Returns: - A GenerateRequest object populated with messages, config, resolved - tools, and output configurations. - - Raises: - Exception: If a tool name provided in options cannot be found in - the registry. - GenkitError: If the options do not contain any messages. - """ - tools: list[Action] = [] - if options.tools: - for tool_name in options.tools: - tool_action = await registry.resolve_action(cast(ActionKind, ActionKind.TOOL), tool_name) - if tool_action is None: - raise GenkitError(status='NOT_FOUND', message=f'Unable to resolve tool {tool_name}') - tools.append(tool_action) - - tool_defs = [to_tool_definition(tool) for tool in tools] if tools else [] - - if not options.messages: - raise GenkitError( - status='INVALID_ARGUMENT', - message='at least one message is required in generate request', - ) - - return GenerateRequest( - messages=options.messages, - config=options.config if options.config is not None else {}, - docs=options.docs, - tools=tool_defs, - tool_choice=options.tool_choice, - output=OutputConfig( - content_type=options.output.content_type if options.output else None, - format=options.output.format if options.output else None, - schema=options.output.json_schema if options.output else None, - constrained=options.output.constrained if options.output else None, - ), - ) - - -def _normalize_prompt_arg( - prompt: str | Part | list[Part] | None, -) -> list[Part]: - """Normalizes different prompt input types into a list of Parts. - - Ensures that the prompt content, whether provided as a string, a single Part, - or a list of Parts, is consistently represented as a list of Part objects. - - Args: - prompt: The prompt input, which can be a string, a Part, a list of Parts, - or None. - - Returns: - A list containing the normalized Part(s). Returns empty list if input is None - or empty. - """ - if not prompt: - return [] - if isinstance(prompt, str): - # Part is a RootModel, so we pass content via 'root' parameter - return [Part(root=TextPart(text=prompt))] - elif isinstance(prompt, list): - return prompt - elif isinstance(prompt, Part): # pyright: ignore[reportUnnecessaryIsInstance] - return [prompt] - else: - return [] # pyright: ignore[reportUnreachable] - defensive fallback - - -async def render_system_prompt( - registry: Registry, - input: dict[str, Any], - options: PromptConfig, - prompt_cache: PromptCache, - context: dict[str, Any] | None = None, -) -> Message: - """Renders the system prompt for a prompt action. - - This function handles rendering system prompts by either: - 1. Processing dotprompt templates if the system prompt is a string - 2. Normalizing the system prompt into a list of Parts if it's a Part or list of Parts - - Args: - registry: Registry instance for resolving models and tools - input: Dictionary of input values for template rendering - options: Configuration options for the prompt - prompt_cache: Cache for compiled prompt templates - context: Optional dictionary of context values for template rendering - - Returns: - Message: A Message object containing the rendered system prompt with Role.SYSTEM - - """ - if isinstance(options.system, str): - if prompt_cache.system is None: - prompt_cache.system = await registry.dotprompt.compile(options.system) - - if options.metadata: - context = {**(context or {}), 'state': options.metadata.get('state')} - - # Cast to list[Part] - Pydantic coerces dicts to Part objects at runtime - rendered_parts = cast( - list[Part], - await render_dotprompt_to_parts( - context or {}, - prompt_cache.system, - input, - PromptMetadata( - input=PromptInputConfig( - schema=to_json_schema(options.input_schema) if options.input_schema else None, - ) - ), - ), - ) - return Message(role=Role.SYSTEM, content=rendered_parts) - - if callable(options.system): - resolved = await ensure_async(options.system)(input, context) - return Message(role=Role.SYSTEM, content=_normalize_prompt_arg(resolved)) - - return Message(role=Role.SYSTEM, content=_normalize_prompt_arg(options.system)) - - -async def render_dotprompt_to_parts( - context: dict[str, Any], - prompt_function: PromptFunction[Any], - input_: dict[str, Any], - options: PromptMetadata[Any] | None = None, -) -> list[dict[str, Any]]: - """Renders a prompt template into a list of content parts using dotprompt. - - Args: - context: Dictionary containing context variables available to the prompt template. - prompt_function: The compiled dotprompt function to execute. - input_: Dictionary containing input variables for the prompt template. - options: Optional prompt metadata configuration. - - Returns: - A list of dictionaries representing Part objects for Pydantic re-validation. - - Raises: - Exception: If the template produces more than one message. - """ - # Flatten input and context for template resolution - flattened_data = {**(context or {}), **(input_ or {})} - rendered = await prompt_function( - data=DataArgument[dict[str, Any]]( - input=flattened_data, - context=context, - ), - options=options, - ) - - if len(rendered.messages) > 1: - raise Exception('parts template must produce only one message') - - # Convert parts to dicts for Pydantic re-validation when creating new Message - part_rendered: list[dict[str, Any]] = [] - for message in rendered.messages: - for part in message.content: - part_rendered.append(part.model_dump()) - - return part_rendered - - -async def render_message_prompt( - registry: Registry, - input: dict[str, Any], - options: PromptConfig, - prompt_cache: PromptCache, - context: dict[str, Any] | None = None, - history: list[Message] | None = None, -) -> list[Message]: - """Render a message prompt using a given registry, input data, options, and a context. - - This function processes different types of message options (string or list) to render - appropriate messages using a prompt registry and cache. If the `messages` option is of type - string, the function compiles the dotprompt messages from the `registry` and applies data - and metadata context. If the `messages` option is of type list, it either validates and - returns the list or processes it for message rendering. The function ensures correct message - output using the provided input, prompt configuration, and caching mechanism. - - Arguments: - registry (Registry): The registry used to compile dotprompt messages. - input (dict[str, Any]): The input data to render messages. - options (PromptConfig): Configuration containing prompt options and message settings. - prompt_cache (PromptCache): Cache to store compiled prompt results. - context (dict[str, Any] | None): Optional additional context to be used for rendering. - Defaults to None. - history (list[Message] | None): Optional conversation history to be passed to message - resolvers/templates. Matches JS behavior where opts.messages is passed as history. - - Returns: - list[Message]: A list of rendered or validated message objects. - """ - if isinstance(options.messages, str): - if prompt_cache.messages is None: - prompt_cache.messages = await registry.dotprompt.compile(options.messages) - - if options.metadata: - context = {**(context or {}), 'state': options.metadata.get('state')} - - # Convert history to dict format for template - messages_ = None - if history: - messages_ = [e.model_dump() for e in history] - - # Flatten input and context for template resolution - flattened_data = {**(context or {}), **(input or {})} - rendered = await prompt_cache.messages( - data=DataArgument[dict[str, Any]]( - input=flattened_data, - context=context, - messages=messages_, # pyright: ignore[reportArgumentType] - ), - options=PromptMetadata( - input=PromptInputConfig( - schema=to_json_schema(options.input_schema) if options.input_schema else None, - ) - ), - ) - return [Message.model_validate(e.model_dump()) for e in rendered.messages] - - elif isinstance(options.messages, list): - return [m if isinstance(m, Message) else Message.model_validate(m) for m in options.messages] - - elif callable(options.messages): - # Pass history to resolver function (matching JS MessagesResolver signature) - resolved = await ensure_async(options.messages)(input, {'context': context, 'history': history}) - return list(resolved) if resolved else [] - - raise TypeError(f'Unsupported type for messages: {type(options.messages)}') - - -async def render_user_prompt( - registry: Registry, - input: dict[str, Any], - options: PromptConfig, - prompt_cache: PromptCache, - context: dict[str, Any] | None = None, -) -> Message: - """Asynchronously renders a user prompt based on the given input, context, and options. - - Utilizes a pre-compiled or dynamically compiled dotprompt template. - - Arguments: - registry: The registry instance used to compile dotprompt templates. - input: The input data used to populate the prompt. - options: The configuration for rendering the prompt, including - the template type and associated metadata. - prompt_cache: A cache that stores pre-compiled prompt templates to - optimize rendering. - context: Optional dynamic context data to override or - supplement in the rendering process. - - Returns: - Message: A Message instance containing the rendered user prompt. - """ - if isinstance(options.prompt, str): - if prompt_cache.user_prompt is None: - prompt_cache.user_prompt = await registry.dotprompt.compile(options.prompt) - - if options.metadata: - context = {**(context or {}), 'state': options.metadata.get('state')} - - # Cast to list[Part] - Pydantic coerces dicts to Part objects at runtime - rendered_parts = cast( - list[Part], - await render_dotprompt_to_parts( - context or {}, - prompt_cache.user_prompt, - input, - PromptMetadata( - input=PromptInputConfig( - schema=to_json_schema(options.input_schema) if options.input_schema else None, - ) - ), - ), - ) - return Message(role=Role.USER, content=rendered_parts) - - if callable(options.prompt): - resolved = await ensure_async(options.prompt)(input, context) - return Message(role=Role.USER, content=_normalize_prompt_arg(resolved)) - - return Message(role=Role.USER, content=_normalize_prompt_arg(options.prompt)) - - -async def render_docs( - input: dict[str, Any], - options: PromptConfig, - context: dict[str, Any] | None = None, -) -> list[DocumentData] | None: - """Renders the docs for a prompt action. - - Args: - input: Dictionary of input values. - options: Configuration options for the prompt. - context: Optional dictionary of context values. - - Returns: - A list of DocumentData objects or None. - """ - if options.docs is None: - return None - - if callable(options.docs): - return await ensure_async(options.docs)(input, context) - - return options.docs - - -def registry_definition_key(name: str, variant: str | None = None, ns: str | None = None) -> str: - """Generate a registry definition key for a prompt. - - Format: "ns/name.variant" where ns and variant are optional. - - Args: - name: The prompt name. - variant: Optional variant name. - ns: Optional namespace. - - Returns: - Registry key string. - """ - parts = [] - if ns: - parts.append(ns) - parts.append(name) - if variant: - parts[-1] = f'{parts[-1]}.{variant}' - return '/'.join(parts) - - -def registry_lookup_key(name: str, variant: str | None = None, ns: str | None = None) -> str: - """Generate a registry lookup key for a prompt. - - Args: - name: The prompt name. - variant: Optional variant name. - ns: Optional namespace. - - Returns: - Registry lookup key string. - """ - return f'/prompt/{registry_definition_key(name, variant, ns)}' - - -def define_partial(registry: Registry, name: str, source: str) -> None: - """Define a partial template in the registry. - - Partials are reusable template fragments that can be included in other prompts. - Files starting with `_` are treated as partials. - - Args: - registry: The registry to register the partial in. - name: The name of the partial. - source: The template source code. - """ - _ = registry.dotprompt.define_partial(name, source) - logger.debug(f'Registered Dotprompt partial "{name}"') - - -def define_helper(registry: Registry, name: str, fn: Callable[..., Any]) -> None: - """Define a Handlebars helper function in the registry. - - Args: - registry: The registry to register the helper in. - name: The name of the helper function. - fn: The helper function to register. - """ - _ = registry.dotprompt.define_helper(name, fn) - logger.debug(f'Registered Dotprompt helper "{name}"') - - -def define_schema(registry: Registry, name: str, schema: type[BaseModel]) -> None: - """Register a Pydantic schema for use in prompts. - - Schemas registered with this function can be referenced by name in - .prompt files using the `output.schema` field. - - Args: - registry: The registry to register the schema in. - name: The name of the schema. - schema: The Pydantic model class to register. - - Example: - ```python - from genkit.blocks.prompt import define_schema - - define_schema(registry, 'Recipe', Recipe) - ``` - - Then in a .prompt file: - ```yaml - output: - schema: Recipe - ``` - """ - json_schema = to_json_schema(schema) - registry.register_schema(name, json_schema, schema_type=schema) - logger.debug(f'Registered schema "{name}"') - - -def load_prompt(registry: Registry, path: Path, filename: str, prefix: str = '', ns: str = '') -> None: - """Load a single prompt file and register it in the registry. - - This function loads a .prompt file, parses it, and registers it as a lazy-loaded - prompt that will only be fully loaded when first accessed. - - Args: - registry: The registry to register the prompt in. - path: Base path to the prompts directory. - filename: Name of the prompt file (e.g., "myPrompt.prompt" or "myPrompt.variant.prompt"). - prefix: Subdirectory prefix (for namespacing). - ns: Namespace for the prompt. - """ - # Extract name and variant from filename - # "myPrompt.prompt" -> name="myPrompt", variant=None - # "myPrompt.variant.prompt" -> name="myPrompt", variant="variant" - # "subdir/myPrompt.prompt" -> name="subdir/myPrompt", variant=None - if not filename.endswith('.prompt'): - raise ValueError(f"Invalid prompt filename: {filename}. Must end with '.prompt'") - - base_name = filename.removesuffix('.prompt') - - name = f'{prefix}{base_name}' if prefix else base_name - variant: str | None = None - - # Extract variant (only takes parts[1], not all remaining parts) - if '.' in name: - parts = name.split('.') - name = parts[0] - variant = parts[1] # Only first part after split - - # Build full file path - # prefix may have trailing slash, so we need to handle it - if prefix: - # Strip trailing slash for path construction (pathlib handles it) - prefix_clean = prefix.rstrip('/') - file_path = path / prefix_clean / filename - else: - file_path = path / filename - - # Read the prompt file - with Path(file_path).open(encoding='utf-8') as f: - source = f.read() - - # Parse the prompt - parsed_prompt = registry.dotprompt.parse(source) - - # Generate registry key - registry_key = registry_definition_key(name, variant, ns) - - # Create a lazy-loaded prompt definition - # The prompt will only be fully loaded when first accessed - async def load_prompt_metadata() -> dict[str, Any]: - """Lazy loader for prompt metadata.""" - prompt_metadata = await registry.dotprompt.render_metadata(parsed_prompt) - - # Convert Pydantic model to dict if needed - prompt_metadata_dict: dict[str, Any] - if hasattr(prompt_metadata, 'model_dump'): - prompt_metadata_dict = prompt_metadata.model_dump(by_alias=True) - elif hasattr(prompt_metadata, 'dict'): - # Fallback for older Pydantic versions - # pyrefly: ignore[deprecated] - Intentional for Pydantic v1 compatibility - prompt_metadata_dict = prompt_metadata.dict(by_alias=True) # pyright: ignore[reportDeprecated] - else: - # Already a dict - cast through object to satisfy type checker - prompt_metadata_dict = cast(dict[str, Any], cast(object, prompt_metadata)) - - # Ensure raw metadata is available (critical for lazy schema resolution) - if hasattr(prompt_metadata, 'raw'): - prompt_metadata_dict['raw'] = prompt_metadata.raw - - if variant: - prompt_metadata_dict['variant'] = variant - - # Fallback for model if not present (Dotprompt issue) - if not prompt_metadata_dict.get('model'): - raw_model = (prompt_metadata_dict.get('raw') or {}).get('model') - if raw_model: - prompt_metadata_dict['model'] = raw_model - - # Clean up null descriptions - output = prompt_metadata_dict.get('output') - schema = None - if output and isinstance(output, dict): - schema = output.get('schema') - if schema and isinstance(schema, dict) and schema.get('description') is None: - schema.pop('description', None) - - if not schema: - # Fallback to raw schema name if schema definition is missing - raw_schema = (prompt_metadata_dict.get('raw') or {}).get('output', {}).get('schema') - if isinstance(raw_schema, str): - schema = raw_schema - # output might be None if it wasn't in parsed config - if not output: - output = {'schema': schema} - prompt_metadata_dict['output'] = output - elif isinstance(output, dict): - output['schema'] = schema - - input_schema = prompt_metadata_dict.get('input') - if input_schema and isinstance(input_schema, dict): - schema = input_schema.get('schema') - if schema and isinstance(schema, dict) and schema.get('description') is None: - schema.pop('description', None) - - # Build metadata structure (prompt_metadata_dict is always dict[str, Any] at this point) - metadata_inner = prompt_metadata_dict.get('metadata', {}) - base_dict = prompt_metadata_dict - metadata = { - **base_dict, - **(metadata_inner if isinstance(metadata_inner, dict) else {}), - 'type': 'prompt', - 'prompt': { - **base_dict, - 'template': parsed_prompt.template, - }, - } - - raw = prompt_metadata_dict.get('raw') - if raw and isinstance(raw, dict) and raw.get('metadata'): - metadata['metadata'] = {**raw['metadata']} - - output = prompt_metadata_dict.get('output') - input_schema = prompt_metadata_dict.get('input') - raw = prompt_metadata_dict.get('raw') - - return { - 'name': registry_key, - 'model': prompt_metadata_dict.get('model'), - 'config': prompt_metadata_dict.get('config'), - 'tools': prompt_metadata_dict.get('tools'), - 'description': prompt_metadata_dict.get('description'), - 'output': { - 'jsonSchema': output.get('schema') if isinstance(output, dict) else None, - 'format': output.get('format') if isinstance(output, dict) else None, - }, - 'input': { - 'default': input_schema.get('default') if isinstance(input_schema, dict) else None, - 'jsonSchema': input_schema.get('schema') if isinstance(input_schema, dict) else None, - }, - 'metadata': metadata, - 'maxTurns': raw.get('maxTurns') if isinstance(raw, dict) else None, - 'toolChoice': raw.get('toolChoice') if isinstance(raw, dict) else None, - 'returnToolRequests': raw.get('returnToolRequests') if isinstance(raw, dict) else None, - 'messages': parsed_prompt.template, - } - - # Cache the ExecutablePrompt to avoid duplicate creation when both PROMPT and - # EXECUTABLE_PROMPT actions trigger lazy loading from the reflection API. - _cached_prompt: ExecutablePrompt[Any, Any] | None = None - - # Create a factory function that will create the ExecutablePrompt when accessed - # Store metadata in a closure to avoid global state - async def create_prompt_from_file() -> ExecutablePrompt[Any, Any]: - """Factory function to create ExecutablePrompt from file metadata (memoized).""" - nonlocal _cached_prompt - if _cached_prompt is not None: - return _cached_prompt - - metadata = await load_prompt_metadata() - - # Create ExecutablePrompt from metadata - # Pass name/ns so as_tool() can look up the action - executable_prompt = ExecutablePrompt( - registry=registry, - variant=metadata.get('variant'), - model=metadata.get('model'), - config=metadata.get('config'), - description=metadata.get('description'), - input_schema=metadata.get('input', {}).get('jsonSchema'), - output_schema=metadata.get('output', {}).get('jsonSchema'), - output_constrained=True if metadata.get('output', {}).get('jsonSchema') else None, - output_format=metadata.get('output', {}).get('format'), - messages=metadata.get('messages'), - max_turns=metadata.get('maxTurns'), - tool_choice=metadata.get('toolChoice'), - return_tool_requests=metadata.get('returnToolRequests'), - metadata=metadata.get('metadata'), - tools=metadata.get('tools'), - _name=name, # Store name for action lookup - _ns=ns, # Store namespace for action lookup - ) - - # Store reference to PROMPT action on the ExecutablePrompt - # Actions are already registered at this point (lazy loading happens after registration) - definition_key = registry_definition_key(name, variant, ns) - prompt_lookup_key = create_action_key(ActionKind.PROMPT, definition_key) - exec_prompt_lookup_key = create_action_key(ActionKind.EXECUTABLE_PROMPT, definition_key) - - # Update PROMPT action - prompt_action = await registry.resolve_action_by_key(prompt_lookup_key) - if prompt_action and prompt_action.kind == ActionKind.PROMPT: - executable_prompt._prompt_action = prompt_action # pyright: ignore[reportPrivateUsage] - setattr(prompt_action, '_executable_prompt', weakref.ref(executable_prompt)) # noqa: B010 - - # Update EXECUTABLE_PROMPT action - exec_prompt_action = await registry.resolve_action_by_key(exec_prompt_lookup_key) - - # Update schemas on BOTH actions for Dev UI reflection - input_json_schema = metadata.get('input', {}).get('jsonSchema') - output_json_schema = metadata.get('output', {}).get('jsonSchema') - - for action in [prompt_action, exec_prompt_action]: - if action is not None: - if input_json_schema: - action.input_schema = input_json_schema - if output_json_schema: - action.output_schema = output_json_schema - - # Cache the result for subsequent calls - _cached_prompt = executable_prompt - return executable_prompt - - # Store the async factory in a way that can be accessed later - # We'll store it in the action metadata - action_metadata: dict[str, object] = { - 'type': 'prompt', - 'lazy': True, - 'source': 'file', - 'prompt': { - 'name': name, - 'variant': variant or '', - }, - } - - # Create two separate action functions : - # 1. PROMPT action - returns GenerateRequest (for rendering prompts) - # 2. EXECUTABLE_PROMPT action - returns GenerateActionOptions (for executing prompts) - - async def prompt_action_fn(input: Any = None) -> GenerateRequest: # noqa: ANN401 - """PROMPT action function - renders prompt and returns GenerateRequest.""" - # Load the prompt (lazy loading) - prompt = await create_prompt_from_file() - - # Render the prompt with input to get GenerateActionOptions - options = await prompt.render(input=input) - - # Convert GenerateActionOptions to GenerateRequest - return await to_generate_request(registry, options) - - async def executable_prompt_action_fn(input: Any = None) -> GenerateActionOptions: # noqa: ANN401 - """EXECUTABLE_PROMPT action function - renders prompt and returns GenerateActionOptions.""" - # Load the prompt (lazy loading) - prompt = await create_prompt_from_file() - - # Render the prompt with input to get GenerateActionOptions - return await prompt.render(input=input) - - # Register the PROMPT action - # Use registry_definition_key for the action name (not registry_lookup_key) - # The action name should be just the definition key (e.g., "dotprompt/testPrompt"), - # not the full lookup key (e.g., "/prompt/dotprompt/testPrompt") - action_name = registry_definition_key(name, variant, ns) - prompt_action = registry.register_action( - kind=cast(ActionKind, ActionKind.PROMPT), - name=action_name, - fn=prompt_action_fn, - metadata=action_metadata, - ) - - # Register the EXECUTABLE_PROMPT action - executable_prompt_action = registry.register_action( - kind=cast(ActionKind, ActionKind.EXECUTABLE_PROMPT), - name=action_name, - fn=executable_prompt_action_fn, - metadata=action_metadata, - ) - - # Store the factory function on both actions for easy access - setattr(prompt_action, '_async_factory', create_prompt_from_file) # noqa: B010 - setattr(executable_prompt_action, '_async_factory', create_prompt_from_file) # noqa: B010 - - # Store ExecutablePrompt reference on actions - # This will be set when the prompt is first accessed (lazy loading) - # We'll update it in create_prompt_from_file after the prompt is created - - logger.debug(f'Registered prompt "{registry_key}" from "{file_path}"') - - -def load_prompt_folder_recursively(registry: Registry, dir_path: Path, ns: str, sub_dir: str = '') -> None: - """Recursively load all prompt files from a directory. - - Args: - registry: The registry to register prompts in. - dir_path: Base path to the prompts directory. - ns: Namespace for prompts. - sub_dir: Current subdirectory being processed (for recursion). - """ - full_path = dir_path / sub_dir if sub_dir else dir_path - - if not full_path.exists() or not full_path.is_dir(): - return - - # Iterate through directory entries - try: - for entry in os.scandir(full_path): - if entry.is_file() and entry.name.endswith('.prompt'): - if entry.name.startswith('_'): - # This is a partial - partial_name = entry.name[1:-7] # Remove "_" prefix and ".prompt" suffix - with Path(entry.path).open(encoding='utf-8') as f: - source = f.read() - - # Strip frontmatter if present - if source.startswith('---'): - end_frontmatter = source.find('---', 3) - if end_frontmatter != -1: - source = source[end_frontmatter + 3 :].strip() - - define_partial(registry, partial_name, source) - logger.debug(f'Registered Dotprompt partial "{partial_name}" from "{entry.path}"') - else: - # This is a regular prompt - prefix_with_slash = f'{sub_dir}/' if sub_dir else '' - load_prompt(registry, dir_path, entry.name, prefix_with_slash, ns) - elif entry.is_dir(): - # Recursively process subdirectories - new_sub_dir = os.path.join(sub_dir, entry.name) if sub_dir else entry.name - load_prompt_folder_recursively(registry, dir_path, ns, new_sub_dir) - except PermissionError: - logger.warning(f'Permission denied accessing directory: {full_path}') - except Exception as e: - logger.exception(f'Error loading prompts from {full_path}', exc_info=e) - - -def load_prompt_folder(registry: Registry, dir_path: str | Path = './prompts', ns: str = '') -> None: - """Load all prompt files from a directory. - - This is the main entry point for loading prompts from a directory. - It recursively processes all `.prompt` files and registers them. - - Args: - registry: The registry to register prompts in. - dir_path: Path to the prompts directory. Defaults to './prompts'. - ns: Namespace for prompts. Defaults to 'dotprompt'. - """ - path = Path(dir_path).resolve() - - if not path.exists(): - logger.warning(f'Prompt directory does not exist: {path}') - return - - if not path.is_dir(): - logger.warning(f'Prompt path is not a directory: {path}') - return - - load_prompt_folder_recursively(registry, path, ns, '') - logger.info(f'Loaded prompts from directory: {path}') - - -async def lookup_prompt(registry: Registry, name: str, variant: str | None = None) -> ExecutablePrompt[Any, Any]: - """Look up a prompt from the registry. - - Args: - registry: The registry to look up the prompt from. - name: The name of the prompt. - variant: Optional variant name. - - Returns: - An ExecutablePrompt instance. - - Raises: - GenkitError: If the prompt is not found. - """ - # Try without namespace first (for programmatic prompts) - # Use create_action_key to build the full key: "/prompt/" - definition_key = registry_definition_key(name, variant, None) - lookup_key = create_action_key(cast(ActionKind, ActionKind.PROMPT), definition_key) - action = await registry.resolve_action_by_key(lookup_key) - - # If not found and no namespace was specified, try with default 'dotprompt' namespace - # (for file-based prompts) - if not action: - definition_key = registry_definition_key(name, variant, 'dotprompt') - lookup_key = create_action_key(cast(ActionKind, ActionKind.PROMPT), definition_key) - action = await registry.resolve_action_by_key(lookup_key) - - if action: - # First check if we've stored the ExecutablePrompt directly - prompt_ref = getattr(action, '_executable_prompt', None) - if prompt_ref is not None: - if isinstance(prompt_ref, weakref.ReferenceType): - resolved = prompt_ref() - if resolved is not None: - return resolved - if isinstance(prompt_ref, ExecutablePrompt): - return prompt_ref - # Otherwise, create it from the factory (lazy loading) - async_factory = getattr(action, '_async_factory', None) - if callable(async_factory): - # Cast to async callable - getattr returns object but we've verified it's callable - async_factory_fn = cast(Callable[[], Awaitable[ExecutablePrompt]], async_factory) - executable_prompt = await async_factory_fn() - if getattr(action, '_executable_prompt', None) is None: - setattr(action, '_executable_prompt', executable_prompt) # noqa: B010 - return executable_prompt - # Fallback: try to get from metadata - factory = action.metadata.get('_async_factory') - if callable(factory): - factory_async = ensure_async(cast(Callable[..., Any], factory)) - executable_prompt = await factory_async() - if getattr(action, '_executable_prompt', None) is None: - setattr(action, '_executable_prompt', executable_prompt) # noqa: B010 - return executable_prompt - # Last resort: this shouldn't happen if prompts are loaded correctly - raise GenkitError( - status='INTERNAL', - message=f'Prompt action found but no ExecutablePrompt available for {name}', - ) - - variant_str = f' (variant {variant})' if variant else '' - raise GenkitError( - status='NOT_FOUND', - message=f'Prompt {name}{variant_str} not found', - ) - - -async def prompt( - registry: Registry, - name: str, - variant: str | None = None, - _dir: str | Path | None = None, # Accepted but not used -) -> ExecutablePrompt[Any, Any]: - """Look up a prompt by name and optional variant. - - Can look up prompts that were: - 1. Defined programmatically using define_prompt() - 2. Loaded from .prompt files using load_prompt_folder() - - Args: - registry: The registry to look up the prompt from. - name: The name of the prompt. - variant: Optional variant name. - dir: Optional directory parameter (accepted for compatibility but not used). - - Returns: - An ExecutablePrompt instance. - - Raises: - GenkitError: If the prompt is not found. - """ - return await lookup_prompt(registry, name, variant) diff --git a/py/packages/genkit/src/genkit/blocks/reranker.py b/py/packages/genkit/src/genkit/blocks/reranker.py deleted file mode 100644 index a5eecd2854..0000000000 --- a/py/packages/genkit/src/genkit/blocks/reranker.py +++ /dev/null @@ -1,440 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Reranker type definitions for the Genkit framework. - -Rerankers and Two-Stage Retrieval -================================= - -A **reranking model** (also known as a cross-encoder) is a type of model that, -given a query and document, outputs a similarity score. This score is used to -reorder documents by relevance to the query. - -Reranker APIs take a list of documents (e.g., the output of a retriever) and -reorder them based on their relevance to the query. This step can be useful -for fine-tuning results and ensuring the most pertinent information is used -in the prompt provided to a generative model. - -Two-Stage Retrieval -------------------- - -In a typical RAG (Retrieval-Augmented Generation) pipeline: - -1. **Stage 1 - Retrieval**: A retriever fetches a large set of candidate - documents using fast vector similarity search. -2. **Stage 2 - Reranking**: A reranker scores and reorders these candidates - using more expensive but accurate cross-encoder models. - -This two-stage approach balances speed and accuracy: -- Retrievers are fast but may not perfectly rank results -- Rerankers are slower but provide superior relevance scoring - -Usage Example -------------- - -Using an existing reranker (e.g., Vertex AI): - -.. code-block:: python - - from genkit.ai import Genkit - - ai = Genkit(plugins=[...]) - - - @ai.flow() - async def rerank_flow(query: str): - documents = [ - Document.from_text('pythagorean theorem'), - Document.from_text('quantum mechanics'), - Document.from_text('pizza'), - ] - - reranked = await ai.rerank( - reranker='vertexai/semantic-ranker-512', - query=query, - documents=documents, - ) - - return [{'text': doc.text(), 'score': doc.score} for doc in reranked] - -Custom Rerankers ----------------- - -You can define custom rerankers for specific use cases: - -.. code-block:: python - - from genkit.ai import Genkit - from genkit.core.typing import ( - RerankerResponse, - RankedDocumentData, - RankedDocumentMetadata, - ) - - ai = Genkit() - - - async def custom_reranker_fn(query, documents, options): - # Your custom reranking logic here - # Example: score by keyword overlap - query_words = set(query.text().lower().split()) - scored = [] - for doc in documents: - doc_words = set(doc.text().lower().split()) - overlap = len(query_words & doc_words) - score = overlap / max(len(query_words), 1) - scored.append((doc, score)) - - # Sort by score descending and take top k - k = options.get('k', 3) if options else 3 - scored.sort(key=lambda x: x[1], reverse=True) - top_k = scored[:k] - - return RerankerResponse( - documents=[ - RankedDocumentData(content=doc.content, metadata=RankedDocumentMetadata(score=score)) - for doc, score in top_k - ] - ) - - - ai.define_reranker('custom/keyword-reranker', custom_reranker_fn) - - - # Use it in a flow - @ai.flow() - async def search_flow(query: str): - docs = await ai.retrieve(retriever='my-retriever', query=query) - return await ai.rerank(reranker='custom/keyword-reranker', query=query, documents=docs, options={'k': 5}) -""" - -from collections.abc import Awaitable, Callable -from typing import Any, ClassVar, TypeVar, cast - -from pydantic import BaseModel, ConfigDict -from pydantic.alias_generators import to_camel - -from genkit.blocks.document import Document -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry -from genkit.core.schema import to_json_schema -from genkit.core.typing import ( - DocumentData, - DocumentPart, - RankedDocumentData, - RerankerRequest, - RerankerResponse, -) - -T = TypeVar('T') - -# Type alias for reranker function -RerankerFn = Callable[[Document, list[Document], T], Awaitable[RerankerResponse]] - - -class RankedDocument(Document): - """A document with a relevance score from reranking. - - This class extends Document to include a score property that represents - the document's relevance to a query as determined by a reranker. - """ - - def __init__( - self, - content: list[DocumentPart], - metadata: dict[str, Any] | None = None, - score: float | None = None, - ) -> None: - """Initializes a RankedDocument object. - - Args: - content: A list of DocumentPart objects representing the document's content. - metadata: An optional dictionary containing metadata about the document. - score: The relevance score from reranking. - """ - md = metadata.copy() if metadata else {} - if score is not None: - md['score'] = score - super().__init__(content=content, metadata=md) - - @property - def score(self) -> float | None: - """Returns the relevance score of the document. - - Returns: - The relevance score as a float, or None if not set. - """ - if self.metadata and 'score' in self.metadata: - return self.metadata['score'] - return None - - @staticmethod - def from_ranked_document_data(data: RankedDocumentData) -> 'RankedDocument': - """Constructs a RankedDocument from RankedDocumentData. - - Args: - data: The RankedDocumentData containing content, metadata with score. - - Returns: - A new RankedDocument instance. - """ - return RankedDocument( - content=data.content, - metadata=data.metadata.model_dump(), - score=data.metadata.score, - ) - - -class RerankerSupports(BaseModel): - """Reranker capability support.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - media: bool | None = None - - -class RerankerInfo(BaseModel): - """Information about a reranker's capabilities.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - label: str | None = None - supports: RerankerSupports | None = None - - -class RerankerOptions(BaseModel): - """Configuration options for a reranker.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, alias_generator=to_camel) - - config_schema: dict[str, Any] | None = None - label: str | None = None - supports: RerankerSupports | None = None - - -class RerankerRef(BaseModel): - """Reference to a reranker with configuration. - - Used to reference a reranker by name with optional configuration - and version information. - """ - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - name: str - config: Any | None = None - version: str | None = None - info: RerankerInfo | None = None - - -def reranker_action_metadata( - name: str, - options: RerankerOptions | None = None, -) -> ActionMetadata: - """Creates action metadata for a reranker. - - Args: - name: The name of the reranker. - options: Optional configuration options for the reranker. - - Returns: - An ActionMetadata instance for the reranker. - """ - options = options if options is not None else RerankerOptions() - reranker_metadata_dict: dict[str, Any] = {'reranker': {}} - - if options.label: - reranker_metadata_dict['reranker']['label'] = options.label - - if options.supports: - reranker_metadata_dict['reranker']['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True) - - reranker_metadata_dict['reranker']['customOptions'] = options.config_schema if options.config_schema else None - - return ActionMetadata( - kind=cast(ActionKind, ActionKind.RERANKER), - name=name, - input_json_schema=to_json_schema(RerankerRequest), - output_json_schema=to_json_schema(RerankerResponse), - metadata=reranker_metadata_dict, - ) - - -def create_reranker_ref( - name: str, - config: dict[str, Any] | None = None, - version: str | None = None, - info: RerankerInfo | None = None, -) -> RerankerRef: - """Creates a RerankerRef instance. - - Args: - name: The name of the reranker. - config: Optional configuration for the reranker. - version: Optional version string. - info: Optional RerankerInfo with capability information. - - Returns: - A RerankerRef instance. - """ - return RerankerRef(name=name, config=config, version=version, info=info) - - -def define_reranker( - registry: Registry, - name: str, - fn: RerankerFn[Any], - options: RerankerOptions | None = None, - description: str | None = None, -) -> Action: - """Defines and registers a reranker action. - - Creates a reranker action from the provided function and registers it - in the given registry. - - Args: - registry: The registry to register the reranker in. - name: The name of the reranker. - fn: The reranker function that implements the reranking logic. - options: Optional configuration options for the reranker. - description: Optional description for the reranker action. - - Returns: - The registered Action instance. - - Example: - >>> async def my_reranker(query, documents, options): - ... # Score and sort documents - ... scored = [(doc, score_doc(query, doc)) for doc in documents] - ... scored.sort(key=lambda x: x[1], reverse=True) - ... return RerankerResponse( - ... documents=[ - ... RankedDocumentData(content=doc.content, metadata=RankedDocumentMetadata(score=score)) - ... for doc, score in scored - ... ] - ... ) - >>> define_reranker(registry, 'my-reranker', my_reranker) - """ - metadata = reranker_action_metadata(name, options) - - async def wrapper( - request: RerankerRequest, - _ctx: Any, # noqa: ANN401 - ) -> RerankerResponse: - query_doc = Document.from_document_data(request.query) - documents = [Document.from_document_data(d) for d in request.documents] - return await fn(query_doc, documents, request.options) - - return registry.register_action( - kind=cast(ActionKind, ActionKind.RERANKER), - name=name, - fn=wrapper, - metadata=metadata.metadata, - span_metadata={'genkit:metadata:reranker:name': name}, - description=description, - ) - - -# Type for reranker argument (can be action, reference, or string name) -RerankerArgument = Action | RerankerRef | str - - -class RerankerParams(BaseModel): - """Parameters for the rerank function. - - Attributes: - reranker: The reranker to use (action, reference, or name string). - query: The query to rank documents against. - documents: The list of documents to rerank. - options: Optional configuration options for this rerank call. - """ - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, arbitrary_types_allowed=True) - - reranker: RerankerArgument - query: str | DocumentData - documents: list[DocumentData] - options: Any | None = None - - -async def rerank( - registry: Registry, - params: RerankerParams | dict[str, Any], -) -> list[RankedDocument]: - """Reranks documents based on the provided query using a reranker. - - This function takes a query and a list of documents, and returns the - documents reordered by relevance to the query as determined by the - specified reranker. - - Args: - registry: The registry to look up the reranker in. - params: Parameters for the rerank operation + including the reranker, - query, documents, and optional configuration. - - Returns: - A list of RankedDocument objects sorted by relevance. - - Raises: - ValueError: If the reranker cannot be resolved. - - Example: - >>> ranked_docs = await rerank( - ... registry, - ... { - ... 'reranker': 'my-reranker', - ... 'query': 'What is machine learning?', - ... 'documents': [doc1, doc2, doc3], - ... }, - ... ) - >>> for doc in ranked_docs: - ... print(f'Score: {doc.score}, Text: {doc.text()}') - """ - # Convert dict to RerankerParams if needed - if isinstance(params, dict): - params = RerankerParams(**params) - - # Resolve the reranker action - reranker_action = None - - if isinstance(params.reranker, str): - reranker_action = await registry.resolve_reranker(params.reranker) - elif isinstance(params.reranker, RerankerRef): - reranker_action = await registry.resolve_reranker(params.reranker.name) - elif isinstance(params.reranker, Action): # pyright: ignore[reportUnnecessaryIsInstance] - reranker_action = params.reranker - - if reranker_action is None: - raise ValueError(f'Unable to resolve reranker: {params.reranker}') - - # Convert query to DocumentData if it's a string - query_data: DocumentData - query_data = Document.from_text(params.query) if isinstance(params.query, str) else params.query - - # Build the request - request = RerankerRequest( - query=query_data, - documents=params.documents, - options=params.options, - ) - - # Call the reranker - action_response = await reranker_action.arun(request) - response: RerankerResponse = action_response.response - - # Convert response to RankedDocument list - return [RankedDocument.from_ranked_document_data(doc) for doc in response.documents] diff --git a/py/packages/genkit/src/genkit/blocks/retriever.py b/py/packages/genkit/src/genkit/blocks/retriever.py deleted file mode 100644 index 31877fd524..0000000000 --- a/py/packages/genkit/src/genkit/blocks/retriever.py +++ /dev/null @@ -1,359 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Retriever type definitions for the Genkit framework. - -This module defines the type interfaces for retrievers and indexers in the -Genkit framework. Retrievers are used for fetching documents from a datastore -given a query, enabling Retrieval-Augmented Generation (RAG) patterns. - -Overview: - Retrievers are a core component of RAG (Retrieval-Augmented Generation) - workflows. They search a document store and return relevant documents - that can be used to ground model responses with factual information. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ RAG Data Flow β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ User β”‚ ───► β”‚ Embedder β”‚ ───► β”‚ Retrieverβ”‚ ───► β”‚ Model β”‚ β”‚ - β”‚ β”‚ Query β”‚ β”‚ β”‚ β”‚ Search β”‚ β”‚ Generate β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Document β”‚ β”‚ - β”‚ β”‚ Store β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Terminology: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Term β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Retriever β”‚ Action that searches a document store and returns β”‚ - β”‚ β”‚ relevant documents based on a query. β”‚ - β”‚ Indexer β”‚ Action that adds documents to a document store, β”‚ - β”‚ β”‚ typically with embeddings for similarity search. β”‚ - β”‚ RetrieverRef β”‚ Reference to a retriever with optional config. β”‚ - β”‚ IndexerRef β”‚ Reference to an indexer with optional config. β”‚ - β”‚ Document β”‚ A structured piece of content with text/media and β”‚ - β”‚ β”‚ metadata. See genkit.blocks.document. β”‚ - β”‚ Query β”‚ The search query, typically as a Document object. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Components: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Component β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Retriever[T] β”‚ Base class for retriever implementations β”‚ - β”‚ RetrieverRequest β”‚ Input model with query and options β”‚ - β”‚ RetrieverOptions β”‚ Configuration for defining retrievers β”‚ - β”‚ RetrieverRef β”‚ Reference bundling name, config, version β”‚ - β”‚ IndexerRequest β”‚ Input model for indexing documents β”‚ - β”‚ IndexerOptions β”‚ Configuration for defining indexers β”‚ - β”‚ define_retriever() β”‚ Factory function for creating retriever actions β”‚ - β”‚ define_indexer() β”‚ Factory function for creating indexer actions β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Defining a simple retriever: - - ```python - from genkit import Genkit, Document - from genkit.blocks.retriever import RetrieverOptions - - ai = Genkit() - - - # Simple in-memory retriever - @ai.retriever(name='my_retriever') - async def my_retriever(query: Document, options: dict) -> list[Document]: - # Search logic here (e.g., vector similarity search) - results = search_documents(query.text(), top_k=options.get('k', 5)) - return [Document.from_text(r['text'], r['metadata']) for r in results] - - - # Use the retriever - docs = await ai.retrieve(retriever='my_retriever', query='What is Genkit?') - ``` - - Using simple_retriever for easier definition: - - ```python - @ai.simple_retriever(name='docs_retriever', configSchema=MyConfigSchema) - async def docs_retriever(query: str, options: MyConfigSchema) -> list[Document]: - # The query is automatically converted to string - return await search_docs(query, limit=options.limit) - ``` - -Caveats: - - Retriever functions receive a Document object, not a raw string - - Use simple_retriever() for a more convenient string-based query API - - Indexers are typically used during data ingestion, not query time - -See Also: - - genkit.blocks.document: Document model - - genkit.blocks.embedding: Embedder for generating document embeddings - - RAG documentation: https://genkit.dev/docs/rag -""" - -import inspect -from collections.abc import Awaitable, Callable -from typing import Any, ClassVar, Generic, TypeVar, cast - -from pydantic import BaseModel, ConfigDict -from pydantic.alias_generators import to_camel - -from genkit.blocks.document import Document -from genkit.core.action import ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry -from genkit.core.schema import to_json_schema -from genkit.core.typing import DocumentData, RetrieverResponse - -T = TypeVar('T') -# type RetrieverFn[T] = Callable[[Document, T], RetrieverResponse | Awaitable[RetrieverResponse]] -RetrieverFn = Callable[[Document, T], RetrieverResponse | Awaitable[RetrieverResponse]] - - -class Retriever(Generic[T]): - """Base class for retrievers in the Genkit framework.""" - - def __init__( - self, - retriever_fn: RetrieverFn[T], - ) -> None: - """Initialize a Retriever. - - Args: - retriever_fn: The function that performs the retrieval. - """ - self.retriever_fn: RetrieverFn[T] = retriever_fn - - -class RetrieverRequest(BaseModel): - """Request model for a retriever execution.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - query: DocumentData - options: Any | None = None - - -class RetrieverSupports(BaseModel): - """Retriever capability support.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - media: bool | None = None - - -class RetrieverInfo(BaseModel): - """Information about a retriever.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - label: str | None = None - supports: RetrieverSupports | None = None - - -class RetrieverOptions(BaseModel): - """Configuration options for a retriever.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, alias_generator=to_camel) - - config_schema: dict[str, Any] | None = None - label: str | None = None - supports: RetrieverSupports | None = None - - -class RetrieverRef(BaseModel): - """Reference to a retriever with configuration.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - name: str - config: Any | None = None - version: str | None = None - info: RetrieverInfo | None = None - - -def retriever_action_metadata( - name: str, - options: RetrieverOptions | None = None, -) -> ActionMetadata: - """Creates action metadata for a retriever.""" - options = options if options is not None else RetrieverOptions() - retriever_metadata_dict: dict[str, object] = {'retriever': {}} - retriever_info = cast(dict[str, object], retriever_metadata_dict['retriever']) - - if options.label: - retriever_info['label'] = options.label - - if options.supports: - retriever_info['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True) - - retriever_info['customOptions'] = options.config_schema if options.config_schema else None - return ActionMetadata( - kind=cast(ActionKind, ActionKind.RETRIEVER), - name=name, - input_json_schema=to_json_schema(RetrieverRequest), - output_json_schema=to_json_schema(RetrieverResponse), - metadata=retriever_metadata_dict, - ) - - -def create_retriever_ref( - name: str, - config: dict[str, Any] | None = None, - version: str | None = None, - info: RetrieverInfo | None = None, -) -> RetrieverRef: - """Creates a RetrieverRef instance.""" - return RetrieverRef(name=name, config=config, version=version, info=info) - - -class IndexerRequest(BaseModel): - """Request model for an indexer execution.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - documents: list[DocumentData] - options: Any | None = None - - -class IndexerInfo(BaseModel): - """Information about an indexer.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - label: str | None = None - supports: RetrieverSupports | None = None - - -class IndexerOptions(BaseModel): - """Configuration options for an indexer.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True, alias_generator=to_camel) - - config_schema: dict[str, Any] | None = None - label: str | None = None - supports: RetrieverSupports | None = None - - -class IndexerRef(BaseModel): - """Reference to an indexer with configuration.""" - - model_config: ClassVar[ConfigDict] = ConfigDict(extra='forbid', populate_by_name=True) - - name: str - config: Any | None = None - version: str | None = None - info: IndexerInfo | None = None - - -def indexer_action_metadata( - name: str, - options: IndexerOptions | None = None, -) -> ActionMetadata: - """Creates action metadata for an indexer.""" - options = options if options is not None else IndexerOptions() - indexer_metadata_dict: dict[str, object] = {'indexer': {}} - indexer_info = cast(dict[str, object], indexer_metadata_dict['indexer']) - - if options.label: - indexer_info['label'] = options.label - - if options.supports: - indexer_info['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True) - - indexer_info['customOptions'] = options.config_schema if options.config_schema else None - - return ActionMetadata( - kind=cast(ActionKind, ActionKind.INDEXER), - name=name, - input_json_schema=to_json_schema(IndexerRequest), - output_json_schema=to_json_schema(None), - metadata=indexer_metadata_dict, - ) - - -def create_indexer_ref( - name: str, - config: dict[str, Any] | None = None, - version: str | None = None, - info: IndexerInfo | None = None, -) -> IndexerRef: - """Creates a IndexerRef instance.""" - return IndexerRef(name=name, config=config, version=version, info=info) - - -def define_retriever( - registry: Registry, - name: str, - fn: RetrieverFn[Any], - options: RetrieverOptions | None = None, -) -> None: - """Defines and registers a retriever action.""" - metadata = retriever_action_metadata(name, options) - - async def wrapper( - request: RetrieverRequest, - _ctx: Any, # noqa: ANN401 - ) -> RetrieverResponse: - query = Document.from_document_data(request.query) - res = fn(query, request.options) - return await res if inspect.isawaitable(res) else res - - _ = registry.register_action( - kind=cast(ActionKind, ActionKind.RETRIEVER), - name=name, - fn=wrapper, - metadata=metadata.metadata, - span_metadata={'genkit:metadata:retriever:name': name}, - ) - - -IndexerFn = Callable[[list[Document], T], None | Awaitable[None]] - - -def define_indexer( - registry: Registry, - name: str, - fn: IndexerFn[Any], - options: IndexerOptions | None = None, -) -> None: - """Defines and registers an indexer action.""" - metadata = indexer_action_metadata(name, options) - - async def wrapper( - request: IndexerRequest, - _ctx: Any, # noqa: ANN401 - ) -> None: - docs = [Document.from_document_data(d) for d in request.documents] - res = fn(docs, request.options) - if inspect.isawaitable(res): - await res - - _ = registry.register_action( - kind=cast(ActionKind, ActionKind.INDEXER), - name=name, - fn=wrapper, - metadata=metadata.metadata, - span_metadata={'genkit:metadata:indexer:name': name}, - ) diff --git a/py/packages/genkit/src/genkit/blocks/tools.py b/py/packages/genkit/src/genkit/blocks/tools.py deleted file mode 100644 index 31954749c0..0000000000 --- a/py/packages/genkit/src/genkit/blocks/tools.py +++ /dev/null @@ -1,228 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tool-specific types and utilities for the Genkit framework. - -Genkit tools are actions that can be called by models during a generation -process. This module provides context and error types for tool execution, -including support for controlled interruptions and specific response formatting. - -Overview: - Tools extend the capabilities of AI models by allowing them to call - external functions during generation. The model decides when to use - a tool based on the conversation context and tool descriptions. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Tool Execution Flow β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Model β”‚ ───► β”‚ Tool β”‚ ───► β”‚ Execute β”‚ ───► β”‚ Model β”‚ β”‚ - β”‚ β”‚ Request β”‚ β”‚ Request β”‚ β”‚ Function β”‚ β”‚ Continue β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό (if interrupt=True) β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Pause β”‚ ────► User confirms ────► Resume β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Terminology: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Term β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Tool β”‚ A function that models can call during generation β”‚ - β”‚ ToolRunContext β”‚ Execution context with interrupt capability β”‚ - β”‚ ToolInterruptError β”‚ Exception for controlled tool execution pause β”‚ - β”‚ Interrupt β”‚ A tool marked to pause for user confirmation β”‚ - β”‚ tool_response() β”‚ Helper to construct response for interrupted tool β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Components: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Component β”‚ Purpose β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ ToolRunContext β”‚ Context for tool execution, extends ActionContext β”‚ - β”‚ ToolInterruptError β”‚ Exception to pause execution for user input β”‚ - β”‚ tool_response() β”‚ Constructs tool response Part for interrupts β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Basic tool definition: - - ```python - from genkit import Genkit - - ai = Genkit() - - - @ai.tool() - def get_weather(city: str) -> str: - '''Get current weather for a city.''' - # Fetch weather data... - return f'Weather in {city}: Sunny, 72Β°F' - - - # Use in generation - response = await ai.generate( - prompt='What is the weather in Paris?', - tools=['get_weather'], - ) - ``` - - Interrupt tool (human-in-the-loop): - - ```python - @ai.tool(interrupt=True) - def book_flight(destination: str, date: str) -> str: - '''Book a flight - requires user confirmation.''' - return f'Booked flight to {destination} on {date}' - - - # First generate - tool call is returned, not executed - response = await ai.generate( - prompt='Book me a flight to Paris next Friday', - tools=['book_flight'], - ) - - # Check for interrupts - if response.interrupts: - interrupt = response.interrupts[0] - # Show user: "Confirm booking to Paris on Friday?" - # Resume after confirmation - response = await ai.generate( - prompt='Book me a flight to Paris next Friday', - tools=['book_flight'], - messages=response.messages, - resume={'respond': tool_response(interrupt, 'Confirmed')}, - ) - ``` - -Caveats: - - Tools receive a ToolRunContext, which extends ActionRunContext - - Interrupt tools must be explicitly resumed to continue generation - - The tool_response() helper is used to respond to interrupted tools - -See Also: - - Interrupts documentation: https://genkit.dev/docs/tool-calling#pause-agentic-loops-with-interrupts - - genkit.core.action: Base action types -""" - -from typing import Any, NoReturn, cast - -from genkit.core.action import ActionRunContext -from genkit.core.typing import Metadata, Part, ToolRequest, ToolRequestPart, ToolResponse, ToolResponsePart - - -class ToolRunContext(ActionRunContext): - """Provides context specific to the execution of a Genkit tool. - - Inherits from ActionRunContext and adds functionality relevant to tools, - such as interrupting the tool's execution flow. - """ - - def __init__( - self, - ctx: ActionRunContext, - ) -> None: - """Initializes the ToolRunContext. - - Args: - ctx: The parent ActionRunContext. - """ - super().__init__( - on_chunk=ctx._on_chunk if ctx.is_streaming else None, - context=ctx.context, - ) - - def interrupt(self, metadata: dict[str, Any] | None = None) -> NoReturn: - """Interrupts the current tool execution. - - Raises a ToolInterruptError, which can be caught by the generation - process to handle controlled interruptions (e.g., asking the user for - clarification). - - Args: - metadata: Optional metadata to associate with the interrupt. - """ - raise ToolInterruptError(metadata=metadata) - - -# TODO(#4346): make this extend GenkitError once it has INTERRUPTED status -class ToolInterruptError(Exception): - """Exception raised to signal a controlled interruption of tool execution. - - This is used as a flow control mechanism within the generation process, - allowing a tool to pause execution and potentially signal back to the - calling flow (e.g., to request user input or clarification) without - causing a hard failure. - """ - - def __init__(self, metadata: dict[str, Any] | None = None) -> None: - """Initializes the ToolInterruptError. - - Args: - metadata: Metadata associated with the interruption. - """ - super().__init__() - self.metadata: dict[str, Any] = metadata or {} - - -def tool_response( - interrupt: Part | ToolRequestPart, - response_data: object | None = None, - metadata: dict[str, Any] | None = None, -) -> Part: - """Constructs a ToolResponse Part, typically for an interrupted request. - - This is often used when a tool's execution was interrupted (e.g., via - ToolInterruptError) and a specific response needs to be formulated and - sent back as part of the tool interaction history. - - Args: - interrupt: The original ToolRequest Part or ToolRequestPart that was interrupted. - response_data: The data to include in the ToolResponse output. Defaults to None. - metadata: Optional metadata to include in the resulting Part, often used - to signal that this response corresponds to an interrupt. - Defaults to {'interruptResponse': True}. - - Returns: - A Part object containing the constructed ToolResponse. - """ - # TODO(#4347): validate against tool schema - tool_request = interrupt.root.tool_request if isinstance(interrupt, Part) else interrupt.tool_request - - interrupt_metadata = True - if isinstance(metadata, Metadata): - interrupt_metadata = metadata.root - elif metadata: - interrupt_metadata = metadata - - tr = cast(ToolRequest, tool_request) - return Part( - root=ToolResponsePart( - tool_response=ToolResponse( - name=tr.name, - ref=tr.ref, - output=response_data, - ), - metadata=Metadata( - root={ - 'interruptResponse': interrupt_metadata, - } - ), - ) - ) diff --git a/py/packages/genkit/src/genkit/codec.py b/py/packages/genkit/src/genkit/codec.py deleted file mode 100644 index 70d69849a8..0000000000 --- a/py/packages/genkit/src/genkit/codec.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Encoding and decoding utilities for the Genkit framework. - -This module provides functions for serializing Genkit objects to dictionaries -and JSON strings, with special handling for Pydantic models and binary data. - -Overview: - Genkit uses Pydantic models extensively for type safety. This module - provides utilities to convert these models to dictionaries and JSON - strings for serialization, API responses, and debugging. - -Key Functions: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Function β”‚ Purpose β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ dump_dict() β”‚ Convert Pydantic model to dict (exclude_none=True) β”‚ - β”‚ dump_json() β”‚ Convert object to JSON string with alias support β”‚ - β”‚ default_serializerβ”‚ Fallback serializer for non-standard types (bytes) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Converting Pydantic models: - - ```python - from genkit.codec import dump_dict, dump_json - from genkit.types import Message, Part, TextPart - - msg = Message(role='user', content=[Part(root=TextPart(text='Hello'))]) - - # Convert to dict (excludes None values, uses aliases) - msg_dict = dump_dict(msg) - # {'role': 'user', 'content': [{'text': 'Hello'}]} - - # Convert to JSON string - msg_json = dump_json(msg, indent=2) - # '{"role": "user", "content": [{"text": "Hello"}]}' - ``` - -Caveats: - - Pydantic models use by_alias=True for JSON Schema compatibility - - None values are excluded from output (exclude_none=True) - - Binary data (bytes) is base64-encoded when serializing - -See Also: - - Pydantic documentation: https://docs.pydantic.dev/ - - genkit.core.typing: Core type definitions using Pydantic -""" - -import base64 -import json -from collections.abc import Callable - -from pydantic import BaseModel - - -def dump_dict(obj: object, fallback: Callable[[object], object] | None = None) -> object: - """Converts an object or Pydantic to a dictionary. - - If the input object is a Pydantic BaseModel, it returns a dictionary - representation of the model, excluding fields with `None` values and using - aliases for field names. For any other object type, it returns the object - unchanged. - - Args: - obj: The object to potentially convert to a dictionary. - fallback: A function to call when an unknown value is encountered. If not provided, error is raised. - - Returns: - A dictionary if the input is a Pydantic BaseModel, otherwise the - original object. - - Raises: - ValueError: If a circular reference is detected. - """ - - def _dump(o: object, seen: set[int]) -> object: - if isinstance(o, (list, dict)): - obj_id = id(o) - if obj_id in seen: - raise ValueError('Circular reference detected') - seen.add(obj_id) - - try: - if isinstance(o, BaseModel): - return o.model_dump(exclude_none=True, by_alias=True, fallback=fallback) - elif isinstance(o, list): - return [_dump(i, seen) for i in o] - elif isinstance(o, dict): - return {k: _dump(v, seen) for k, v in o.items()} - else: - return o - finally: - if isinstance(o, (list, dict)): - seen.remove(id(o)) - - return _dump(obj, set()) - - -def default_serializer(obj: object) -> object: - """Default serializer for objects not handled by json.dumps. - - Args: - obj: The object to serialize. - - Returns: - A serializable representation of the object. - """ - if isinstance(obj, bytes): - try: - return base64.b64encode(obj).decode('utf-8') - except Exception: - return '' - return str(obj) - - -def dump_json( - obj: object, - indent: int | None = None, - fallback: Callable[[object], object] | None = None, -) -> str: - """Dumps an object to a JSON string. - - If the object is a Pydantic BaseModel, it will be dumped using the - model_dump_json method using the by_alias flag set to True. Otherwise, the - object will be dumped using the json.dumps method. - - Args: - obj: The object to dump. - indent: The indentation level for the JSON string. - fallback: A function to call when an unknown value is encountered. If not provided, error is raised. - - Returns: - A JSON string. - """ - if isinstance(obj, BaseModel): - return obj.model_dump_json(by_alias=True, exclude_none=True, indent=indent, fallback=fallback) - else: - separators = (',', ':') if indent is None else None - return json.dumps(obj, indent=indent, default=fallback or default_serializer, separators=separators) diff --git a/py/packages/genkit/src/genkit/core/__init__.py b/py/packages/genkit/src/genkit/core/__init__.py deleted file mode 100644 index e908696bfa..0000000000 --- a/py/packages/genkit/src/genkit/core/__init__.py +++ /dev/null @@ -1,122 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Core foundations for the Genkit framework. - -This package provides the fundamental building blocks and abstractions used -throughout the Genkit framework. It includes the action system, plugin -architecture, registry, tracing, and schema types. - -Architecture Overview: - The core package forms the foundation layer upon which all Genkit - functionality is built. It provides primitives that are used by both - the framework internals and user-facing APIs. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ User Application β”‚ - β”‚ (flows, prompts, tools, chat) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ genkit.ai Layer β”‚ - β”‚ (Genkit class, GenkitRegistry) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ genkit.core Layer β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Action β”‚ Plugin β”‚ Registry β”‚ Trace β”‚ Schema β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Modules: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Module β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ action β”‚ Action class and execution context β”‚ - β”‚ plugin β”‚ Plugin base class for extending Genkit β”‚ - β”‚ registry β”‚ Central repository for actions and resources β”‚ - β”‚ tracing β”‚ OpenTelemetry-based tracing infrastructure β”‚ - β”‚ typing β”‚ Core type definitions (Part, Message, etc.) β”‚ - β”‚ schema β”‚ JSON Schema generation and validation β”‚ - β”‚ error β”‚ GenkitError and error handling utilities β”‚ - β”‚ logging β”‚ Structured logging via structlog β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Concepts: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Action β”‚ Strongly-typed, traceable, callable function. β”‚ - β”‚ β”‚ All models, tools, flows, embedders are actions. β”‚ - β”‚ ActionKind β”‚ Type of action: MODEL, TOOL, FLOW, EMBEDDER, etc. β”‚ - β”‚ ActionRunContext β”‚ Execution context with streaming and user context β”‚ - β”‚ Plugin β”‚ Extension mechanism for adding capabilities β”‚ - β”‚ Registry β”‚ Central storage for actions, schemas, and plugins β”‚ - β”‚ Trace β”‚ OpenTelemetry span for observability β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Usage: - Most users interact with the core layer indirectly through the ``Genkit`` - class. Direct usage is typically for plugin authors or advanced use cases: - - ```python - from genkit.core.action import Action - from genkit.core.plugin import Plugin - from genkit.core.registry import Registry - - # Create a registry - registry = Registry() - - # Register an action directly - action = registry.register_action( - kind=ActionKind.TOOL, - name='my_tool', - fn=my_tool_function, - ) - ``` - -See Also: - - genkit.ai: High-level user-facing API - - genkit.blocks: Building blocks for models, prompts, embedders - - genkit.types: Re-exported type definitions from core.typing -""" - -from .constants import GENKIT_CLIENT_HEADER, GENKIT_VERSION -from .http_client import clear_client_cache, close_cached_clients, get_cached_client -from .logging import Logger, get_logger - - -def package_name() -> str: - """Get the fully qualified package name. - - Returns: - The string 'genkit.core', which is the fully qualified package name. - """ - return 'genkit.core' - - -__all__ = [ - 'GENKIT_CLIENT_HEADER', - 'GENKIT_VERSION', - 'Logger', - 'clear_client_cache', - 'close_cached_clients', - 'get_cached_client', - 'get_logger', - 'package_name', -] diff --git a/py/packages/genkit/src/genkit/core/_compat.py b/py/packages/genkit/src/genkit/core/_compat.py deleted file mode 100644 index cc89ebe9f9..0000000000 --- a/py/packages/genkit/src/genkit/core/_compat.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Compatibility module for Python version differences. - -This module provides imports that work across different Python versions, -allowing the codebase to support Python 3.10+ while using newer typing features. -""" - -import sys - -# StrEnum - Added in Python 3.11 -# Used for string enums throughout the codebase -if sys.version_info >= (3, 11): - from enum import StrEnum -else: - from strenum import StrEnum - -# override decorator - Added in Python 3.12 -# We use this throughout the codebase to mark methods that override parent methods -if sys.version_info >= (3, 12): - from typing import override -else: - from typing_extensions import override - -__all__ = ['StrEnum', 'override'] diff --git a/py/packages/genkit/src/genkit/core/action/__init__.py b/py/packages/genkit/src/genkit/core/action/__init__.py deleted file mode 100644 index 4e47132989..0000000000 --- a/py/packages/genkit/src/genkit/core/action/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Action module for defining and managing RPC-over-HTTP functions.""" - -from ._action import Action, ActionMetadata, ActionRunContext -from ._key import create_action_key, parse_action_key -from ._tracing import SpanAttributeValue -from ._util import parse_plugin_name_from_action_name -from .types import ActionKind, ActionResponse - -__all__ = [ - 'Action', - 'ActionKind', - 'ActionMetadata', - 'ActionResponse', - 'ActionRunContext', - 'SpanAttributeValue', - 'create_action_key', - 'parse_action_key', - 'parse_plugin_name_from_action_name', -] diff --git a/py/packages/genkit/src/genkit/core/action/_action.py b/py/packages/genkit/src/genkit/core/action/_action.py deleted file mode 100644 index ad15da802b..0000000000 --- a/py/packages/genkit/src/genkit/core/action/_action.py +++ /dev/null @@ -1,685 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Action module for defining and managing remotely callable functions. - -This module provides the core `Action` class, the fundamental building block for -defining operations within the Genkit framework. - -## What is an Action? - -An Action represents a named, observable, and strongly-typed unit of work that -wraps a standard Python function (either synchronous or asynchronous). It serves -as a consistent interface for executing logic like calling models, running tools, -or orchestrating flows. - -## How it Works: - -1. Initialization - - * An `Action` is created with a unique `name`, a `kind` (e.g., MODEL, - TOOL, FLOW), and the Python function (`fn`) containing the core logic. - - * It automatically inspects the function's type hints (specifically the - first argument for input and the return annotation for output) using - Pydantic's `TypeAdapter` to generate JSON schemas (`input_schema`, - `output_schema`). These are stored for validation and metadata. It - raises a `TypeError` if the function signature has more than two - arguments (input, context). - - * It internally creates tracing wrappers (`_fn`, `_afn`) around the - original function using the `_make_tracing_wrappers` helper. These - wrappers handle OpenTelemetry span creation, recording input/output - metadata, and standardizing error handling by raising `GenkitError` - with a trace ID. The `_afn` wrapper ensures even synchronous functions - can be awaited. - -2. Execution Methods - - * `run()`: Executes the action synchronously. It calls the internal - synchronous tracing wrapper (`_fn`). - - * `arun()`: Executes the action asynchronously. It calls the internal - asynchronous tracing wrapper (`_afn`). This wrapper handles - awaiting the original function if it was async or running it via - `ensure_async` if it was sync. - - * `arun_raw()`: Similar to `arun`, but performs Pydantic validation on - the `raw_input` before calling `arun`. - - * `stream()`: Initiates an asynchronous execution via `arun` and returns - an `AsyncIterator` (via `Channel`) for receiving chunks and an - `asyncio.Future` that resolves with the final `ActionResponse`. - -3. Streaming and Context - - * During execution (`run`/`arun`/`stream`), an `ActionRunContext` instance - is created. - - * This context holds an `on_chunk` callback (provided by the caller, e.g., - by `stream()`) and any user-provided `context` dictionary. - - * If the wrapped function (`fn`) accepts a context argument (`ctx`), this - `ActionRunContext` instance is passed, allowing the function to send - intermediate chunks using `ctx.send_chunk()`. - - * A `ContextVar` (`_action_context`) is also used to propagate the user - context dictionary implicitly. - -The `Action` class provides a robust way to define executable units, abstracting -away the complexities of sync/async handling (for async callers), schema -generation, tracing, and streaming mechanics. -""" - -import asyncio -import inspect -import time -from collections.abc import AsyncIterator, Awaitable, Callable -from contextvars import ContextVar -from functools import cached_property -from typing import Any, Generic, Protocol, cast, get_type_hints - -from pydantic import BaseModel, TypeAdapter, ValidationError -from typing_extensions import Never, TypeVar - -from genkit.aio import Channel, ensure_async -from genkit.core.error import GenkitError -from genkit.core.tracing import tracer - -from ._tracing import ( - SpanAttributeValue, - record_input_metadata, - record_output_metadata, - save_parent_path, -) -from ._util import extract_action_args_and_types, noop_streaming_callback -from .types import ActionKind, ActionMetadataKey, ActionResponse - -InputT = TypeVar('InputT', default=Any) -OutputT = TypeVar('OutputT', default=Any) -ChunkT = TypeVar('ChunkT', default=Never) - -StreamingCallback = Callable[[object], None] - -_action_context: ContextVar[dict[str, object] | None] = ContextVar('context') -_ = _action_context.set(None) - - -class _LatencyTrackable(Protocol): - """Protocol for objects that support latency tracking.""" - - latency_ms: float - - -class _ModelCopyable(Protocol): - """Protocol for objects that support model_copy.""" - - def model_copy(self, *, update: dict[str, Any] | None = None) -> Any: # noqa: ANN401 - """Copy the model with optional updates.""" - ... - - -class ActionRunContext: - """Provides context for an action's execution, including streaming support. - - This class holds context information relevant to a single execution of an - Action. It manages the streaming callback (`on_chunk`) and any additional - context dictionary provided during the action call. - - Attributes: - context: A dictionary containing arbitrary context data. - is_streaming: Whether the action is being executed in streaming mode. - """ - - def __init__( - self, - on_chunk: StreamingCallback | None = None, - context: dict[str, object] | None = None, - on_trace_start: Callable[[str, str], None] | None = None, - ) -> None: - """Initializes an ActionRunContext instance. - - Sets up the context with an optional streaming callback and an optional - context dictionary. If `on_chunk` is None, a no-op callback is used. - - Args: - on_chunk: A callable to be invoked when a chunk of data is ready - during streaming execution. Defaults to a no-op function. - context: An optional dictionary containing context data to be made - available within the action execution. Defaults to an empty - dictionary. - on_trace_start: A callable to be invoked with the trace ID and span - ID when the trace is started. - """ - self._on_chunk: StreamingCallback = on_chunk if on_chunk is not None else noop_streaming_callback - self._context: dict[str, object] = context if context is not None else {} - self._on_trace_start: Callable[[str, str], None] = on_trace_start if on_trace_start else lambda _t, _s: None - - @property - def context(self) -> dict[str, object]: - return self._context - - @cached_property - def is_streaming(self) -> bool: - """Indicates whether the action is being executed in streaming mode. - - This property checks if a valid streaming callback (`on_chunk`) - was provided during initialization. - - Returns: - True if a streaming callback (other than the no-op default) is set, - False otherwise. - """ - return self._on_chunk != noop_streaming_callback - - def send_chunk(self, chunk: object) -> None: - """Send a chunk to from the action to the client. - - Args: - chunk: The chunk to send to the client. - """ - self._on_chunk(chunk) - - @staticmethod - def _current_context() -> dict[str, object] | None: - """Obtains current context if running within an action. - - Returns: - The current context if running within an action, None otherwise. - """ - return _action_context.get(None) - - -class Action(Generic[InputT, OutputT, ChunkT]): - """Represents a strongly-typed, remotely callable function within Genkit. - - Actions are the fundamental building blocks for defining operations in Genkit. - They are named, observable (via tracing), and support both streaming and - non-streaming execution modes. An Action wraps a Python function, handling - input validation, execution, tracing, and output serialization. - - Attributes: - name: A unique identifier for the action. - kind: The type category of the action (e.g., MODEL, TOOL, FLOW). - description: An optional human-readable description. - input_schema: The JSON schema definition for the expected input type. - output_schema: The JSON schema definition for the expected output type. - metadata: A dictionary for storing arbitrary metadata associated with the action. - is_async: Whether the action is asynchronous. - """ - - def __init__( - self, - kind: ActionKind, - name: str, - fn: Callable[..., OutputT | Awaitable[OutputT]], - metadata_fn: Callable[..., object] | None = None, - description: str | None = None, - metadata: dict[str, object] | None = None, - span_metadata: dict[str, SpanAttributeValue] | None = None, - ) -> None: - """Initialize an Action. - - Args: - kind: The kind of action (e.g., TOOL, MODEL, etc.). - name: Unique name identifier for this action. - fn: The function to call when the action is executed. - metadata_fn: The function to be used to infer metadata (e.g. - schemas). - description: Optional human-readable description of the action. - metadata: Optional dictionary of metadata about the action. - span_metadata: Optional dictionary of tracing span metadata. - """ - self._kind: ActionKind = kind - self._name: str = name - self._metadata: dict[str, object] = metadata if metadata else {} - self._description: str | None = description - self._is_async: bool = inspect.iscoroutinefunction(fn) - # Optional matcher function for resource actions - self.matches: Callable[[object], bool] | None = None - - input_spec = inspect.getfullargspec(metadata_fn if metadata_fn else fn) - try: - resolved_annotations = get_type_hints(metadata_fn if metadata_fn else fn) - except (NameError, TypeError, AttributeError): - resolved_annotations = input_spec.annotations - action_args, arg_types = extract_action_args_and_types(input_spec, resolved_annotations) - n_action_args = len(action_args) - fn_pair = _make_tracing_wrappers(name, kind, span_metadata or {}, n_action_args, fn) - self._fn: Callable[..., ActionResponse[OutputT]] = fn_pair[0] - self._afn: Callable[..., Awaitable[ActionResponse[OutputT]]] = fn_pair[1] - self._initialize_io_schemas(action_args, arg_types, resolved_annotations, input_spec) - - @property - def kind(self) -> ActionKind: - return self._kind - - @property - def name(self) -> str: - return self._name - - @property - def description(self) -> str | None: - return self._description - - @property - def metadata(self) -> dict[str, object]: - return self._metadata - - @property - def input_type(self) -> TypeAdapter[InputT] | None: - return self._input_type - - @property - def input_schema(self) -> dict[str, object]: - return self._input_schema - - @input_schema.setter - def input_schema(self, value: dict[str, object]) -> None: - """Update input schema (used by lazy-loaded prompts to set schema after registration).""" - self._input_schema = value - self._metadata[ActionMetadataKey.INPUT_KEY] = value - - @property - def output_schema(self) -> dict[str, object]: - return self._output_schema - - @output_schema.setter - def output_schema(self, value: dict[str, object]) -> None: - """Update output schema (used by lazy-loaded prompts to set schema after registration).""" - self._output_schema = value - self._metadata[ActionMetadataKey.OUTPUT_KEY] = value - - @property - def is_async(self) -> bool: - return self._is_async - - def run( - self, - input: InputT | None = None, - on_chunk: StreamingCallback | None = None, - context: dict[str, object] | None = None, - _telemetry_labels: dict[str, object] | None = None, - ) -> ActionResponse[OutputT]: - """Executes the action synchronously with the given input. - - This method runs the action's underlying function synchronously. - It handles input validation, tracing, and output serialization. - If the action's function is async, it will be run in the current event loop. - - Args: - input: The input data for the action. It should conform to the action's - input schema. - on_chunk: An optional callback function to receive streaming output chunks. - Note: For synchronous execution of streaming actions, chunks - will be delivered synchronously via this callback. - context: An optional dictionary containing context data for the execution. - telemetry_labels: Optional labels for telemetry. - - Returns: - An ActionResponse object containing the final result and trace ID. - - Raises: - GenkitError: If an error occurs during action execution. - """ - # TODO(#4348): handle telemetry_labels - - if context: - _ = _action_context.set(context) - - return self._fn( - input, - ActionRunContext(on_chunk=on_chunk, context=_action_context.get(None)), - ) - - async def arun( - self, - input: InputT | None = None, - on_chunk: StreamingCallback | None = None, - context: dict[str, object] | None = None, - on_trace_start: Callable[[str, str], None] | None = None, - _telemetry_labels: dict[str, object] | None = None, - ) -> ActionResponse[OutputT]: - """Executes the action asynchronously with the given input. - - This method runs the action's underlying function asynchronously. - It handles input validation, tracing, and output serialization. - If the action's function is synchronous, it will be wrapped to run - asynchronously. - - Args: - input: The input data for the action. It should conform to the action's - input schema. - on_chunk: An optional callback function to receive streaming output chunks. - context: An optional dictionary containing context data for the execution. - on_trace_start: An optional callback to be invoked with the trace ID - and span ID when the trace is started. - telemetry_labels: Optional labels for telemetry. - - Returns: - An awaitable ActionResponse object containing the final result and trace ID. - - Raises: - GenkitError: If an error occurs during action execution. - """ - # TODO(#4348): handle telemetry_labels - - if context: - _ = _action_context.set(context) - - return await self._afn( - input, - ActionRunContext(on_chunk=on_chunk, context=_action_context.get(None), on_trace_start=on_trace_start), - ) - - async def arun_raw( - self, - raw_input: InputT | None = None, - on_chunk: StreamingCallback | None = None, - context: dict[str, object] | None = None, - on_trace_start: Callable[[str, str], None] | None = None, - telemetry_labels: dict[str, object] | None = None, - ) -> ActionResponse[OutputT]: - """Executes the action asynchronously with raw, unvalidated input. - - This method bypasses the Pydantic input validation and calls the underlying - action function directly with the provided `raw_input`. It still handles - tracing and context management. - - Use this method when you need to work with input that may not conform - to the defined schema or when you have already validated the input. - - Args: - raw_input: The raw input data to pass directly to the action function. - on_chunk: An optional callback function to receive streaming output chunks. - context: An optional dictionary containing context data for the execution. - on_trace_start: An optional callback to be invoked with the trace ID - and span ID when the trace is started. - telemetry_labels: Optional labels for telemetry. - - Returns: - An awaitable ActionResponse object containing the final result and trace ID. - - Raises: - GenkitError: If an error occurs during action execution, or if - the action requires input but none was provided. - """ - input_action: InputT | None = None - if self._input_type is not None: - if raw_input is None: - raise GenkitError( - message=( - f"Action '{self.name}' requires input but none was provided. " - 'Please supply a valid input payload.' - ), - status='INVALID_ARGUMENT', - ) - input_action = self._input_type.validate_python(raw_input) - - return await self.arun( - input=input_action, - on_chunk=on_chunk, - context=context, - on_trace_start=on_trace_start, - _telemetry_labels=telemetry_labels, - ) - - def stream( - self, - input: InputT | None = None, - context: dict[str, object] | None = None, - telemetry_labels: dict[str, object] | None = None, - timeout: float | None = None, - ) -> tuple[AsyncIterator[ChunkT], asyncio.Future[ActionResponse[OutputT]]]: - """Executes the action asynchronously and provides a streaming response. - - This method initiates an asynchronous action execution and returns immediately - with a tuple containing an async iterator for the chunks and a future for the - final response. - - Args: - input: The input data for the action. It should conform to the action's - input schema. - context: An optional dictionary containing context data for the execution. - telemetry_labels: Optional labels for telemetry. - timeout: Optional timeout for the stream. - - Returns: - A tuple: (chunk_iterator, final_response_future) - - chunk_iterator: An AsyncIterator yielding output chunks as they become available. - - final_response_future: An asyncio.Future that will resolve to the - complete ActionResponse when the action finishes. - """ - stream: Channel[ChunkT, ActionResponse[OutputT]] = Channel(timeout=timeout) - - def send_chunk(c: object) -> None: - stream.send(cast(ChunkT, c)) - - resp = self.arun( - input=input, - context=context, - _telemetry_labels=telemetry_labels, - on_chunk=send_chunk, - ) - stream.set_close_future(asyncio.create_task(resp)) - - result_future: asyncio.Future[ActionResponse[OutputT]] = asyncio.Future() - stream.closed.add_done_callback(lambda _: result_future.set_result(stream.closed.result())) - - return (stream, result_future) - - def _initialize_io_schemas( - self, - action_args: list[str], - arg_types: list[type], - annotations: dict[str, Any], - _input_spec: inspect.FullArgSpec, - ) -> None: - """Initializes input/output schemas based on function signature and hints. - - Uses Pydantic's TypeAdapter to generate JSON schemas for the first - argument (if present) and the return type annotation (if present). - Stores schemas on the instance (input_schema, output_schema) and in - the metadata dictionary. - - Args: - action_args: List of detected argument names. - arg_types: List of detected argument types. - annotations: Type annotations dict from function signature. - _input_spec: The FullArgSpec object from inspecting the function. - - Raises: - TypeError: If the function has more than two arguments. - """ - if len(action_args) > 2: - raise TypeError(f'can only have up to 2 arg: {action_args}') - - if len(action_args) > 0: - type_adapter = TypeAdapter(arg_types[0]) - self._input_schema: dict[str, object] = type_adapter.json_schema() - self._input_type: TypeAdapter[Any] | None = type_adapter - self._metadata[ActionMetadataKey.INPUT_KEY] = self._input_schema - else: - self._input_schema = TypeAdapter(object).json_schema() - self._input_type = None - self._metadata[ActionMetadataKey.INPUT_KEY] = self._input_schema - - if ActionMetadataKey.RETURN in annotations: - type_adapter = TypeAdapter(annotations[ActionMetadataKey.RETURN]) - self._output_schema: dict[str, object] = type_adapter.json_schema() - self._metadata[ActionMetadataKey.OUTPUT_KEY] = self._output_schema - else: - self._output_schema = TypeAdapter(object).json_schema() - self._metadata[ActionMetadataKey.OUTPUT_KEY] = self._output_schema - - -class ActionMetadata(BaseModel): - """Metadata for actions.""" - - kind: ActionKind - name: str - description: str | None = None - input_schema: object | None = None - input_json_schema: dict[str, object] | None = None - output_schema: object | None = None - output_json_schema: dict[str, object] | None = None - stream_schema: object | None = None - metadata: dict[str, object] | None = None - - -_SyncTracingWrapper = Callable[[object | None, ActionRunContext], ActionResponse[Any]] -_AsyncTracingWrapper = Callable[[object | None, ActionRunContext], Awaitable[ActionResponse[Any]]] - - -def _make_tracing_wrappers( - name: str, - kind: ActionKind, - span_metadata: dict[str, SpanAttributeValue], - n_action_args: int, - fn: Callable[..., object], -) -> tuple[_SyncTracingWrapper, _AsyncTracingWrapper]: - """Make the sync and async tracing wrappers for an action function. - - Args: - name: The name of the action. - kind: The kind of action. - span_metadata: The span metadata for the action. - n_action_args: The arguments of the action. - fn: The function to wrap. - """ - - def _record_latency(output: object, start_time: float) -> object: - """Record latency for the action if the output supports it. - - Args: - output: The action output. - start_time: The start time of the action execution. - - Returns: - The updated action output. - """ - latency_ms = (time.perf_counter() - start_time) * 1000 - if hasattr(output, 'latency_ms'): - try: - cast(_LatencyTrackable, output).latency_ms = latency_ms - except (TypeError, ValidationError, AttributeError): - # If immutable (e.g. Pydantic model with frozen=True), try model_copy - if hasattr(output, 'model_copy'): - output = cast(_ModelCopyable, output).model_copy(update={'latency_ms': latency_ms}) - return output - - async def async_tracing_wrapper(input: object | None, ctx: ActionRunContext) -> ActionResponse[Any]: - """Wrap the function in an async tracing wrapper. - - Args: - input: The input to the action. - ctx: The context to pass to the action. - - Returns: - The action response. - """ - afn = ensure_async(fn) - start_time = time.perf_counter() - - with save_parent_path(): - with tracer.start_as_current_span(name) as span: - # Format trace_id and span_id as hex strings (OpenTelemetry standard format) - trace_id = format(span.get_span_context().trace_id, '032x') - span_id = format(span.get_span_context().span_id, '016x') - ctx._on_trace_start(trace_id, span_id) # pyright: ignore[reportPrivateUsage] - record_input_metadata( - span=span, - kind=kind, - name=name, - span_metadata=span_metadata, - input=input, - ) - - try: - match n_action_args: - case 0: - output = await afn() - case 1: - output = await afn(input) - case 2: - output = await afn(input, ctx) - case _: - raise ValueError('action fn must have 0-2 args...') - except Exception as e: - # Re-raise existing GenkitError instances to avoid double-wrapping - if isinstance(e, GenkitError): - raise - raise GenkitError( - cause=e, - message=f'Error while running action {name}', - trace_id=trace_id, - ) from e - - output = _record_latency(output, start_time) - record_output_metadata(span, output=output) - return ActionResponse(response=output, trace_id=trace_id, span_id=span_id) - - def sync_tracing_wrapper(input: object | None, ctx: ActionRunContext) -> ActionResponse[Any]: - """Wrap the function in a sync tracing wrapper. - - Args: - input: The input to the action. - ctx: The context to pass to the action. - - Returns: - The action response. - """ - start_time = time.perf_counter() - - with save_parent_path(): - with tracer.start_as_current_span(name) as span: - # Format trace_id and span_id as hex strings (OpenTelemetry standard format) - trace_id = format(span.get_span_context().trace_id, '032x') - span_id = format(span.get_span_context().span_id, '016x') - ctx._on_trace_start(trace_id, span_id) # pyright: ignore[reportPrivateUsage] - record_input_metadata( - span=span, - kind=kind, - name=name, - span_metadata=span_metadata, - input=input, - ) - - try: - match n_action_args: - case 0: - output = fn() - case 1: - output = fn(input) - case 2: - output = fn(input, ctx) - case _: - raise ValueError('action fn must have 0-2 args...') - except Exception as e: - # Re-raise existing GenkitError instances to avoid double-wrapping - if isinstance(e, GenkitError): - raise - raise GenkitError( - cause=e, - message=f'Error while running action {name}', - trace_id=trace_id, - ) from e - - output = _record_latency(output, start_time) - record_output_metadata(span, output=output) - return ActionResponse(response=output, trace_id=trace_id, span_id=span_id) - - return sync_tracing_wrapper, async_tracing_wrapper diff --git a/py/packages/genkit/src/genkit/core/action/_key.py b/py/packages/genkit/src/genkit/core/action/_key.py deleted file mode 100644 index d4ef94deec..0000000000 --- a/py/packages/genkit/src/genkit/core/action/_key.py +++ /dev/null @@ -1,61 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Action key module for defining and managing action keys.""" - -from .types import ActionKind - - -def parse_action_key(key: str) -> tuple[ActionKind, str]: - """Parse an action key into its kind and name components. - - Args: - key: The action key to parse, in the format "/kind/name". - - Returns: - A tuple containing the ActionKind and name. - - Raises: - ValueError: If the key format is invalid or if the kind is not a valid - ActionKind. - """ - tokens = key.split('/') - if len(tokens) < 3 or not tokens[1] or not tokens[2]: - msg = f'Invalid action key format: `{key}`.Expected format: `//`' - raise ValueError(msg) - - kind_str = tokens[1] - name = '/'.join(tokens[2:]) - try: - kind = ActionKind(kind_str) - except ValueError as e: - msg = f'Invalid action kind: `{kind_str}`' - raise ValueError(msg) from e - # pyrefly: ignore[bad-return] - ActionKind is StrEnum subclass, pyrefly doesn't narrow properly - return kind, name - - -def create_action_key(kind: ActionKind, name: str) -> str: - """Create an action key from its kind and name components. - - Args: - kind: The kind of action. - name: The name of the action. - - Returns: - The action key in the format `//`. - """ - return f'/{kind}/{name}' diff --git a/py/packages/genkit/src/genkit/core/action/_tracing.py b/py/packages/genkit/src/genkit/core/action/_tracing.py deleted file mode 100644 index 3869a6b161..0000000000 --- a/py/packages/genkit/src/genkit/core/action/_tracing.py +++ /dev/null @@ -1,192 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Action tracing module for defining and managing action tracing.""" - -from collections.abc import Generator -from contextlib import contextmanager -from contextvars import ContextVar -from urllib.parse import quote - -from opentelemetry.trace import Span -from opentelemetry.util import types as otel_types - -from genkit.codec import dump_json - -# Type alias for span attribute values -SpanAttributeValue = otel_types.AttributeValue - -# Context variable to track parent path across nested spans -_parent_path_context: ContextVar[str] = ContextVar('genkit_parent_path', default='') - - -def build_path( - name: str, - parent_path: str, - type_str: str, - subtype: str | None = None, -) -> str: - """Build a hierarchical path with type annotations. - - Args: - name: The name of the action/flow/step. - parent_path: The path of the parent span (empty string for root). - type_str: The type (e.g., 'flow', 'action', 'flowStep'). - subtype: Optional subtype (e.g., 'tool', 'model', 'flow'). - - Returns: - Annotated path string. - - Examples: - >>> build_path('myFlow', '', 'flow') - '/{myFlow,t:flow}' - - >>> build_path('myTool', '/{myFlow,t:flow}', 'action', 'tool') - '/{myFlow,t:flow}/{myTool,t:action,s:tool}' - """ - # URL-encode name to handle special characters - name = quote(name, safe='') - - # Build the path segment with type annotation - if type_str: - path_segment = f'{name},t:{type_str}' - else: - path_segment = name - - # Add subtype if provided - if subtype: - path_segment = f'{path_segment},s:{subtype}' - - # Wrap in braces and append to parent path - path_segment = '{' + path_segment + '}' - return parent_path + '/' + path_segment - - -def decorate_path_with_subtype(path: str, subtype: str) -> str: - """Add subtype annotation to the leaf node of a path. - - Args: - path: The path to decorate. - subtype: The subtype to add (e.g., 'tool', 'model', 'flow'). - - Returns: - Decorated path string. - - Examples: - >>> decorate_path_with_subtype('/{myFlow,t:flow}/{myTool,t:action}', 'tool') - '/{myFlow,t:flow}/{myTool,t:action,s:tool}' - """ - if not path or not subtype: - return path - - # Find the last opening brace - last_brace_idx = path.rfind('{') - if last_brace_idx == -1: - return path # No braces found - - # Find the closing brace after the last opening brace - closing_brace_idx = path.find('}', last_brace_idx) - if closing_brace_idx == -1: - return path # No closing brace found - - # Extract the content of the last segment (without braces) - segment_content = path[last_brace_idx + 1 : closing_brace_idx] - - # Check if subtype already exists - if any(p.strip().startswith('s:') for p in segment_content.split(',')[1:]): - return path - - # Add subtype annotation - decorated_content = segment_content + ',s:' + subtype - - # Rebuild the path with the decorated last segment - return path[: last_brace_idx + 1] + decorated_content + path[closing_brace_idx:] - - -@contextmanager -def save_parent_path() -> Generator[None, None, None]: - """Context manager to save and restore parent path. - - Yields: - None - """ - saved = _parent_path_context.get() - try: - yield - finally: - _parent_path_context.set(saved) - - -def record_input_metadata( - span: Span, - kind: str, - name: str, - span_metadata: dict[str, SpanAttributeValue] | None, - input: object | None, -) -> None: - """Records input metadata onto an OpenTelemetry span for a Genkit action. - - Sets standard Genkit attributes like action type, subtype (kind), name, - path, qualifiedPath, and the JSON representation of the input. Also adds - any custom span metadata provided. - - Args: - span: The OpenTelemetry Span object to add attributes to. - kind: The kind (e.g., 'model', 'flow', 'tool') of the action. - name: The specific name of the action. - span_metadata: An optional dictionary of custom key-value pairs to add - as span attributes. - input: The input data provided to the action. - """ - span.set_attribute('genkit:type', 'action') - span.set_attribute('genkit:metadata:subtype', kind) - span.set_attribute('genkit:name', name) - if input is not None: - span.set_attribute('genkit:input', dump_json(input)) - - # Build and set path attributes (qualified path with full annotations) - parent_path = _parent_path_context.get() - qualified_path = build_path(name, parent_path, 'action', kind) - - # IMPORTANT: Span attributes store the QUALIFIED path (full annotated) - # Telemetry handlers will derive display path using to_display_path() - span.set_attribute('genkit:path', qualified_path) - span.set_attribute('genkit:qualifiedPath', qualified_path) - - # Update context for nested spans - _parent_path_context.set(qualified_path) - - if span_metadata is not None: - for meta_key in span_metadata: - span.set_attribute(meta_key, span_metadata[meta_key]) - - -def record_output_metadata(span: Span, output: object) -> None: - """Records output metadata onto an OpenTelemetry span for a Genkit action. - - Marks the span state as 'success' and records the JSON representation of - the action's output. - - Args: - span: The OpenTelemetry Span object to add attributes to. - output: The output data returned by the action. - """ - span.set_attribute('genkit:state', 'success') - try: - span.set_attribute('genkit:output', dump_json(output)) - except Exception: - # Fallback for non-serializable output - span.set_attribute('genkit:output', str(output)) diff --git a/py/packages/genkit/src/genkit/core/action/_util.py b/py/packages/genkit/src/genkit/core/action/_util.py deleted file mode 100644 index e690289b16..0000000000 --- a/py/packages/genkit/src/genkit/core/action/_util.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Action utility module for defining and managing action utilities.""" - -import inspect -from collections.abc import Mapping -from typing import Any - - -def noop_streaming_callback(_chunk: Any) -> None: # noqa: ANN401 - """A no-op streaming callback. - - This callback does nothing and is used when no streaming is desired. - - Args: - _chunk: The chunk (unused, callback is a no-op). - - Returns: - None. - """ - pass - - -def parse_plugin_name_from_action_name(name: str) -> str | None: - """Parses the plugin name from an action name. - - As per convention, the plugin name is optional. If present, it's the first - part of the action name, separated by a forward slash: `pluginname/*`. - - Args: - name: The action name string. - - Returns: - The plugin name, or None if no plugin name is found in the action name. - """ - tokens = name.split('/') - if len(tokens) > 1: - return tokens[0] - return None - - -def extract_action_args_and_types( - input_spec: inspect.FullArgSpec, - annotations: Mapping[str, Any] | None = None, -) -> tuple[list[str], list[Any]]: - """Extracts relevant argument names and types from a function's FullArgSpec. - - Specifically handles the case where the first argument might be 'self' - (for methods) and determines the type hint for each argument. - - Args: - input_spec: The FullArgSpec object obtained from - inspect.getfullargspec(). - annotations: Optional override for type annotations. If not provided, - uses input_spec.annotations. - - Returns: - A tuple containing: - - A list of argument names (potentially excluding 'self'). - - A list of corresponding argument types (using Any if no - annotation). - """ - arg_types = [] - action_args = input_spec.args.copy() - resolved_annotations = annotations or input_spec.annotations - - # Special case when using a method as an action, we ignore first "self" - # arg. (Note: The original condition `len(action_args) <= 3` is preserved - # from the source snippet). - if len(action_args) > 0 and len(action_args) <= 3 and action_args[0] == 'self': - del action_args[0] - - for arg in action_args: - arg_types.append(resolved_annotations.get(arg, Any)) - - return action_args, arg_types diff --git a/py/packages/genkit/src/genkit/core/action/types.py b/py/packages/genkit/src/genkit/core/action/types.py deleted file mode 100644 index 6232522134..0000000000 --- a/py/packages/genkit/src/genkit/core/action/types.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Action types module for defining and managing action types.""" - -from collections.abc import Callable -from typing import ClassVar, Generic, TypeVar - -from pydantic import BaseModel, ConfigDict -from pydantic.alias_generators import to_camel - -from genkit.core._compat import StrEnum - -# Type alias for action name. -# type ActionName = str -ActionName = str - - -# type ActionResolver = Callable[[ActionKind, str], None] -ActionResolver = Callable[['ActionKind', str], None] - - -class ActionKind(StrEnum): - """Enumerates all the types of action that can be registered. - - This enum defines the various types of actions supported by the framework, - including chat models, embedders, evaluators, and other utility functions. - """ - - BACKGROUND_MODEL = 'background-model' - CANCEL_OPERATION = 'cancel-operation' - CHECK_OPERATION = 'check-operation' - CUSTOM = 'custom' - DYNAMIC_ACTION_PROVIDER = 'dynamic-action-provider' - EMBEDDER = 'embedder' - EVALUATOR = 'evaluator' - EXECUTABLE_PROMPT = 'executable-prompt' - FLOW = 'flow' - INDEXER = 'indexer' - MODEL = 'model' - PROMPT = 'prompt' - RERANKER = 'reranker' - RESOURCE = 'resource' - RETRIEVER = 'retriever' - TOOL = 'tool' - UTIL = 'util' - - -ResponseT = TypeVar('ResponseT') - - -class ActionResponse(BaseModel, Generic[ResponseT]): - """The response from an action. - - Attributes: - response: The actual response data from the action execution. - trace_id: A unique identifier for tracing the action execution. - span_id: The span ID of the root action span. - """ - - model_config: ClassVar[ConfigDict] = ConfigDict( - extra='forbid', populate_by_name=True, alias_generator=to_camel, arbitrary_types_allowed=True - ) - - response: ResponseT - trace_id: str - span_id: str = '' - - -class ActionMetadataKey(StrEnum): - """Enumerates all the keys of the action metadata. - - Attributes: - INPUT_KEY: Key for the input schema metadata. - OUTPUT_KEY: Key for the output schema metadata. - RETURN: Key for the return type metadata. - """ - - INPUT_KEY = 'inputSchema' - OUTPUT_KEY = 'outputSchema' - RETURN = 'return' diff --git a/py/packages/genkit/src/genkit/core/environment.py b/py/packages/genkit/src/genkit/core/environment.py deleted file mode 100644 index e945bc86e2..0000000000 --- a/py/packages/genkit/src/genkit/core/environment.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Convenience functionality to determine the running environment.""" - -import os -from typing import cast - -from genkit.core._compat import StrEnum - - -class EnvVar(StrEnum): - """Enumerates all the environment variables used by Genkit.""" - - GENKIT_ENV = 'GENKIT_ENV' - - -class GenkitEnvironment(StrEnum): - """Enumerates all the environments Genkit can run in.""" - - DEV = 'dev' - PROD = 'prod' - - -def is_dev_environment() -> bool: - """Returns True if the current environment is a development environment. - - Returns: - True if the current environment is a development environment. - """ - return get_current_environment() == GenkitEnvironment.DEV - - -def is_prod_environment() -> bool: - """Returns True if the current environment is a production environment. - - Returns: - True if the current environment is a production environment. - """ - return get_current_environment() == GenkitEnvironment.PROD - - -def get_current_environment() -> GenkitEnvironment: - """Returns the current environment. - - Returns: - The current environment. - """ - env = os.getenv(EnvVar.GENKIT_ENV) - if env is None: - return cast(GenkitEnvironment, GenkitEnvironment.PROD) - try: - # pyrefly: ignore[bad-return] - GenkitEnvironment is StrEnum subclass, pyrefly doesn't narrow properly - return GenkitEnvironment(env) - except ValueError: - return cast(GenkitEnvironment, GenkitEnvironment.PROD) diff --git a/py/packages/genkit/src/genkit/core/extract.py b/py/packages/genkit/src/genkit/core/extract.py deleted file mode 100644 index df638c9f9a..0000000000 --- a/py/packages/genkit/src/genkit/core/extract.py +++ /dev/null @@ -1,263 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Utility functions for extracting JSON data from text and markdown.""" - -from typing import Any - -import json5 -from partial_json_parser import loads - -CHAR_NON_BREAKING_SPACE = '\u00a0' - - -def parse_partial_json(json_string: str) -> Any: # noqa: ANN401 - """Parses a partially complete JSON string and returns the parsed object. - - This function attempts to parse the given JSON string, even if it is not - a complete or valid JSON document. - - Args: - json_string: The string to parse as JSON. - - Returns: - The parsed JSON object. - - Raises: - AssertionError: If the string cannot be parsed as JSON. - """ - # TODO(#4350): add handling for malformed JSON cases. - return loads(json_string) - - -def extract_json(text: str, throw_on_bad_json: bool = True) -> Any: # noqa: ANN401 - """Extracts JSON from a string with lenient parsing. - - This function attempts to extract a valid JSON object or array from a - string, even if the string contains extraneous characters or minor - formatting issues. It uses a combination of basic parsing and - `json5` and `partial-json` libraries to maximize the chance of - successful extraction. - - Args: - text: The string to extract JSON from. - throw_on_bad_json: If True, raises a ValueError if no valid JSON - can be extracted. If False, returns None in such cases. - - Returns: - The extracted JSON object (dict or list), or None if no valid - JSON is found and `throw_on_bad_json` is False. - - Raises: - ValueError: If `throw_on_bad_json` is True and no valid JSON - can be extracted, or if parsing an incomplete structure fails. - - Examples: - >>> extract_json(' { "key" : "value" } ') - {'key': 'value'} - - >>> extract_json('{"key": "value",}') # Trailing comma - {'key': 'value'} - - >>> extract_json('some text {"key": "value"} more text') - {'key': 'value'} - - >>> extract_json('invalid json', throw_on_bad_json=False) - None - """ - if text.strip() == '': - return None - - # Explicit type annotations for loop variables to help type checkers - opening_char: str | None = None - closing_char: str | None = None - start_pos: int | None = None - nesting_count: int = 0 - in_string: bool = False - escape_next: bool = False - - for i in range(len(text)): - char = text[i].replace(CHAR_NON_BREAKING_SPACE, ' ') - - if escape_next: - escape_next = False - continue - - if char == '\\': - escape_next = True - continue - - if char == '"': - in_string = not in_string - continue - - if in_string: - continue - - if not opening_char and (char == '{' or char == '['): - # Look for opening character - opening_char = char - closing_char = '}' if char == '{' else ']' - start_pos = i - nesting_count += 1 - elif char == opening_char: - # Increment nesting for matching opening character - nesting_count += 1 - elif char == closing_char: - # Decrement nesting for matching closing character - nesting_count -= 1 - if not nesting_count: - # Reached end of target element - return json5.loads(text[start_pos or 0 : i + 1]) - if start_pos is not None and nesting_count > 0: - # If an incomplete JSON structure is detected - try: - # Parse the incomplete JSON structure using partial-json for lenient parsing - return parse_partial_json(text[start_pos:]) - except Exception as e: - # If parsing fails, throw an error - if throw_on_bad_json: - raise ValueError(f'Invalid JSON extracted from model output: {text}') from e - return None - - if throw_on_bad_json: - raise ValueError(f'Invalid JSON extracted from model output: {text}') - return None - - -class ExtractItemsResult: - """Holds the result of extracting items from a text array. - - Attributes: - items: A list of the extracted JSON objects. - cursor: The index in the original text immediately after the last - processed character. - """ - - def __init__(self, items: list[Any], cursor: int) -> None: - """Initialize an ExtractItemsResult. - - Args: - items: A list of the extracted JSON objects. - cursor: The updated cursor position. - """ - self.items: list[Any] = items - self.cursor: int = cursor - - -def extract_items(text: str, cursor: int = 0) -> ExtractItemsResult: - """Extracts complete JSON objects from the first array found in the text. - - This function searches for the first JSON array within the input string, - starting from an optional cursor position. It extracts complete JSON - objects from this array and returns them along with an updated cursor - position, indicating how much of the string has been processed. - - Args: - text: The string to extract items from. - cursor: The starting position for searching the array (default: 0). - Useful for processing large strings in chunks. - - Returns: - An `ExtractItemsResult` object containing: - - `items`: A list of extracted JSON objects (dictionaries). - - `cursor`: The updated cursor position, which is the index - immediately after the last processed character. If no array is - found, the cursor will be the length of the text. - - Examples: - >>> text = '[{"a": 1}, {"b": 2}, {"c": 3}]' - >>> result = extract_items(text) - >>> result.items - [{'a': 1}, {'b': 2}, {'c': 3}] - >>> result.cursor - 29 - - >>> text = ' [ {"x": 10}, {"y": 20} ] ' - >>> result = extract_items(text) - >>> result.items - [{'x': 10}, {'y': 20}] - >>> result.cursor - 25 - - >>> text = 'some text [ {"p": 100} , {"q": 200} ] more text' - >>> result = extract_items(text, cursor=10) - >>> result.items - [{'p': 100}, {'q': 200}] - >>> result.cursor - 35 - - >>> text = 'no array here' - >>> result = extract_items(text) - >>> result.items - [] - >>> result.cursor - 13 - """ - items = [] - current_cursor = cursor - - # Find the first array start if we haven't already processed any text - if cursor == 0: - array_start = text.find('[') - if array_start == -1: - return ExtractItemsResult(items=[], cursor=len(text)) - current_cursor = array_start + 1 - - object_start = -1 - brace_count = 0 - in_string = False - escape_next = False - - # Process the text from the cursor position - for i in range(current_cursor, len(text)): - char = text[i] - - if escape_next: - escape_next = False - continue - - if char == '\\': - escape_next = True - continue - - if char == '"': - in_string = not in_string - continue - - if in_string: - continue - - if char == '{': - if brace_count == 0: - object_start = i - brace_count += 1 - elif char == '}': - brace_count -= 1 - if brace_count == 0 and object_start != -1: - try: - obj = json5.loads(text[object_start : i + 1]) - items.append(obj) - current_cursor = i + 1 - object_start = -1 - except Exception: # noqa: S110 - intentionally silent, parsing partial JSON - # If parsing fails, continue trying next position - pass - elif char == ']' and brace_count == 0: - # End of array - break - - return ExtractItemsResult(items=items, cursor=current_cursor) diff --git a/py/packages/genkit/src/genkit/core/flows.py b/py/packages/genkit/src/genkit/core/flows.py deleted file mode 100644 index 442722f487..0000000000 --- a/py/packages/genkit/src/genkit/core/flows.py +++ /dev/null @@ -1,370 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Flows API module for GenKit. - -This module implements the Flows API server functionality for GenKit, providing -endpoints to execute both streaming and standard flows. The Flows API allows -clients to invoke registered flows over HTTP, with support for: - -1. **Standard request/response flows**: Synchronous execution with a single - response. - -2. **Streaming flows**: Asynchronous execution that streams partial results to - clients. - -The module provides: - -- An ASGI application factory (create_flows_asgi_app) that creates a Starlette - application with appropriate routes. -- Request handlers for both streaming and standard flows. -- Proper error handling and response formatting. -- Context providers for request-specific execution contexts. -- Endpoints for health checks and flow execution, supporting both streaming and - non-streaming responses. - -## Example usage: - - ```python - from genkit.core.flows import create_flows_asgi_app - from genkit.core.registry import Registry - - # Create a registry with your flows - registry = Registry() - registry.register(my_flow, name='my_flow') - - # Create the ASGI app - app = create_flows_asgi_app(registry) - - # Run with any ASGI server - import uvicorn - - uvicorn.run(app, host='localhost', port=3400) - ``` - -For a higher-level server implementation, see the FlowsServer class in -genkit.core.flow_server. -""" - -from __future__ import annotations - -import asyncio -import json -from collections.abc import AsyncGenerator, Callable -from typing import Any - -from sse_starlette.sse import EventSourceResponse -from starlette.applications import Starlette -from starlette.middleware import Middleware -from starlette.middleware.cors import CORSMiddleware -from starlette.requests import Request -from starlette.responses import JSONResponse -from starlette.routing import Route - -from genkit.codec import dump_dict -from genkit.core.action import Action -from genkit.core.constants import DEFAULT_GENKIT_VERSION -from genkit.core.error import get_callable_json -from genkit.core.logging import get_logger -from genkit.core.registry import Registry -from genkit.web.requests import ( - is_streaming_requested, -) -from genkit.web.typing import ( - Application, - StartupHandler, -) - -logger = get_logger(__name__) - - -# TODO(#4351): This is a work in progress and may change. Do not use. -def create_flows_asgi_app( - registry: Registry, - context_providers: list[Callable[..., Any]] | None = None, - on_app_startup: StartupHandler | None = None, - on_app_shutdown: StartupHandler | None = None, - version: str = DEFAULT_GENKIT_VERSION, -) -> Application: - """Create an ASGI application for flows. - - Args: - registry: The registry to use for the flows server. - context_providers: Optional list of context providers to process - requests. Each provider is a callable that takes a context and - request data and returns an enriched context. - on_app_startup: Optional callback to execute when the app's - lifespan starts. Must be an async function. - on_app_shutdown: Optional callback to execute when the app's - lifespan ends. Must be an async function. - version: The version of the Genkit server to use. - - Returns: - An ASGI application. - """ - routes = [] - logger = get_logger(__name__) - - async def health_check(_request: Request) -> JSONResponse: - """Handle health check requests. - - Args: - _request: The Starlette request object (unused). - - Returns: - A JSON response with status code 200. - """ - return JSONResponse(content={'status': 'OK'}) - - async def handle_run_flows( - request: Request, - ) -> JSONResponse | EventSourceResponse: - """Handle flow execution. - - Flow: - 1. Extracts flow name from the path - 2. Reads and validates the request payload - 3. Looks up the requested flow action - 4. Executes the flow with the provided input - 5. Returns the flow result as JSON - - Args: - request: The Starlette request object. - - Returns: - A JSON or EventSourceResponse with the flow result, or an error - response. - """ - flow_name = request.path_params.get('flow_name') - if not flow_name: - return JSONResponse( - content={'error': 'Flow name not provided'}, - status_code=400, - ) - - try: - # Look up the flow action. - action = await registry.resolve_action_by_key(flow_name) - if action is None: - await logger.aerror( - 'Flow not found', - error=f'Flow not found: {flow_name}', - ) - return JSONResponse( - content={'error': f'Flow not found: {flow_name}'}, - status_code=404, - ) - - # Parse request body. - try: - input_data = {} - if await request.body(): - payload = await request.json() - input_data = payload.get('data', {}) - except json.JSONDecodeError as e: - await logger.aerror( - 'Invalid JSON', - error=f'Invalid JSON: {e!s}', - ) - return JSONResponse( - content={'error': f'Invalid JSON: {e!s}'}, - status_code=400, - ) - - # Set up context. - ctx = {} - if context_providers: - headers = {k.lower(): v for k, v in request.headers.items()} - request_data = { - 'method': request.method, - 'headers': headers, - 'input': input_data, - } - - for provider in context_providers: - try: - provider_ctx = await provider(request.app.state.context, request_data) - ctx.update(provider_ctx) - except Exception as e: - await logger.aerror( - 'context provider error', - error=str(e), - ) - return JSONResponse( - content={'error': f'Unauthorized: {e!s}'}, - status_code=401, - ) - - # Run the flow. - stream = is_streaming_requested(request) - handler = handle_streaming_flow if stream else handle_standard_flow - return await handler(action, input_data, ctx, version) - except Exception as e: - await logger.aerror('error executing flow', error=str(e)) - error_response = {'error': str(e)} - return JSONResponse( - content=error_response, - status_code=500, - ) - - async def handle_streaming_flow( - action: Action, - input_data: dict[str, Any], - context: dict[str, Any], - version: str, - ) -> EventSourceResponse: - """Handle streaming flow execution. - - Args: - action: The flow action to execute. - input_data: Input data for the flow. - context: Execution context. - version: The Genkit version header value. - - Returns: - An EventSourceResponse with the flow result or error. - """ - - async def stream_generator() -> AsyncGenerator[dict[str, str], None]: - """Generate stream of data dictionaries for the SSE response.""" - # Use an asyncio.Queue for true streaming - chunks are yielded as they arrive - chunk_queue: asyncio.Queue[dict[str, str] | None] = asyncio.Queue() - result_holder: list[object] = [] - error_holder: list[Exception] = [] - - def chunk_callback(chunk: object) -> None: - # Put chunk into queue (non-blocking since queue is unbounded) - # Use dump_dict to properly serialize Pydantic models with field aliases - chunk_queue.put_nowait({ - 'event': 'message', - 'data': json.dumps({'message': dump_dict(chunk)}), - }) - - async def run_action() -> None: - try: - output = await action.arun_raw( - raw_input=input_data, - on_chunk=chunk_callback, - context=context, - ) - result_holder.append(output.response) - except Exception as e: - error_holder.append(e) - finally: - # Signal completion - await chunk_queue.put(None) - - # Start the action in the background - action_task = asyncio.create_task(run_action()) - - # Yield chunks as they arrive - while True: - item = await chunk_queue.get() - if item is None: - break - yield item - - # Wait for task to complete (should already be done) - await action_task - - # Handle result or error - if error_holder: - error_msg = str(error_holder[0]) - await logger.aerror('error in stream', error=error_msg) - yield { - 'event': 'error', - 'data': json.dumps({ - 'error': { - 'status': 'INTERNAL', - 'message': 'stream flow error', - 'details': error_msg, - } - }), - } - elif result_holder: - result = dump_dict(result_holder[0]) - yield { - 'event': 'result', - 'data': json.dumps({'result': result}), - } - - return EventSourceResponse( - stream_generator(), - headers={ - 'x-genkit-version': version, - 'Cache-Control': 'no-cache', - 'Connection': 'keep-alive', - }, - ) - - async def handle_standard_flow( - action: Action, - input_data: object, - context: dict[str, Any], - version: str, - ) -> JSONResponse: - """Handle standard (non-streaming) flow execution. - - Args: - action: The flow action to execute. - input_data: Input data for the flow. - context: Execution context. - version: The Genkit version header value. - - Returns: - A JSONResponse with the flow result or error. - """ - try: - output = await action.arun_raw(raw_input=input_data, context=context) - - result = dump_dict(output.response) - response = {'result': result} - - return JSONResponse( - content=response, - status_code=200, - headers={'x-genkit-version': version}, - ) - except Exception as e: - error_response = get_callable_json(e).model_dump(by_alias=True) - await logger.aerror('error executing flow', error=error_response) - return JSONResponse( - content={'error': error_response}, - status_code=500, - ) - - routes = [ - Route('/__health', health_check, methods=['GET']), - Route('/{flow_name:path}', handle_run_flows, methods=['POST']), - ] - - app = Starlette( - routes=routes, - middleware=[ - Middleware( - CORSMiddleware, # type: ignore[arg-type] - allow_origins=['*'], - allow_methods=['*'], - allow_headers=['*'], - ) - ], - on_startup=[on_app_startup] if on_app_startup else [], - on_shutdown=[on_app_shutdown] if on_app_shutdown else [], - ) - - app.state.context = {} - - return app # pyright: ignore[reportReturnType] diff --git a/py/packages/genkit/src/genkit/core/http_client.py b/py/packages/genkit/src/genkit/core/http_client.py deleted file mode 100644 index 3ecad01da7..0000000000 --- a/py/packages/genkit/src/genkit/core/http_client.py +++ /dev/null @@ -1,263 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Shared HTTP client utilities for Genkit plugins. - -This module provides utilities for managing httpx.AsyncClient instances -across different event loops. The key feature is per-event-loop caching -that ensures: - -1. Clients are reused within the same event loop (avoiding connection - setup overhead) -2. Each event loop gets its own client (avoiding "bound to different - event loop" errors) -3. Automatic cleanup when event loops are garbage collected - -Example Usage -------------- - -Basic usage with default settings:: - - from genkit.core.http_client import get_cached_client - - - async def my_async_function(): - client = get_cached_client( - cache_key='my-plugin', - headers={'Authorization': 'Bearer token'}, - ) - response = await client.get('https://api.example.com') - -With custom timeout:: - - client = get_cached_client( - cache_key='my-plugin', - headers={'Authorization': 'Bearer token'}, - timeout=httpx.Timeout(120.0, connect=30.0), - ) - -Multiple cache keys for different configurations:: - - # One client for API calls - api_client = get_cached_client( - cache_key='my-plugin/api', - headers={'Authorization': 'Bearer api-token'}, - ) - - # Another client for media fetching (different headers) - media_client = get_cached_client( - cache_key='my-plugin/media', - headers={'User-Agent': 'MyApp/1.0'}, - ) - -Implementation Notes --------------------- - -The cache uses a two-level structure: - -1. Outer level: WeakKeyDictionary keyed by event loop - - Automatically cleans up when event loop is garbage collected - - Prevents memory leaks from abandoned event loops - -2. Inner level: Regular dict keyed by cache_key string - - Allows multiple clients per event loop with different configurations - - Cache key should include any configuration that affects the client - -The client configuration (headers, timeout, etc.) is only used when -creating a new client. If a cached client exists, the provided -configuration is ignored. This is intentional - changing configuration -requires a new cache key. - -Thread Safety -------------- - -This module is thread-safe. A threading.Lock protects all modifications -and multi-step read operations on the cache. In multi-threaded applications -where different threads manage different event loops, this prevents race -conditions between concurrent cache operations. -""" - -import asyncio -import threading -import weakref -from collections.abc import MutableMapping -from typing import Any - -import httpx - -from genkit.core.logging import get_logger - -logger = get_logger(__name__) - -# Two-level cache: event_loop -> (cache_key -> client) -# Using WeakKeyDictionary for event loops ensures automatic cleanup -_loop_clients: MutableMapping[asyncio.AbstractEventLoop, dict[str, httpx.AsyncClient]] = weakref.WeakKeyDictionary() - -# Lock for thread-safe cache operations -_cache_lock = threading.Lock() - - -def get_cached_client( - cache_key: str, - headers: dict[str, str] | None = None, - timeout: httpx.Timeout | float | None = None, - **httpx_kwargs: Any, # noqa: ANN401 -) -> httpx.AsyncClient: - """Get or create a cached httpx.AsyncClient for the current event loop. - - This function provides per-event-loop client caching, which: - - Reuses clients within the same event loop (reduces connection overhead) - - Creates separate clients for different event loops (avoids binding errors) - - Automatically cleans up when event loops are garbage collected - - Args: - cache_key: Unique identifier for this client configuration. - Use a consistent key for the same configuration to benefit from - caching. Include any distinguishing factors in the key (e.g., - 'my-plugin/api' vs 'my-plugin/media' for different use cases). - headers: HTTP headers to include in all requests. - timeout: Request timeout. Can be a float (total timeout in seconds) - or an httpx.Timeout object for fine-grained control. - Defaults to 60s total with 10s connect timeout. - **httpx_kwargs: Additional arguments passed to httpx.AsyncClient(). - - Returns: - A cached or newly created httpx.AsyncClient instance. - - Raises: - RuntimeError: If called outside of an async context (no running loop). - - Note: - The client configuration is only used when creating a new client. - If a cached client exists, the provided configuration is ignored. - To use different configurations, use different cache_key values. - - Example:: - - # In an async function - client = get_cached_client( - cache_key='vertex-ai-evaluator', - headers={'Authorization': f'Bearer {token}'}, - timeout=60.0, - ) - response = await client.post(url, json=data) - """ - try: - loop = asyncio.get_running_loop() - except RuntimeError as e: - raise RuntimeError( - 'get_cached_client() must be called from within an async context ' - '(inside an async function with a running event loop)' - ) from e - - with _cache_lock: - # Get or create the cache dict for this event loop - loop_cache = _loop_clients.setdefault(loop, {}) - - # Check if we have a cached client for this key - if cache_key in loop_cache: - client = loop_cache[cache_key] - # Verify the client is still usable - if not client.is_closed: - return client - # Client was closed, remove from cache - logger.debug('Cached client was closed, creating new one', cache_key=cache_key) - del loop_cache[cache_key] - - # Create a new client - if timeout is None: - timeout = httpx.Timeout(60.0, connect=10.0) - elif isinstance(timeout, (int, float)): - timeout = httpx.Timeout(float(timeout)) - - client = httpx.AsyncClient( - headers=headers or {}, - timeout=timeout, - **httpx_kwargs, - ) - - loop_cache[cache_key] = client - logger.debug('Created new httpx client', cache_key=cache_key, loop_id=id(loop)) - - return client - - -async def close_cached_clients(cache_key: str | None = None) -> None: - """Close and remove cached clients. - - This is useful for cleanup in tests or when reconfiguring clients. - In normal usage, clients are automatically cleaned up when their - event loop is garbage collected. - - Args: - cache_key: If provided, only close clients with this key. - If None, close all clients in the current event loop's cache. - - Example:: - - # Close specific client - await close_cached_clients('my-plugin') - - # Close all clients in current event loop - await close_cached_clients() - """ - try: - loop = asyncio.get_running_loop() - except RuntimeError: - return # No running loop, nothing to close - - # Collect clients to close under the lock, then close outside the lock - # to avoid blocking other threads during async I/O - clients_to_close: dict[str, httpx.AsyncClient] = {} - - with _cache_lock: - if loop not in _loop_clients: - return - - loop_cache = _loop_clients[loop] - - if cache_key is not None: - # Close specific client - if cache_key in loop_cache: - clients_to_close[cache_key] = loop_cache.pop(cache_key) - else: - # Close all clients in this loop's cache - clients_to_close.update(loop_cache) - loop_cache.clear() - - # Close clients outside the lock to avoid blocking - for key, client in clients_to_close.items(): - try: - await client.aclose() - logger.debug('Closed cached client', cache_key=key) - except Exception as e: - logger.warning('Failed to close cached client', cache_key=key, error=e) - - -def clear_client_cache() -> None: - """Clear all cached clients across all event loops. - - This is primarily for testing purposes. Clients are NOT closed, - just removed from the cache. Use close_cached_clients() to properly - close clients before clearing. - - Warning: - This will cause any existing client references to become orphaned. - Only use this in tests or when you're sure no clients are in use. - """ - with _cache_lock: - _loop_clients.clear() - logger.debug('Cleared all client caches') diff --git a/py/packages/genkit/src/genkit/core/logging.py b/py/packages/genkit/src/genkit/core/logging.py deleted file mode 100644 index 84f5f2bfa0..0000000000 --- a/py/packages/genkit/src/genkit/core/logging.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright 2025 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -"""Typed logging utilities for Genkit. - -This module provides a typed wrapper around structlog to eliminate type warnings -and improve IDE support. The `Logger` protocol defines the interface used -throughout Genkit, and `get_logger()` returns a properly typed logger instance. - -Usage: - from genkit.core.logging import get_logger - - logger = get_logger(__name__) - logger.info("Server started", port=8080) - await logger.ainfo("Async operation complete") -""" - -from typing import Protocol - -import structlog - - -class Logger(Protocol): - """Protocol defining the logger interface used throughout Genkit. - - This protocol matches structlog's BoundLogger interface, providing type hints - for all standard logging methods including async variants. - """ - - # Synchronous logging methods - def debug(self, event: str | None = None, **kw: object) -> None: - """Log a debug message.""" - ... - - def info(self, event: str | None = None, **kw: object) -> None: - """Log an info message.""" - ... - - def warning(self, event: str | None = None, **kw: object) -> None: - """Log a warning message.""" - ... - - def warn(self, event: str | None = None, **kw: object) -> None: - """Log a warning message (alias for warning).""" - ... - - def error(self, event: str | None = None, **kw: object) -> None: - """Log an error message.""" - ... - - def exception(self, event: str | None = None, **kw: object) -> None: - """Log an exception with traceback.""" - ... - - def critical(self, event: str | None = None, **kw: object) -> None: - """Log a critical message.""" - ... - - def fatal(self, event: str | None = None, **kw: object) -> None: - """Log a fatal message (alias for critical).""" - ... - - # Async logging methods - async def adebug(self, event: str | None = None, **kw: object) -> None: - """Log a debug message asynchronously.""" - ... - - async def ainfo(self, event: str | None = None, **kw: object) -> None: - """Log an info message asynchronously.""" - ... - - async def awarning(self, event: str | None = None, **kw: object) -> None: - """Log a warning message asynchronously.""" - ... - - async def awarn(self, event: str | None = None, **kw: object) -> None: - """Log a warning message asynchronously (alias for awarning).""" - ... - - async def aerror(self, event: str | None = None, **kw: object) -> None: - """Log an error message asynchronously.""" - ... - - async def aexception(self, event: str | None = None, **kw: object) -> None: - """Log an exception with traceback asynchronously.""" - ... - - async def acritical(self, event: str | None = None, **kw: object) -> None: - """Log a critical message asynchronously.""" - ... - - async def afatal(self, event: str | None = None, **kw: object) -> None: - """Log a fatal message asynchronously (alias for acritical).""" - ... - - # Context binding - def bind(self, **new_values: object) -> 'Logger': - """Return a new logger with bound context values.""" - ... - - def unbind(self, *keys: str) -> 'Logger': - """Return a new logger with specified keys removed from context.""" - ... - - def try_unbind(self, *keys: str) -> 'Logger': - """Return a new logger with specified keys removed (ignoring missing).""" - ... - - def new(self, **new_values: object) -> 'Logger': - """Return a new logger with only the specified context values.""" - ... - - -def get_logger(name: str | None = None) -> Logger: - """Get a typed logger instance. - - This is a typed wrapper around structlog.get_logger() that provides - proper type hints for IDE support and type checking. - - Args: - name: Optional logger name (typically __name__). - - Returns: - A typed logger instance. - - Example: - >>> logger = get_logger(__name__) - >>> logger.info('Server started', port=8080) - >>> await logger.ainfo('Async operation complete') - """ - # The cast is safe because structlog's BoundLogger implements these methods - return structlog.get_logger(name) diff --git a/py/packages/genkit/src/genkit/core/plugin.py b/py/packages/genkit/src/genkit/core/plugin.py deleted file mode 100644 index 510a4dab3b..0000000000 --- a/py/packages/genkit/src/genkit/core/plugin.py +++ /dev/null @@ -1,220 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Abstract base class for Genkit plugins. - -This module defines the base plugin interface that all plugins must implement. -Plugins extend Genkit by registering models, embedders, retrievers, and -other action types. - -Overview: - Plugins are the primary extension mechanism in Genkit. They allow adding - support for AI providers (Google AI, Anthropic, OpenAI), vector stores, - and other capabilities. Each plugin implements a standard interface for - initialization, action resolution, and action discovery. - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Plugin Lifecycle β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Register β”‚ ───► β”‚ init β”‚ ───► β”‚ resolve β”‚ ───► β”‚ Action β”‚ β”‚ - β”‚ β”‚ Plugin β”‚ β”‚ (async) β”‚ β”‚ (lazy) β”‚ β”‚ Returns β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ β–Ό β”‚ - β”‚ β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ └─────────►│ List β”‚ β”‚ - β”‚ β”‚ Actions β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Terminology: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Term β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Plugin β”‚ Abstract base class that defines the plugin API β”‚ - β”‚ name β”‚ Plugin namespace (e.g., 'googleai', 'anthropic') β”‚ - β”‚ init() β”‚ Async method called once on first action resolve β”‚ - β”‚ resolve() β”‚ Lazy resolution of actions by kind and name β”‚ - β”‚ list_actions() β”‚ Returns metadata for plugin's available actions β”‚ - β”‚ model() β”‚ Helper to create namespaced ModelReference β”‚ - β”‚ embedder() β”‚ Helper to create namespaced EmbedderRef β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Methods: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Method β”‚ Purpose β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ init() β”‚ One-time initialization; return pre-registered β”‚ - β”‚ β”‚ actions (e.g., known models to pre-register) β”‚ - β”‚ resolve() β”‚ Create/return Action for a given kind and name; β”‚ - β”‚ β”‚ called when action is first used β”‚ - β”‚ list_actions() β”‚ Return ActionMetadata for dev UI action discovery β”‚ - β”‚ β”‚ (should be fast, no heavy initialization) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Implementing a custom plugin: - - ```python - from genkit.core.plugin import Plugin - from genkit.core.action import Action, ActionMetadata - from genkit.core.action.types import ActionKind - - - class MyPlugin(Plugin): - name = 'myplugin' - - def __init__(self, api_key: str) -> None: - self.api_key = api_key - - async def init(self) -> list[Action]: - # Return actions to pre-register (optional) - return [] - - async def resolve(self, action_type: ActionKind, name: str) -> Action | None: - if action_type == ActionKind.MODEL: - model_name = name.replace(f'{self.name}/', '') - return self._create_model_action(model_name) - return None - - async def list_actions(self) -> list[ActionMetadata]: - # Return metadata for available models (for dev UI) - return [ - ActionMetadata(kind=ActionKind.MODEL, name='myplugin/my-model'), - ] - - def _create_model_action(self, name: str) -> Action: - # Create and return the model action - ... - ``` - - Using the plugin: - - ```python - from genkit import Genkit - - ai = Genkit(plugins=[MyPlugin(api_key='...')]) - response = await ai.generate(model='myplugin/my-model', prompt='Hello!') - ``` - -Caveats: - - init() is called lazily on first action resolution, not at registration - - Plugin names must be unique within a Genkit instance - - Actions returned from init() are pre-registered with the plugin namespace - - resolve() receives the fully namespaced name (e.g., 'plugin/model') - -See Also: - - Built-in plugins: genkit.plugins.google_genai, genkit.plugins.anthropic - - Registry: genkit.core.registry -""" - -import abc -from typing import TYPE_CHECKING - -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionKind - -if TYPE_CHECKING: - from genkit.blocks.embedding import EmbedderRef - from genkit.blocks.model import ModelReference - - -class Plugin(abc.ABC): - """Abstract base class for implementing Genkit plugins. - - This class defines the async plugin interface that all plugins must implement. - Plugins provide a way to extend functionality by registering new actions, models, - or other capabilities. - """ - - name: str # plugin namespace - - @abc.abstractmethod - async def init(self) -> list[Action]: - """Lazy warm-up called once per plugin per registry instance. - - This method is called lazily when the first action resolution attempt - involving this plugin occurs. It should return a list of Action objects - to pre-register. - - Returns: - list[Action]: A list of Action instances to register. - """ - ... - - @abc.abstractmethod - async def resolve(self, action_type: ActionKind, name: str) -> Action | None: - """Resolve a single action. - - The registry will call this with a namespaced name (e.g., "plugin/model-name"). - - Args: - action_type: The kind of action to resolve. - name: The namespaced name of the action to resolve. - - Returns: - Action | None: The Action instance if found, None otherwise. - """ - ... - - @abc.abstractmethod - async def list_actions(self) -> list[ActionMetadata]: - """Advertised set for dev UI/reflection listing endpoint. - - This method should be safe and ideally inexpensive. It returns the set - of actions that this plugin advertises without triggering full initialization. - - Returns: - list[ActionMetadata]: A list of ActionMetadata objects describing - available actions. - """ - ... - - def model(self, name: str) -> 'ModelReference': - """Creates a model reference. - - Prefixes local name with plugin namespace. - - Args: - name: The model name (local or namespaced). - - Returns: - ModelReference: A reference to the model. - """ - # Deferred import: avoid circular import with genkit.blocks.model - from genkit.blocks.model import ModelReference # noqa: PLC0415 - - target = name if '/' in name else f'{self.name}/{name}' - return ModelReference(name=target) - - def embedder(self, name: str) -> 'EmbedderRef': - """Creates an embedder reference. - - Prefixes local name with plugin namespace. - - Args: - name: The embedder name (local or namespaced). - - Returns: - EmbedderRef: A reference to the embedder. - """ - # Deferred import: avoid circular import with genkit.blocks.embedding - from genkit.blocks.embedding import EmbedderRef # noqa: PLC0415 - - target = name if '/' in name else f'{self.name}/{name}' - return EmbedderRef(name=target) diff --git a/py/packages/genkit/src/genkit/core/reflection.py b/py/packages/genkit/src/genkit/core/reflection.py deleted file mode 100644 index b9c699a9b1..0000000000 --- a/py/packages/genkit/src/genkit/core/reflection.py +++ /dev/null @@ -1,598 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Development API for inspecting and interacting with Genkit. - -This module provides a reflection API server for inspection and interaction -during development. It exposes endpoints for health checks, action discovery, -and action execution. - -## Caveats - -The reflection API server predates the flows server implementation and differs -in the protocol it uses to interface with the Dev UI. The streaming protocol -uses unadorned JSON per streamed chunk. This may change in the future to use -Server-Sent Events (SSE). - -## Key endpoints - - | Method | Path | Handler | - |--------|---------------------|-----------------------| - | GET | /api/__health | Health check | - | GET | /api/actions | List actions | - | POST | /api/__quitquitquit | Trigger shutdown | - | POST | /api/notify | Handle notification | - | POST | /api/runAction | Run action (streaming)| -""" - -from __future__ import annotations - -import asyncio -import json -import os -import signal -from collections.abc import AsyncGenerator, Callable -from typing import Any, cast - -from starlette.applications import Starlette -from starlette.middleware import Middleware -from starlette.middleware.cors import CORSMiddleware -from starlette.requests import Request -from starlette.responses import JSONResponse, StreamingResponse -from starlette.routing import Route - -from genkit.codec import dump_dict, dump_json -from genkit.core.action import Action -from genkit.core.action.types import ActionKind -from genkit.core.constants import DEFAULT_GENKIT_VERSION -from genkit.core.error import get_reflection_json -from genkit.core.logging import get_logger -from genkit.core.registry import Registry -from genkit.web.requests import ( - is_streaming_requested, -) -from genkit.web.typing import ( - Application, - StartupHandler, -) - -logger = get_logger(__name__) - - -async def _list_registered_actions(registry: Registry) -> dict[str, Action]: - """Return all locally registered actions keyed as `//`. - - Uses resolve_actions_by_kind() to trigger lazy loading for any actions with - deferred metadata (e.g., file-based prompts), ensuring schemas are available - for the Dev UI. - """ - registered: dict[str, Action] = {} - for kind in ActionKind.__members__.values(): - for name, action in (await registry.resolve_actions_by_kind(kind)).items(): - registered[f'/{kind.value}/{name}'] = action - return registered - - -def _build_actions_payload( - *, - registered_actions: dict[str, Action], - plugin_metas: list[Any], -) -> dict[str, dict[str, Any]]: - """Build payload for GET /api/actions.""" - actions: dict[str, dict[str, Any]] = {} - - # 1) Registered actions (flows/tools/etc). - for key, action in registered_actions.items(): - actions[key] = { - 'key': key, - 'name': action.name, - 'type': action.kind.value, - 'description': action.description, - 'inputSchema': action.input_schema, - 'outputSchema': action.output_schema, - 'metadata': action.metadata, - } - - # 2) Plugin-advertised actions (may not be registered yet). - for meta in plugin_metas or []: - try: - key = f'/{meta.kind.value}/{meta.name}' - except Exception as exc: - # Defensive: skip unexpected plugin metadata objects. - logger.warning('Skipping invalid plugin action metadata', error=str(exc)) - continue - - advertised = { - 'key': key, - 'name': meta.name, - 'type': meta.kind.value, - 'description': getattr(meta, 'description', None), - 'inputSchema': getattr(meta, 'input_json_schema', None), - 'outputSchema': getattr(meta, 'output_json_schema', None), - 'metadata': getattr(meta, 'metadata', None), - } - - if key not in actions: - actions[key] = advertised - continue - - # Merge into the existing (registered) action entry; prefer registered data. - existing = actions[key] - - if not existing.get('description') and advertised.get('description'): - existing['description'] = advertised['description'] - - if not existing.get('inputSchema') and advertised.get('inputSchema'): - existing['inputSchema'] = advertised['inputSchema'] - - if not existing.get('outputSchema') and advertised.get('outputSchema'): - existing['outputSchema'] = advertised['outputSchema'] - - existing_meta = existing.get('metadata') or {} - advertised_meta = advertised.get('metadata') or {} - if isinstance(existing_meta, dict) and isinstance(advertised_meta, dict): - # Prefer registered action metadata on key conflicts. - existing['metadata'] = {**advertised_meta, **existing_meta} - - return actions - - -def create_reflection_asgi_app( - registry: Registry, - on_app_startup: StartupHandler | None = None, - on_app_shutdown: StartupHandler | None = None, - version: str = DEFAULT_GENKIT_VERSION, - _encoding: str = 'utf-8', -) -> Application: - """Create and return a ASGI application for the Genkit reflection API. - - Caveats: - - The reflection API server predates the flows server implementation and - differs in the protocol it uses to interface with the Dev UI. The - streaming protocol uses unadorned JSON per streamed chunk. This may - change in the future to use Server-Sent Events (SSE). - - Key endpoints: - - | Method | Path | Handler | - |--------|---------------------|-----------------------| - | GET | /api/__health | Health check | - | GET | /api/actions | List actions | - | POST | /api/__quitquitquit | Trigger shutdown | - | POST | /api/notify | Handle notification | - | POST | /api/runAction | Run action (streaming)| - - Args: - registry: The registry to use for the reflection server. - on_app_startup: Optional callback to execute when the app's - lifespan starts. Must be an async function. - on_app_shutdown: Optional callback to execute when the app's - lifespan ends. Must be an async function. - version: The version string to use when setting the value of - the X-GENKIT-VERSION HTTP header. - encoding: The text encoding to use; default 'utf-8'. - - Returns: - An ASGI application configured with the given registry. - """ - - async def handle_health_check(_request: Request) -> JSONResponse: - """Handle health check requests. - - Args: - _request: The Starlette request object (unused). - - Returns: - A JSON response with status code 200. - """ - return JSONResponse(content={'status': 'OK'}) - - async def handle_terminate(_request: Request) -> JSONResponse: - """Handle the quit endpoint. - - Args: - _request: The Starlette request object (unused). - - Returns: - An empty JSON response with status code 200. - """ - await logger.ainfo('Shutting down servers...') - asyncio.get_running_loop().call_soon(os.kill, os.getpid(), signal.SIGTERM) - return JSONResponse(content={'status': 'OK'}) - - async def handle_list_actions(_request: Request) -> JSONResponse: - """Handle the request for listing available actions. - - Args: - _request: The Starlette request object (unused). - - Returns: - A JSON response containing all serializable actions. - """ - registered = await _list_registered_actions(registry) - metas = await registry.list_actions() - actions = _build_actions_payload(registered_actions=registered, plugin_metas=metas) - - return JSONResponse( - content=actions, - status_code=200, - headers={'x-genkit-version': version}, - ) - - async def handle_list_values(request: Request) -> JSONResponse: - """Handle the request for listing registered values. - - Args: - request: The Starlette request object. - - Returns: - A JSON response containing value names. - """ - kind = request.query_params.get('type') - if not kind: - return JSONResponse(content={'error': 'Query parameter "type" is required.'}, status_code=400) - - if kind != 'defaultModel': - return JSONResponse( - content={'error': f"'type' {kind} is not supported. Only 'defaultModel' is supported"}, status_code=400 - ) - - values = registry.list_values(kind) - return JSONResponse(content=values, status_code=200) - - async def handle_list_envs(_request: Request) -> JSONResponse: - """Handle the request for listing environments. - - Args: - _request: The Starlette request object (unused). - - Returns: - A JSON response containing environments. - """ - return JSONResponse(content=['dev'], status_code=200) - - async def handle_notify(_request: Request) -> JSONResponse: - """Handle the notification endpoint. - - Args: - _request: The Starlette request object (unused). - - Returns: - An empty JSON response with status code 200. - """ - return JSONResponse( - content={}, - status_code=200, - headers={'x-genkit-version': version}, - ) - - # Map of active actions indexed by trace ID for cancellation support. - active_actions: dict[str, asyncio.Task[Any]] = {} - - async def handle_cancel_action(request: Request) -> JSONResponse: - """Handle the cancelAction endpoint. - - Args: - request: The Starlette request object. - - Returns: - A JSON response. - """ - try: - payload = await request.json() - trace_id = payload.get('traceId') - if not trace_id: - return JSONResponse(content={'error': 'traceId is required'}, status_code=400) - - task = active_actions.get(trace_id) - if task: - _ = task.cancel() - return JSONResponse(content={'message': 'Action cancelled'}, status_code=200) - else: - return JSONResponse(content={'message': 'Action not found or already completed'}, status_code=404) - except Exception as e: - logger.error(f'Error cancelling action: {e}', exc_info=True) - return JSONResponse( - content={'error': 'An unexpected error occurred while cancelling the action.'}, - status_code=500, - ) - - async def handle_run_action( - request: Request, - ) -> JSONResponse | StreamingResponse: - """Handle the runAction endpoint for executing registered actions. - - Flow: - 1. Reads and validates the request payload - 2. Looks up the requested action - 3. Executes the action with the provided input - 4. Returns the action result as JSON with trace ID - - Args: - request: The Starlette request object. - - Returns: - A JSON or StreamingResponse with the action result, or an error - response. - """ - # Get the action using async resolve. - payload = await request.json() - action = await registry.resolve_action_by_key(payload['key']) - if action is None: - return JSONResponse( - content={'error': f'Action not found: {payload["key"]}'}, - status_code=404, - ) - - # Run the action. - context = payload.get('context', {}) - action_input = payload.get('input') - stream = is_streaming_requested(request) - - # Wrap execution to track the task for cancellation support - task = asyncio.current_task() - - def on_trace_start(trace_id: str, span_id: str) -> None: - if task: - active_actions[trace_id] = task - - handler = run_streaming_action if stream else run_standard_action - - try: - return await handler(action, payload, action_input, context, version, on_trace_start) - except asyncio.CancelledError: - logger.info('Action execution cancelled.') - # Can't really send response if cancelled? Starlette/uvicorn closes connection? - # Just raise. - raise - - async def run_streaming_action( - action: Action, - payload: dict[str, Any], - _action_input: object, - context: dict[str, Any], - version: str, - on_trace_start: Callable[[str, str], None], - ) -> StreamingResponse: - """Handle streaming action execution with early header flushing. - - Uses early header flushing to send X-Genkit-Trace-Id immediately when - the trace starts, enabling the Dev UI to subscribe to SSE for real-time - trace updates. - - Args: - action: The action to execute. - payload: Request payload with input data. - action_input: The input for the action. - context: Execution context. - version: The Genkit version header value. - on_trace_start: Callback for trace start. - - Returns: - A StreamingResponse with JSON chunks containing result or error - events. - """ - # Use a queue to pass chunks from the callback to the generator - chunk_queue: asyncio.Queue[str | None] = asyncio.Queue() - - # Event to signal when trace ID is available - trace_id_event: asyncio.Event = asyncio.Event() - run_trace_id: str | None = None - run_span_id: str | None = None - - def wrapped_on_trace_start(tid: str, sid: str) -> None: - nonlocal run_trace_id, run_span_id - run_trace_id = tid - run_span_id = sid - on_trace_start(tid, sid) - trace_id_event.set() # Signal that trace ID is ready - - async def run_action_task() -> None: - """Run the action and put chunks on the queue.""" - try: - - def send_chunk(chunk: Any) -> None: # noqa: ANN401 - """Callback that puts chunks on the queue.""" - out = dump_json(chunk) - chunk_queue.put_nowait(f'{out}\n') - - output = await action.arun_raw( - raw_input=payload.get('input'), - on_chunk=send_chunk, - context=context, - on_trace_start=wrapped_on_trace_start, - ) - final_response = { - 'result': dump_dict(output.response), - 'telemetry': {'traceId': output.trace_id, 'spanId': output.span_id}, - } - chunk_queue.put_nowait(json.dumps(final_response)) - - except Exception as e: - error_response = get_reflection_json(e).model_dump(by_alias=True) - # Log with exc_info for pretty exception output via rich/structlog - logger.exception('Error streaming action', exc_info=e) - # Error response also should not have trailing newline (final message) - # Wrap error in an 'error' field to match JS SDK format - chunk_queue.put_nowait(json.dumps({'error': error_response})) - # Ensure trace_id_event is set even on error - trace_id_event.set() - - finally: - if not trace_id_event.is_set(): - trace_id_event.set() - # Signal end of stream - chunk_queue.put_nowait(None) - if run_trace_id: - _ = active_actions.pop(run_trace_id, None) - - # Start the action task immediately so trace ID becomes available ASAP - action_task = asyncio.create_task(run_action_task()) - - # Wait for trace ID before returning response - this enables early header flushing - _ = await trace_id_event.wait() - - # Now we have the trace ID, include it in headers - headers = { - 'x-genkit-version': version, - 'Transfer-Encoding': 'chunked', - } - if run_trace_id: - headers['X-Genkit-Trace-Id'] = run_trace_id # pyright: ignore[reportUnreachable] - if run_span_id: - headers['X-Genkit-Span-Id'] = run_span_id # pyright: ignore[reportUnreachable] - - async def stream_generator() -> AsyncGenerator[str, None]: - """Yield chunks from the queue as they arrive.""" - try: - while True: - chunk = await chunk_queue.get() - if chunk is None: - break - yield chunk - finally: - # Cancel task if still running (no-op if already done) - _ = action_task.cancel() - - return StreamingResponse( - stream_generator(), - # Reflection server uses text/plain for streaming (not SSE format) - # to match Go implementation - media_type='text/plain', - headers=headers, - ) - - async def run_standard_action( - action: Action, - payload: dict[str, Any], - _action_input: object, - context: dict[str, Any], - version: str, - on_trace_start: Callable[[str, str], None], - ) -> StreamingResponse: - """Handle standard (non-streaming) action execution with early header flushing. - - Uses StreamingResponse to enable sending the X-Genkit-Trace-Id header - immediately when the trace starts, allowing the Dev UI to subscribe to - the SSE stream for real-time trace updates. - - Args: - action: The action to execute. - payload: Request payload with input data. - action_input: The input for the action. - context: Execution context. - version: The Genkit version header value. - on_trace_start: Callback for trace start. - - Returns: - A StreamingResponse that flushes headers early. - """ - # Event to signal when trace ID is available - trace_id_event: asyncio.Event = asyncio.Event() - run_trace_id: str | None = None - run_span_id: str | None = None - action_result: dict[str, Any] | None = None - action_error: Exception | None = None - - def wrapped_on_trace_start(tid: str, sid: str) -> None: - nonlocal run_trace_id, run_span_id - run_trace_id = tid - run_span_id = sid - on_trace_start(tid, sid) - trace_id_event.set() # Signal that trace ID is ready - - async def run_action_and_get_result() -> None: - nonlocal action_result, action_error - try: - output = await action.arun_raw( - raw_input=payload.get('input'), - context=context, - on_trace_start=wrapped_on_trace_start, - ) - action_result = { - 'result': dump_dict(output.response), - 'telemetry': {'traceId': output.trace_id, 'spanId': output.span_id}, - } - except Exception as e: - action_error = e - finally: - if not trace_id_event.is_set(): - trace_id_event.set() - - # Start the action immediately so trace ID becomes available ASAP - action_task = asyncio.create_task(run_action_and_get_result()) - - # Wait for trace ID before returning response - _ = await trace_id_event.wait() - - # Now return streaming response - headers will include trace ID - async def body_generator() -> AsyncGenerator[bytes, None]: - # Wait for action to complete - await action_task - - if action_error: - error_response = get_reflection_json(action_error).model_dump(by_alias=True) - # Log with exc_info for pretty exception output via rich/structlog - logger.exception('Error executing action', exc_info=action_error) - # Wrap error in an 'error' field to match JS SDK format - yield json.dumps({'error': error_response}).encode('utf-8') - else: - yield json.dumps(action_result).encode('utf-8') - - if run_trace_id: - _ = active_actions.pop(run_trace_id, None) - - headers = { - 'x-genkit-version': version, - 'Transfer-Encoding': 'chunked', - } - if run_trace_id: - headers['X-Genkit-Trace-Id'] = run_trace_id # pyright: ignore[reportUnreachable] - if run_span_id: - headers['X-Genkit-Span-Id'] = run_span_id # pyright: ignore[reportUnreachable] - - return StreamingResponse( - body_generator(), - media_type='application/json', - headers=headers, - ) - - app = Starlette( - routes=[ - Route('/api/__health', handle_health_check, methods=['GET']), - Route('/api/__quitquitquit', handle_terminate, methods=['GET', 'POST']), # Support both for parity - Route('/api/actions', handle_list_actions, methods=['GET']), - Route('/api/values', handle_list_values, methods=['GET']), - Route('/api/envs', handle_list_envs, methods=['GET']), - Route('/api/notify', handle_notify, methods=['POST']), - Route('/api/runAction', handle_run_action, methods=['POST']), - Route('/api/cancelAction', handle_cancel_action, methods=['POST']), - ], - middleware=[ - Middleware( - CORSMiddleware, # type: ignore[arg-type] - allow_origins=['*'], - allow_methods=['*'], - allow_headers=['*'], - expose_headers=['X-Genkit-Trace-Id', 'X-Genkit-Span-Id', 'x-genkit-version'], - ) - ], - on_startup=[on_app_startup] if on_app_startup else [], - on_shutdown=[on_app_shutdown] if on_app_shutdown else [], - ) - app.active_actions = active_actions # type: ignore[attr-defined] - return cast(Application, app) diff --git a/py/packages/genkit/src/genkit/core/schema.py b/py/packages/genkit/src/genkit/core/schema.py deleted file mode 100644 index a5674dca40..0000000000 --- a/py/packages/genkit/src/genkit/core/schema.py +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Functions for working with schema.""" - -from typing import Any - -from pydantic import TypeAdapter - - -def to_json_schema(schema: type | dict[str, Any] | str | None) -> dict[str, Any]: - """Converts a Python type to a JSON schema. - - If the input `schema` is already a dictionary (assumed json schema), it is - returned directly. Otherwise, it is assumed to be a Python type, and a - Pydantic `TypeAdapter` is used to generate the corresponding JSON schema. - - Args: - schema: A Python type or a dictionary representing a JSON schema. - - Returns: - A dictionary representing the JSON schema. - - Examples: - Assuming you have a Pydantic model like this: - - >>> from pydantic import BaseModel - >>> class MyModel(BaseModel): - ... id: int - ... name: str - - You can generate the JSON schema: - - >>> schema = to_json_schema(MyModel) - >>> print(schema) - { - 'properties': {'id': {'title': 'Id', 'type': 'integer'}, 'name': {'title': 'Name', 'type': 'string'}}, - 'required': ['id', 'name'], - 'title': 'MyModel', - 'type': 'object', - } - - If you pass in a dictionary: - - >>> existing_schema = {'type': 'string'} - >>> result = to_json_schema(existing_schema) - >>> print(result) - {'type': 'string'} - """ - if schema is None: - return {'type': 'null'} - if isinstance(schema, dict): - return schema - type_adapter = TypeAdapter(schema) - return type_adapter.json_schema() diff --git a/py/packages/genkit/src/genkit/core/status_types.py b/py/packages/genkit/src/genkit/core/status_types.py deleted file mode 100644 index 2a65451596..0000000000 --- a/py/packages/genkit/src/genkit/core/status_types.py +++ /dev/null @@ -1,249 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Enumeration of response status codes and their corresponding messages.""" - -from enum import IntEnum -from typing import ClassVar, Literal - -from pydantic import BaseModel, ConfigDict, Field - - -class StatusCodes(IntEnum): - """Enumeration of response status codes.""" - - # Not an error; returned on success. - # - # HTTP Mapping: 200 OK - OK = 0 - - # The operation was cancelled, typically by the caller. - # - # HTTP Mapping: 499 Client Closed Request - CANCELLED = 1 - - # Unknown error. For example, this error may be returned when - # a `Status` value received from another address space belongs to - # an error space that is not known in this address space. Also - # errors raised by APIs that do not return enough error information - # may be converted to this error. - # - # HTTP Mapping: 500 Internal Server Error - UNKNOWN = 2 - - # The client specified an invalid argument. Note that this differs - # from `FAILED_PRECONDITION`. `INVALID_ARGUMENT` indicates arguments - # that are problematic regardless of the state of the system - # (e.g., a malformed file name). - # - # HTTP Mapping: 400 Bad Request - INVALID_ARGUMENT = 3 - - # The deadline expired before the operation could complete. For operations - # that change the state of the system, this error may be returned - # even if the operation has completed successfully. For example, a - # successful response from a server could have been delayed long - # enough for the deadline to expire. - # - # HTTP Mapping: 504 Gateway Timeout - DEADLINE_EXCEEDED = 4 - - # Some requested entity (e.g., file or directory) was not found. - # - # Note to server developers: if a request is denied for an entire class - # of users, such as gradual feature rollout or undocumented allowlist, - # `NOT_FOUND` may be used. If a request is denied for some users within - # a class of users, such as user-based access control, `PERMISSION_DENIED` - # must be used. - # - # HTTP Mapping: 404 Not Found - NOT_FOUND = 5 - - # The entity that a client attempted to create (e.g., file or directory) - # already exists. - # - # HTTP Mapping: 409 Conflict - ALREADY_EXISTS = 6 - - # The caller does not have permission to execute the specified - # operation. `PERMISSION_DENIED` must not be used for rejections - # caused by exhausting some resource (use `RESOURCE_EXHAUSTED` - # instead for those errors). `PERMISSION_DENIED` must not be - # used if the caller can not be identified (use `UNAUTHENTICATED` - # instead for those errors). This error code does not imply the - # request is valid or the requested entity exists or satisfies - # other pre-conditions. - # - # HTTP Mapping: 403 Forbidden - PERMISSION_DENIED = 7 - - # The request does not have valid authentication credentials for the - # operation. - # - # HTTP Mapping: 401 Unauthorized - UNAUTHENTICATED = 16 - - # Some resource has been exhausted, perhaps a per-user quota, or - # perhaps the entire file system is out of space. - # - # HTTP Mapping: 429 Too Many Requests - RESOURCE_EXHAUSTED = 8 - - # The operation was rejected because the system is not in a state - # required for the operation's execution. For example, the directory - # to be deleted is non-empty, an rmdir operation is applied to - # a non-directory, etc. - # - # Service implementors can use the following guidelines to decide - # between `FAILED_PRECONDITION`, `ABORTED`, and `UNAVAILABLE`: - # (a) Use `UNAVAILABLE` if the client can retry just the failing call. - # (b) Use `ABORTED` if the client should retry at a higher level. For - # example, when a client-specified test-and-set fails, indicating the - # client should restart a read-modify-write sequence. - # (c) Use `FAILED_PRECONDITION` if the client should not retry until the - # system state has been explicitly fixed. For example, if an "rmdir" - # fails because the directory is non-empty, `FAILED_PRECONDITION` - # should be returned since the client should not retry unless the files - # are deleted from the directory. - # - # HTTP Mapping: 400 Bad Request - FAILED_PRECONDITION = 9 - - # The operation was aborted, typically due to a concurrency issue such as - # a sequencer check failure or transaction abort. - # - # See the guidelines above for deciding between `FAILED_PRECONDITION`, - # `ABORTED`, and `UNAVAILABLE`. - # - # HTTP Mapping: 409 Conflict - ABORTED = 10 - - # The operation was attempted past the valid range. E.g., seeking or - # reading past end-of-file. - # - # Unlike `INVALID_ARGUMENT`, this error indicates a problem that may - # be fixed if the system state changes. For example, a 32-bit file - # system will generate `INVALID_ARGUMENT` if asked to read at an - # offset that is not in the range [0,2^32-1], but it will generate - # `OUT_OF_RANGE` if asked to read from an offset past the current - # file size. - # - # There is a fair bit of overlap between `FAILED_PRECONDITION` and - # `OUT_OF_RANGE`. We recommend using `OUT_OF_RANGE` (the more specific - # error) when it applies so that callers who are iterating through - # a space can easily look for an `OUT_OF_RANGE` error to detect when - # they are done. - # - # HTTP Mapping: 400 Bad Request - OUT_OF_RANGE = 11 - - # The operation is not implemented or is not supported/enabled in this - # service. - # - # HTTP Mapping: 501 Not Implemented - UNIMPLEMENTED = 12 - - # Internal errors. This means that some invariants expected by the - # underlying system have been broken. This error code is reserved - # for serious errors. - # - # HTTP Mapping: 500 Internal Server Error - INTERNAL = 13 - - # The service is currently unavailable. This is most likely a - # transient condition, which can be corrected by retrying with - # a backoff. Note that it is not always safe to retry - # non-idempotent operations. - # - # See the guidelines above for deciding between `FAILED_PRECONDITION`, - # `ABORTED`, and `UNAVAILABLE`. - # - # HTTP Mapping: 503 Service Unavailable - UNAVAILABLE = 14 - - # Unrecoverable data loss or corruption. - # - # HTTP Mapping: 500 Internal Server Error - DATA_LOSS = 15 - - -# Type alias for status names -StatusName = Literal[ - 'OK', - 'CANCELLED', - 'UNKNOWN', - 'INVALID_ARGUMENT', - 'DEADLINE_EXCEEDED', - 'NOT_FOUND', - 'ALREADY_EXISTS', - 'PERMISSION_DENIED', - 'UNAUTHENTICATED', - 'RESOURCE_EXHAUSTED', - 'FAILED_PRECONDITION', - 'ABORTED', - 'OUT_OF_RANGE', - 'UNIMPLEMENTED', - 'INTERNAL', - 'UNAVAILABLE', - 'DATA_LOSS', -] - -# Mapping of status names to HTTP status codes -_STATUS_CODE_MAP: dict[StatusName, int] = { - 'OK': 200, - 'CANCELLED': 499, - 'UNKNOWN': 500, - 'INVALID_ARGUMENT': 400, - 'DEADLINE_EXCEEDED': 504, - 'NOT_FOUND': 404, - 'ALREADY_EXISTS': 409, - 'PERMISSION_DENIED': 403, - 'UNAUTHENTICATED': 401, - 'RESOURCE_EXHAUSTED': 429, - 'FAILED_PRECONDITION': 400, - 'ABORTED': 409, - 'OUT_OF_RANGE': 400, - 'UNIMPLEMENTED': 501, - 'INTERNAL': 500, - 'UNAVAILABLE': 503, - 'DATA_LOSS': 500, -} - - -def http_status_code(status: StatusName) -> int: - """Gets the HTTP status code for a given status name. - - Args: - status: The status name to get the HTTP code for. - - Returns: - The corresponding HTTP status code. - """ - return _STATUS_CODE_MAP[status] - - -class Status(BaseModel): - """Represents a status with a name and optional message.""" - - model_config: ClassVar[ConfigDict] = ConfigDict( - frozen=True, - validate_assignment=True, - extra='forbid', - populate_by_name=True, - ) - - name: StatusName - message: str = Field(default='') diff --git a/py/packages/genkit/src/genkit/core/trace/__init__.py b/py/packages/genkit/src/genkit/core/trace/__init__.py deleted file mode 100644 index 1fe92b7030..0000000000 --- a/py/packages/genkit/src/genkit/core/trace/__init__.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Trace module for defining opentelemetry managing.""" - -from .adjusting_exporter import AdjustingTraceExporter -from .default_exporter import ( - TelemetryServerSpanExporter, - create_span_processor, - init_telemetry_server_exporter, - is_realtime_telemetry_enabled, -) -from .realtime_processor import RealtimeSpanProcessor -from .types import ( - GenkitSpan, -) - -__all__ = [ - 'AdjustingTraceExporter', - 'GenkitSpan', - 'RealtimeSpanProcessor', - 'TelemetryServerSpanExporter', - 'create_span_processor', - 'init_telemetry_server_exporter', - 'is_realtime_telemetry_enabled', -] diff --git a/py/packages/genkit/src/genkit/core/trace/adjusting_exporter.py b/py/packages/genkit/src/genkit/core/trace/adjusting_exporter.py deleted file mode 100644 index 3534a39a0e..0000000000 --- a/py/packages/genkit/src/genkit/core/trace/adjusting_exporter.py +++ /dev/null @@ -1,470 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Adjusting trace exporter for PII redaction and span enhancement. - -This module provides an exporter wrapper that adjusts spans before exporting, -primarily for redacting sensitive input/output data (PII protection) and -augmenting span attributes for cloud observability platforms like Google Cloud Trace. - -Overview: - When exporting traces to cloud services like Google Cloud Trace, you often - want to redact model inputs and outputs to protect potentially sensitive - user data (PII). The AdjustingTraceExporter wraps any SpanExporter and - modifies spans before they're exported. - -Key Features: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Feature β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ PII Redaction β”‚ Replaces genkit:input/output with β”‚ - β”‚ Error Marking β”‚ Adds /http/status_code:599 for GCP red marker β”‚ - β”‚ Label Normalization β”‚ Replaces : with / in attribute keys for GCP β”‚ - β”‚ Failed Span Marking β”‚ Marks failure source with genkit:failedSpan β”‚ - β”‚ Feature Marking β”‚ Marks root spans with genkit:feature β”‚ - β”‚ Model Marking β”‚ Marks model spans with genkit:model β”‚ - β”‚ Configurable Logging β”‚ Optional: keep input/output for debugging β”‚ - β”‚ Error Callbacks β”‚ Custom handling for export errors β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Usage: - ```python - from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter - from genkit.core.trace import AdjustingTraceExporter - - # Wrap the cloud exporter with redaction - base_exporter = CloudTraceSpanExporter() - exporter = AdjustingTraceExporter( - exporter=base_exporter, - log_input_and_output=False, # Redact by default - ) - - # Use with a span processor - processor = BatchSpanProcessor(exporter) - ``` - -Caveats: - - Redaction is shallow - only top-level genkit:input/output are affected - - Setting log_input_and_output=True disables redaction (use with caution) - - The wrapped exporter must implement the standard SpanExporter interface - -See Also: - - JavaScript AdjustingTraceExporter: js/plugins/google-cloud/src/gcpOpenTelemetry.ts -""" - -from __future__ import annotations - -from collections.abc import Callable, Sequence -from typing import Any, ClassVar, cast - -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import Event, ReadableSpan -from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult -from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope -from opentelemetry.trace import Link, SpanContext, SpanKind, Status, StatusCode -from opentelemetry.util.types import Attributes - -from genkit.core._compat import override - - -class RedactedSpan(ReadableSpan): - """A span wrapper that redacts sensitive attributes. - - This class wraps a ReadableSpan and overrides the attributes property - to return redacted values for sensitive fields like genkit:input and - genkit:output. - """ - - def __init__(self, span: ReadableSpan, redacted_attributes: dict[str, Any]) -> None: - """Initialize a RedactedSpan. - - Args: - span: The original span to wrap. - redacted_attributes: The attributes with redacted values. - """ - self._span = span - self._redacted_attributes = redacted_attributes - - @property - @override - def name(self) -> str: - """Return the span name.""" - return self._span.name - - @property - @override - def context(self) -> SpanContext: - """Return the span context.""" - return cast(SpanContext, self._span.context) - - @override - def get_span_context(self) -> SpanContext: - """Return the span context.""" - return cast(SpanContext, self._span.get_span_context()) - - @property - @override - def parent(self) -> SpanContext | None: - """Return the parent span context.""" - return self._span.parent - - @property - @override - def start_time(self) -> int | None: - """Return the span start time.""" - return self._span.start_time - - @property - @override - def end_time(self) -> int | None: - """Return the span end time.""" - return self._span.end_time - - @property - @override - def status(self) -> Status: - """Return the span status.""" - return self._span.status - - @property - @override - def attributes(self) -> Attributes: - """Return the redacted attributes.""" - return self._redacted_attributes - - @property - @override - def events(self) -> Sequence[Event]: - """Return the span events.""" - return self._span.events - - @property - @override - def links(self) -> Sequence[Link]: - """Return the span links.""" - return self._span.links - - @property - @override - def kind(self) -> SpanKind: - """Return the span kind.""" - return self._span.kind - - @property - @override - def resource(self) -> Resource: - """Return the span resource.""" - return self._span.resource - - @property - @override - def instrumentation_info(self) -> InstrumentationInfo | None: - """Return the instrumentation info.""" - # pyrefly: ignore[deprecated] - Required override for ReadableSpan interface compatibility - return self._span.instrumentation_info - - @property - @override - def instrumentation_scope(self) -> InstrumentationScope | None: - """Return the instrumentation scope.""" - return self._span.instrumentation_scope - - @property - @override - def dropped_attributes(self) -> int: - """Return the number of dropped attributes from the wrapped span. - - The base ReadableSpan implementation accesses ``self._attributes`` - which is set by ``ReadableSpan.__init__``. Since RedactedSpan - intentionally skips ``super().__init__()`` (to avoid duplicating - span state), the private ``_attributes`` field does not exist. - This override delegates to the wrapped span instead. - """ - return self._span.dropped_attributes - - @property - @override - def dropped_events(self) -> int: - """Return the number of dropped events from the wrapped span. - - Delegates to the wrapped span for the same reason as - :pyattr:`dropped_attributes`. - """ - return self._span.dropped_events - - @property - @override - def dropped_links(self) -> int: - """Return the number of dropped links from the wrapped span. - - Delegates to the wrapped span for the same reason as - :pyattr:`dropped_attributes`. - """ - return self._span.dropped_links - - -class AdjustingTraceExporter(SpanExporter): - """Adjusts spans before exporting for PII redaction and enhancement. - - This exporter wraps another SpanExporter and modifies spans before they - are exported. The primary use case is redacting model input/output to - protect user privacy when exporting to cloud observability platforms. - - Attributes: - exporter: The wrapped SpanExporter. - log_input_and_output: If True, don't redact input/output. - project_id: Optional project ID for cloud-specific features. - error_handler: Optional callback for export errors. - - Example: - ```python - from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter - from genkit.core.trace import AdjustingTraceExporter - - exporter = AdjustingTraceExporter( - exporter=CloudTraceSpanExporter(), - log_input_and_output=False, - ) - ``` - """ - - REDACTED_VALUE: ClassVar[str] = '' - - def __init__( - self, - exporter: SpanExporter, - log_input_and_output: bool = False, - project_id: str | None = None, - error_handler: Callable[[Exception], None] | None = None, - ) -> None: - """Initialize the AdjustingTraceExporter. - - Args: - exporter: The underlying SpanExporter to wrap. - log_input_and_output: If True, preserve input/output in spans. - Defaults to False (redact for privacy). - project_id: Optional project ID for cloud-specific features. - error_handler: Optional callback invoked when export errors occur. - """ - self._exporter: SpanExporter = exporter - self._log_input_and_output: bool = log_input_and_output - self._project_id: str | None = project_id - self._error_handler: Callable[[Exception], None] | None = error_handler - - @override - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: - """Export spans after adjusting them. - - Applies transformations to each span (redaction, marking, etc.) - before passing them to the underlying exporter. - - Args: - spans: The spans to export. - - Returns: - The result from the underlying exporter. - """ - adjusted_spans = [self._adjust(span) for span in spans] - - try: - result = self._exporter.export(adjusted_spans) - return result - except Exception as e: - if self._error_handler: - self._error_handler(e) - raise - - @override - def shutdown(self) -> None: - """Shut down the underlying exporter.""" - self._exporter.shutdown() - - @override - def force_flush(self, timeout_millis: int = 30000) -> bool: - """Force the underlying exporter to flush. - - Args: - timeout_millis: Maximum time to wait for flush. - - Returns: - True if flush succeeded. - """ - if hasattr(self._exporter, 'force_flush'): - return self._exporter.force_flush(timeout_millis) - return True - - def _adjust(self, span: ReadableSpan) -> ReadableSpan: - """Apply all adjustments to a span. - - This method applies the same transformations as the JavaScript - implementation in gcpOpenTelemetry.ts: - 1. Redact input/output (if not logging) - 2. Mark error spans with HTTP status code for GCP - 3. Mark failed spans with failure source info - 4. Mark root spans with genkit:feature - 5. Mark model spans with genkit:model - 6. Normalize attribute labels (: -> /) - - Args: - span: The span to adjust. - - Returns: - The adjusted span (possibly a RedactedSpan wrapper). - """ - span = self._redact_input_output(span) - span = self._mark_error_span_as_error(span) - span = self._mark_failed_span(span) - span = self._mark_genkit_feature(span) - span = self._mark_genkit_model(span) - span = self._normalize_labels(span) - return span - - def _redact_input_output(self, span: ReadableSpan) -> ReadableSpan: - """Redact genkit:input and genkit:output attributes. - - If log_input_and_output is True, the span is returned unchanged. - Otherwise, these sensitive fields are replaced with ''. - - Args: - span: The span to potentially redact. - - Returns: - The original span or a RedactedSpan wrapper. - """ - if self._log_input_and_output: - return span - - attributes = dict(span.attributes) if span.attributes else {} - has_input = 'genkit:input' in attributes - has_output = 'genkit:output' in attributes - - if not has_input and not has_output: - return span - - # Create redacted attributes - redacted = {**attributes} - if has_input: - redacted['genkit:input'] = self.REDACTED_VALUE - if has_output: - redacted['genkit:output'] = self.REDACTED_VALUE - - return RedactedSpan(span, redacted) - - def _mark_error_span_as_error(self, span: ReadableSpan) -> ReadableSpan: - """Mark error spans with HTTP status code for GCP Trace display. - - This is a workaround for GCP Trace to mark a span with a red - exclamation mark indicating that it is an error. GCP requires - an HTTP status code to show the error indicator. - - Args: - span: The span to potentially mark. - - Returns: - The span with /http/status_code: 599 if it's an error span. - """ - if not span.status or span.status.status_code != StatusCode.ERROR: - return span - - attributes = dict(span.attributes) if span.attributes else {} - attributes['/http/status_code'] = '599' - return RedactedSpan(span, attributes) - - def _mark_failed_span(self, span: ReadableSpan) -> ReadableSpan: - """Mark spans that are the source of a failure. - - Adds genkit:failedSpan and genkit:failedPath attributes to spans - that have genkit:isFailureSource set. - - Args: - span: The span to potentially mark. - - Returns: - The span with failure markers if applicable. - """ - attributes = dict(span.attributes) if span.attributes else {} - - if not attributes.get('genkit:isFailureSource'): - return span - - attributes['genkit:failedSpan'] = attributes.get('genkit:name', '') - attributes['genkit:failedPath'] = attributes.get('genkit:path', '') - return RedactedSpan(span, attributes) - - def _mark_genkit_feature(self, span: ReadableSpan) -> ReadableSpan: - """Mark root spans with the genkit:feature attribute. - - This helps identify the top-level feature being executed. - - Args: - span: The span to potentially mark. - - Returns: - The span with genkit:feature if it's a root span. - """ - attributes = dict(span.attributes) if span.attributes else {} - - is_root = attributes.get('genkit:isRoot') - name = attributes.get('genkit:name') - - if not is_root or not name: - return span - - attributes['genkit:feature'] = name - return RedactedSpan(span, attributes) - - def _mark_genkit_model(self, span: ReadableSpan) -> ReadableSpan: - """Mark model spans with the genkit:model attribute. - - This helps identify which model was used in a span. - - Args: - span: The span to potentially mark. - - Returns: - The span with genkit:model if it's a model action. - """ - attributes = dict(span.attributes) if span.attributes else {} - - subtype = attributes.get('genkit:metadata:subtype') - name = attributes.get('genkit:name') - - if subtype != 'model' or not name: - return span - - attributes['genkit:model'] = name - return RedactedSpan(span, attributes) - - def _normalize_labels(self, span: ReadableSpan) -> ReadableSpan: - """Normalize attribute labels by replacing : with /. - - GCP Cloud Trace has specific requirements for label keys. - This ensures compatibility by replacing colons with slashes. - - Args: - span: The span with attributes to normalize. - - Returns: - The span with normalized attribute keys. - """ - attributes = dict(span.attributes) if span.attributes else {} - - # Replace : with / in all attribute keys - normalized: dict[str, Any] = {} - for key, value in attributes.items(): - normalized[key.replace(':', '/')] = value - - return RedactedSpan(span, normalized) diff --git a/py/packages/genkit/src/genkit/core/trace/default_exporter.py b/py/packages/genkit/src/genkit/core/trace/default_exporter.py deleted file mode 100644 index f29be5d711..0000000000 --- a/py/packages/genkit/src/genkit/core/trace/default_exporter.py +++ /dev/null @@ -1,237 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Telemetry and tracing default exporter for the Genkit framework. - -This module provides functionality for collecting and exporting telemetry data -from Genkit operations. It uses OpenTelemetry for tracing and exports span -data to a telemetry server for monitoring and debugging purposes. - -The module includes: - - A custom span exporter for sending trace data to a telemetry server - - Utility functions for converting and formatting trace attributes -""" - -from __future__ import annotations - -import os -import sys -from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, cast -from urllib.parse import urljoin - -import httpx -from opentelemetry import trace as trace_api -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.trace.export import ( - BatchSpanProcessor, - SimpleSpanProcessor, - SpanExporter, - SpanExportResult, -) - -from genkit.core._compat import override -from genkit.core.environment import is_dev_environment -from genkit.core.logging import get_logger - -from .realtime_processor import RealtimeSpanProcessor - -if TYPE_CHECKING: - from opentelemetry.sdk.trace import SpanProcessor - from opentelemetry.trace import SpanContext - -ATTR_PREFIX = 'genkit' -logger = get_logger(__name__) - - -def extract_span_data(span: ReadableSpan) -> dict[str, Any]: - """Extract span data from a ReadableSpan object. - - This function extracts the span data from a ReadableSpan object and returns - a dictionary containing the span data. - """ - # Format trace_id and span_id as hex strings (OpenTelemetry standard format) - context = cast('SpanContext', span.context) - trace_id_hex = format(context.trace_id, '032x') - span_id_hex = format(context.span_id, '016x') - parent_span_id_hex = format(span.parent.span_id, '016x') if span.parent else None - - span_data: dict[str, Any] = {'traceId': trace_id_hex, 'spans': {}} - start_time = (span.start_time / 1000000) if span.start_time is not None else 0 - end_time = (span.end_time / 1000000) if span.end_time is not None else 0 - - span_data['spans'][span_id_hex] = { - 'spanId': span_id_hex, - 'traceId': trace_id_hex, - 'startTime': start_time, - 'endTime': end_time, - 'attributes': {**(span.attributes or {})}, - 'displayName': span.name, - # "links": span.links, - 'spanKind': trace_api.SpanKind(span.kind).name, - 'parentSpanId': parent_span_id_hex, - 'status': ( - { - 'code': trace_api.StatusCode(span.status.status_code).value, - 'description': span.status.description, - } - if span.status - else None - ), - 'instrumentationLibrary': { - 'name': 'genkit-tracer', - 'version': 'v1', - }, - } - if not span_data['spans'][span_id_hex]['parentSpanId']: - del span_data['spans'][span_id_hex]['parentSpanId'] - - if not span.parent: - span_data['displayName'] = span.name - span_data['startTime'] = start_time - span_data['endTime'] = end_time - - return span_data - - -class TelemetryServerSpanExporter(SpanExporter): - """Exports spans to a Genkit telemetry server. - - This exporter sends span data in a specific JSON format to a telemetry server, - typically running locally during development, for visualization and debugging. - - Attributes: - telemetry_server_url: The URL of the telemetry server endpoint. - """ - - def __init__(self, telemetry_server_url: str, telemetry_server_endpoint: str | None = None) -> None: - """Initializes the TelemetryServerSpanExporter. - - Args: - telemetry_server_url: The URL of the telemetry server. - telemetry_server_endpoint (optional): The telemetry server's trace endpoint. - """ - self.telemetry_server_url: str = telemetry_server_url - if telemetry_server_endpoint is None: - self.telemetry_server_endpoint: str = '/api/traces' - else: - self.telemetry_server_endpoint = telemetry_server_endpoint - - @override - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: - """Exports a sequence of ReadableSpans to the configured telemetry server. - - Iterates through the provided spans, extracts relevant data using - `extract_span_data`, converts it to JSON, and sends it via an HTTP POST - request to the `telemetry_server_url`. - - Args: - spans: A sequence of OpenTelemetry ReadableSpan objects to export. - - Returns: - SpanExportResult.SUCCESS upon successful processing (does not guarantee - server-side success). - """ - with httpx.Client() as client: - for span in spans: - _ = client.post( - urljoin(self.telemetry_server_url, self.telemetry_server_endpoint), - json=extract_span_data(span), - headers={ - 'Content-Type': 'application/json', - 'Accept': 'application/json', - }, - ) - - _ = sys.stdout.flush() - - return SpanExportResult.SUCCESS - - @override - def force_flush(self, timeout_millis: int = 30000) -> bool: - """Forces the exporter to flush any buffered spans. - - Since this exporter sends spans immediately in the `export` method, - this method currently does nothing but return True. - - Args: - timeout_millis: The maximum time in milliseconds to wait for the flush. - This parameter is ignored in the current implementation. - - Returns: - True, indicating the flush operation is considered complete. - """ - return True - - -def init_telemetry_server_exporter() -> SpanExporter | None: - """Initializes tracing with a provider and optional exporter. - - Returns: - A SpanExporter configured for the telemetry server, or None if - GENKIT_TELEMETRY_SERVER is not set. - - Environment Variables: - GENKIT_TELEMETRY_SERVER: URL of the telemetry server. - GENKIT_ENABLE_REALTIME_TELEMETRY: Set to 'true' to enable realtime - span processing (exports spans on start and end). - """ - telemetry_server_url = os.environ.get('GENKIT_TELEMETRY_SERVER') - processor = None - - if telemetry_server_url: - processor = TelemetryServerSpanExporter( - telemetry_server_url=telemetry_server_url, - ) - else: - logger.warn( - 'GENKIT_TELEMETRY_SERVER is not set. If running with `genkit start`, make sure `genkit-cli` is up to date.' - ) - - return processor - - -def is_realtime_telemetry_enabled() -> bool: - """Check if realtime telemetry is enabled. - - Returns: - True if GENKIT_ENABLE_REALTIME_TELEMETRY is set to 'true'. - """ - return os.environ.get('GENKIT_ENABLE_REALTIME_TELEMETRY', '').lower() == 'true' - - -def create_span_processor(exporter: SpanExporter) -> SpanProcessor: - """Create an appropriate SpanProcessor for the given exporter. - - Uses RealtimeSpanProcessor when in dev mode AND GENKIT_ENABLE_REALTIME_TELEMETRY - is set to 'true'. Otherwise uses SimpleSpanProcessor for dev or BatchSpanProcessor - for production. - - This matches the JavaScript implementation in node-telemetry-provider.ts. - - Args: - exporter: The SpanExporter to wrap. - - Returns: - A SpanProcessor configured for the current environment. - """ - # Match JS: RealtimeSpanProcessor requires BOTH dev mode AND env var - if is_dev_environment() and is_realtime_telemetry_enabled(): - return RealtimeSpanProcessor(exporter) - elif is_dev_environment(): - return SimpleSpanProcessor(exporter) - else: - return BatchSpanProcessor(exporter) diff --git a/py/packages/genkit/src/genkit/core/trace/realtime_processor.py b/py/packages/genkit/src/genkit/core/trace/realtime_processor.py deleted file mode 100644 index c8fcf00ea4..0000000000 --- a/py/packages/genkit/src/genkit/core/trace/realtime_processor.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Realtime span processor for live trace visualization. - -This module provides a SpanProcessor that exports spans both when they start -and when they end, enabling real-time trace visualization in the DevUI. - -Overview: - Standard OpenTelemetry processors (SimpleSpanProcessor, BatchSpanProcessor) - only export spans when they complete. This is efficient but means the DevUI - cannot show in-progress operations. - - RealtimeSpanProcessor exports spans immediately on start (without endTime), - then exports again when the span completes with full data. This enables: - - - Live progress visualization in DevUI - - Real-time debugging during development - - Immediate feedback on long-running operations - -Key Concepts: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Processor Type β”‚ Export on Start β”‚ Export on End β”‚ Use Case β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ SimpleSpanProcessor β”‚ No β”‚ Yes β”‚ Dev testing β”‚ - β”‚ BatchSpanProcessor β”‚ No β”‚ Yes (batched) β”‚ Production β”‚ - β”‚ RealtimeSpanProcessor β”‚ Yes β”‚ Yes β”‚ DevUI live β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Usage: - Enable realtime telemetry by setting the environment variable: - - ```bash - export GENKIT_ENABLE_REALTIME_TELEMETRY=true - genkit start -- python main.py - ``` - - Or programmatically configure the processor: - - ```python - from opentelemetry.sdk.trace import TracerProvider - from genkit.core.trace import RealtimeSpanProcessor, TelemetryServerSpanExporter - - exporter = TelemetryServerSpanExporter(telemetry_server_url='http://localhost:4000') - processor = RealtimeSpanProcessor(exporter) - - provider = TracerProvider() - provider.add_span_processor(processor) - ``` - -Caveats: - - Doubles network traffic (each span exported twice) - - Not recommended for production use - - Should only be used with GENKIT_ENABLE_REALTIME_TELEMETRY=true - -See Also: - - JavaScript RealtimeSpanProcessor: js/core/src/tracing/realtime-span-processor.ts -""" - -from opentelemetry.context import Context -from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor -from opentelemetry.sdk.trace.export import SpanExporter - -from genkit.core._compat import override -from genkit.core.logging import get_logger - -logger = get_logger(__name__) - - -class RealtimeSpanProcessor(SpanProcessor): - """Exports spans both when they start and when they end. - - This processor enables real-time trace visualization by exporting spans - immediately when they start (without endTime), then again when they - complete with full timing and status data. - - Attributes: - exporter: The SpanExporter to use for exporting span data. - - Example: - ```python - from opentelemetry.sdk.trace import TracerProvider - from genkit.core.trace import RealtimeSpanProcessor - - exporter = TelemetryServerSpanExporter(url='http://localhost:4000') - processor = RealtimeSpanProcessor(exporter) - - provider = TracerProvider() - provider.add_span_processor(processor) - ``` - """ - - def __init__(self, exporter: SpanExporter) -> None: - """Initialize the RealtimeSpanProcessor. - - Args: - exporter: The SpanExporter to use for exporting spans. - """ - self._exporter: SpanExporter = exporter - - @override - def on_start(self, span: Span, parent_context: Context | None = None) -> None: - """Called when a span is started. - - Exports the span immediately for real-time updates. The span will - not have endTime set yet, allowing the DevUI to show it as in-progress. - - Args: - span: The span that was just started. - parent_context: The parent context (unused). - """ - # Export the span immediately (it won't have endTime yet). - # Catch all exceptions so a failing exporter (e.g. Jaeger not - # running) never propagates into the application call stack. - # Without this guard, a ConnectionError here bubbles through - # the OTel gRPC server interceptor and kills the actual RPC. - try: - self._exporter.export([span]) - except ConnectionError: - logger.debug( - 'RealtimeSpanProcessor: export failed on_start (collector unreachable)', - exc_info=True, - ) - except Exception: # noqa: BLE001 β€” must never crash the caller - logger.warning( - 'RealtimeSpanProcessor: unexpected error during export on_start', - exc_info=True, - ) - - @override - def on_end(self, span: ReadableSpan) -> None: - """Called when a span ends. - - Exports the completed span with full timing and status data. - - Args: - span: The span that just ended. - """ - try: - self._exporter.export([span]) - except ConnectionError: - logger.debug( - 'RealtimeSpanProcessor: export failed on_end (collector unreachable)', - exc_info=True, - ) - except Exception: # noqa: BLE001 β€” must never crash the caller - logger.warning( - 'RealtimeSpanProcessor: unexpected error during export on_end', - exc_info=True, - ) - - @override - def force_flush(self, timeout_millis: int = 30000) -> bool: - """Force the exporter to flush any buffered spans. - - Args: - timeout_millis: Maximum time to wait for flush in milliseconds. - - Returns: - True if flush succeeded, False otherwise. - """ - if hasattr(self._exporter, 'force_flush'): - return self._exporter.force_flush(timeout_millis) - return True - - @override - def shutdown(self) -> None: - """Shut down the processor and exporter.""" - self._exporter.shutdown() diff --git a/py/packages/genkit/src/genkit/core/trace/types.py b/py/packages/genkit/src/genkit/core/trace/types.py deleted file mode 100644 index db3c94a7c0..0000000000 --- a/py/packages/genkit/src/genkit/core/trace/types.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Telemetry and tracing types for the Genkit framework. - -This module defines the core tracing types for the Genkit framework. - -Genkit flows are instrumented to trace key data points, primarily the inputs -sent to and the outputs received from Language Learning Models (LLMs) or -other models. Outputs may be streamed as chunks. This module establishes -how these interactions are recorded and exported for observability. - -Key Features: - -- **GenkitSpan Class:** Provides a core class, `GenkitSpan`, for representing - and exporting telemetry data related to Genkit operations. This class - is intended for internal use by the Genkit framework. -""" - -import json -from collections.abc import Mapping -from typing import Any - -from opentelemetry import trace as trace_api -from opentelemetry.util import types -from pydantic import BaseModel - -from genkit.core.logging import get_logger - -ATTR_PREFIX = 'genkit' - -logger = get_logger(__name__) - - -class GenkitSpan: - """Light wrapper for Span, specific to Genkit.""" - - is_root: bool - _span: trace_api.Span - - def __init__(self, span: trace_api.Span, labels: dict[str, str] | None = None) -> None: - """Create GenkitSpan.""" - self._span = span - parent = getattr(span, 'parent', None) - self.is_root = False - if parent is None: - self.is_root = True - if labels is not None: - self.set_attributes(labels) - - def __getattr__(self, name: str) -> Any: # noqa: ANN401 - """Passthrough for all OpenTelemetry Span attributes.""" - return getattr(self._span, name) - - def set_genkit_attribute(self, key: str, value: types.AttributeValue) -> None: - """Set Genkit specific attribute, with the `genkit` prefix.""" - if key == 'metadata' and isinstance(value, dict) and value: - for meta_key, meta_value in value.items(): - self._span.set_attribute(f'{ATTR_PREFIX}:metadata:{meta_key}', str(meta_value)) - elif isinstance(value, dict): - self._span.set_attribute(f'{ATTR_PREFIX}:{key}', json.dumps(value)) - else: - self._span.set_attribute(f'{ATTR_PREFIX}:{key}', str(value)) - - def set_genkit_attributes(self, attributes: Mapping[str, types.AttributeValue]) -> None: - """Set Genkit specific attributes, with the `genkit` prefix.""" - for key, value in attributes.items(): - self.set_genkit_attribute(key, value) - - @property - def span_id(self) -> str: - """Returns the span_id.""" - return str(self._span.get_span_context().span_id) - - @property - def trace_id(self) -> str: - """Returns the trace_id.""" - return str(self._span.get_span_context().trace_id) - - def set_input(self, input: object) -> None: - """Set Genkit Span input, visible in the trace viewer.""" - value = None - if isinstance(input, BaseModel): - value = input.model_dump_json(by_alias=True, exclude_none=True) - else: - value = json.dumps(input) - self.set_genkit_attribute('input', value) - - def set_output(self, output: object) -> None: - """Set Genkit Span output, visible in the trace viewer.""" - value = None - if isinstance(output, BaseModel): - value = output.model_dump_json(by_alias=True, exclude_none=True) - else: - value = json.dumps(output) - self.set_genkit_attribute('output', value) diff --git a/py/packages/genkit/src/genkit/embedder/__init__.py b/py/packages/genkit/src/genkit/embedder/__init__.py new file mode 100644 index 0000000000..c7f7f2e37a --- /dev/null +++ b/py/packages/genkit/src/genkit/embedder/__init__.py @@ -0,0 +1,57 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Embedder namespace module for Genkit. + +This module provides embedder-related types and utilities for plugin authors +and advanced users who need access to the embedder protocol types. + +Example: + from genkit.embedder import ( + EmbedRequest, + EmbedResponse, + embedder_action_metadata, + EmbedderRef, + ) +""" + +from genkit._ai._embedding import ( + EmbedderOptions, + EmbedderRef, + EmbedderSupports, + create_embedder_ref as embedder_ref, + embedder_action_metadata, +) +from genkit._core._typing import ( + Embedding, + EmbedRequest, + EmbedResponse, +) + +__all__ = [ + # Request/Response types + 'EmbedRequest', + 'EmbedResponse', + 'Embedding', + # Factory functions and metadata + 'embedder_action_metadata', + 'embedder_ref', + # Reference types + 'EmbedderRef', + # Options and capabilities + 'EmbedderSupports', + 'EmbedderOptions', +] diff --git a/py/packages/genkit/src/genkit/evaluator/__init__.py b/py/packages/genkit/src/genkit/evaluator/__init__.py new file mode 100644 index 0000000000..991ea262d3 --- /dev/null +++ b/py/packages/genkit/src/genkit/evaluator/__init__.py @@ -0,0 +1,64 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Evaluator namespace module for Genkit. + +This module provides evaluator-related types and utilities for plugin authors +and advanced users who need access to the evaluator protocol types. + +Example: + from genkit.evaluator import ( + EvalRequest, + EvalResponse, + evaluator_action_metadata, + ) +""" + +from genkit._ai._evaluator import ( + EvaluatorRef, + evaluator_action_metadata, + evaluator_ref, +) +from genkit._core._typing import ( + BaseDataPoint, + BaseEvalDataPoint, + Details, + EvalFnResponse, + EvalRequest, + EvalResponse, + EvalStatusEnum, + Score, +) + +__all__ = [ + # Request/Response types + 'EvalRequest', + 'EvalResponse', + 'EvalFnResponse', + # Score types + 'Score', + 'Details', + # Data point types + 'BaseEvalDataPoint', + 'BaseDataPoint', + # Status + 'EvalStatusEnum', + # Factory functions and metadata + 'evaluator_action_metadata', + 'evaluator_ref', + # Reference types + 'EvaluatorRef', +] diff --git a/py/packages/genkit/src/genkit/lang/__init__.py b/py/packages/genkit/src/genkit/lang/__init__.py deleted file mode 100644 index 7cc63404cd..0000000000 --- a/py/packages/genkit/src/genkit/lang/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Language utilities.""" diff --git a/py/packages/genkit/src/genkit/model/__init__.py b/py/packages/genkit/src/genkit/model/__init__.py new file mode 100644 index 0000000000..fad8dbef7a --- /dev/null +++ b/py/packages/genkit/src/genkit/model/__init__.py @@ -0,0 +1,83 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Model protocol types for plugin authors.""" + +from genkit._ai._model import ( + ModelConfig, + model_action_metadata, + model_ref, +) +from genkit._core._background import BackgroundAction +from genkit._core._model import ( + Message, + ModelRef, + ModelRequest, + ModelResponse, + ModelResponseChunk, + ModelUsage, + get_basic_usage_stats, +) +from genkit._core._typing import ( + Candidate, + Constrained, + Error, + FinishReason, + GenerateActionOptions, + ModelInfo, + Operation, + Stage, + Supports, + ToolDefinition, + ToolRequest, + ToolResponse, +) + +__all__ = [ + # Request/Response types + 'BackgroundAction', + 'ModelRequest', + 'ModelResponse', + 'ModelResponseChunk', + # Usage and metadata + 'ModelUsage', + 'Candidate', + 'FinishReason', + 'GenerateActionOptions', + # Error and operation + 'Error', + 'Operation', + # Tool types + 'ToolRequest', + 'ToolDefinition', + 'ToolResponse', + # Model info + 'ModelInfo', + 'Supports', + 'Constrained', + 'Stage', + # Factory functions and metadata + 'model_action_metadata', + 'model_ref', + # Reference types + 'ModelRef', + # Config + 'ModelConfig', + # Message + 'Message', + # Usage + 'get_basic_usage_stats', +] diff --git a/py/packages/genkit/src/genkit/model_types.py b/py/packages/genkit/src/genkit/model_types.py deleted file mode 100644 index 2bf9cad3b3..0000000000 --- a/py/packages/genkit/src/genkit/model_types.py +++ /dev/null @@ -1,76 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Python SDK-specific model type extensions. - -Extends auto-generated core model types with Python SDK runtime fields and -shared helpers for plugin authors. -""" - -from collections.abc import Mapping -from typing import cast - -from pydantic import Field - -from genkit.core.typing import GenerationCommonConfig as CoreGenerationCommonConfig - - -class GenerationCommonConfig(CoreGenerationCommonConfig): - """Common generation config with Python SDK runtime extensions.""" - - api_key: str | None = Field( - default=None, - alias='apiKey', - description='API Key to use for the model call, overrides API key provided in plugin config.', - ) - - -def get_request_api_key(config: Mapping[str, object] | GenerationCommonConfig | object | None) -> str | None: - """Extract a request-scoped API key from config. - - Supports both typed config objects and dict payloads with either snake_case - or camelCase keys. - """ - if config is None: - return None - - if isinstance(config, GenerationCommonConfig): - return config.api_key - - if isinstance(config, Mapping): - config_mapping = cast(Mapping[str, object], config) - api_key = config_mapping.get('api_key') or config_mapping.get('apiKey') - if isinstance(api_key, str) and api_key: - return api_key - else: - # Defensive fallback for plugin-specific config classes that inherit from - # GenerationCommonConfig or expose an api_key attribute. - api_key_attr = getattr(config, 'api_key', None) - if isinstance(api_key_attr, str) and api_key_attr: - return api_key_attr - - return None - - -def get_effective_api_key( - config: Mapping[str, object] | GenerationCommonConfig | object | None, - plugin_api_key: str | None, -) -> str | None: - """Resolve effective API key using request-over-plugin precedence.""" - return get_request_api_key(config) or plugin_api_key - - -__all__ = ['GenerationCommonConfig', 'get_request_api_key', 'get_effective_api_key'] diff --git a/py/packages/genkit/src/genkit/plugin_api/__init__.py b/py/packages/genkit/src/genkit/plugin_api/__init__.py new file mode 100644 index 0000000000..e80d49d934 --- /dev/null +++ b/py/packages/genkit/src/genkit/plugin_api/__init__.py @@ -0,0 +1,97 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Framework primitives for plugin authors.""" + +# Base class and framework primitives +from genkit._core._action import Action, ActionKind, ActionMetadata, ActionRunContext +from genkit._core._constants import GENKIT_CLIENT_HEADER, GENKIT_VERSION +from genkit._core._context import ContextProvider, RequestData +from genkit._core._environment import is_dev_environment +from genkit._core._error import GenkitError, StatusCodes, StatusName, get_callable_json +from genkit._core._http_client import get_cached_client +from genkit._core._loop_cache import _loop_local_client as loop_local_client +from genkit._core._plugin import Plugin +from genkit._core._schema import to_json_schema +from genkit._core._trace._adjusting_exporter import AdjustingTraceExporter, RedactedSpan +from genkit._core._trace._path import to_display_path +from genkit._core._tracing import add_custom_exporter, tracer + +# Embedder domain re-exports +from genkit.embedder import ( + EmbedderRef, + embedder_action_metadata, + embedder_ref, +) + +# Evaluator domain re-exports +from genkit.evaluator import ( + EvaluatorRef, + evaluator_action_metadata, + evaluator_ref, +) + +# Model domain re-exports +from genkit.model import ( + ModelRef, + model_action_metadata, + model_ref, +) + +__all__ = [ + # Base class and framework primitives + 'Plugin', + 'Action', + 'ActionMetadata', + 'ActionKind', + 'ActionRunContext', + 'StatusCodes', + 'StatusName', + 'GenkitError', + # HTTP / version stamping + 'GENKIT_CLIENT_HEADER', + 'GENKIT_VERSION', + # Loop-local caching + 'loop_local_client', + # Tracing + 'tracer', + 'add_custom_exporter', + 'AdjustingTraceExporter', + 'RedactedSpan', + 'to_display_path', + # Schema utilities + 'to_json_schema', + # HTTP client + 'get_cached_client', + # Error serialization + 'get_callable_json', + # Environment detection + 'is_dev_environment', + # Model domain + 'model_action_metadata', + 'model_ref', + 'ModelRef', + # Embedder domain + 'embedder_action_metadata', + 'embedder_ref', + 'EmbedderRef', + # Evaluator domain + 'evaluator_action_metadata', + 'evaluator_ref', + 'EvaluatorRef', + 'ContextProvider', + 'RequestData', +] diff --git a/py/packages/genkit/src/genkit/plugins/.gitignore b/py/packages/genkit/src/genkit/plugins/.gitignore deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/constant.py b/py/packages/genkit/src/genkit/plugins/__init__.py similarity index 63% rename from py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/constant.py rename to py/packages/genkit/src/genkit/plugins/__init__.py index dd0156e303..faac15e8d3 100644 --- a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/constant.py +++ b/py/packages/genkit/src/genkit/plugins/__init__.py @@ -14,16 +14,14 @@ # # SPDX-License-Identifier: Apache-2.0 +"""Namespace package for Genkit plugins. -"""Constants for dev-local-vectorstore.""" +This is a namespace package that allows plugins to be discovered from +multiple installed packages. Each plugin can be imported as: -from pydantic import BaseModel + from genkit.plugins. import -from genkit.types import DocumentData, Embedding - - -class DbValue(BaseModel): - """Value stored in the local filestore.""" - - doc: DocumentData - embedding: Embedding +For example: + from genkit.plugins.google_genai import GoogleGenai + from genkit.plugins.anthropic import Anthropic +""" diff --git a/py/packages/genkit/src/genkit/testing.py b/py/packages/genkit/src/genkit/testing.py deleted file mode 100644 index 1a06c28156..0000000000 --- a/py/packages/genkit/src/genkit/testing.py +++ /dev/null @@ -1,724 +0,0 @@ -#!/usr/bin/env python3 -# -# Copyright 2025 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -"""Testing utilities for Genkit applications. - -This module provides mock models, test utilities, and a model test suite for -testing Genkit applications without making actual API calls to AI providers. - -Key Components -============== - -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ Testing Components β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ Component β”‚ Purpose β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ EchoModel β”‚ Echoes input back - verify request formatting β”‚ -β”‚ ProgrammableModel β”‚ Returns configurable responses - test scenarios β”‚ -β”‚ StaticResponseModel β”‚ Always returns same response - simple tests β”‚ -β”‚ test_models() β”‚ Run standard test suite against models β”‚ -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - ```python - from genkit.ai import Genkit - from genkit.testing import ( - define_echo_model, - define_programmable_model, - test_models, - ) - - ai = Genkit() - - # Echo model - useful for verifying request formatting - echo, echo_action = define_echo_model(ai) - response = await ai.generate(model='echoModel', prompt='Hello') - # response contains: "[ECHO] user: "Hello"" - - # Programmable model - useful for testing specific scenarios - pm, pm_action = define_programmable_model(ai) - pm.responses = [GenerateResponse(message=Message(...))] - response = await ai.generate(model='programmableModel', prompt='test') - assert pm.last_request is not None - - # Test suite - validate model implementations - report = await test_models(ai, ['googleai/gemini-2.0-flash']) - for test in report: - print(f'{test["description"]}: {test["models"]}') - ``` - -Cross-Language Parity: - - JavaScript: js/ai/src/testing/model-tester.ts - - Go: go/ai/testutil_test.go - -See Also: - - https://genkit.dev for Genkit documentation -""" - -from copy import deepcopy -from typing import Any, TypedDict - -from pydantic import BaseModel, Field - -from genkit.ai import Genkit, Output -from genkit.codec import dump_json -from genkit.core.action import Action, ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.tracing import run_in_new_span -from genkit.core.typing import ( - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - Media, - MediaPart, - Message, - ModelInfo, - Part, - Role, - SpanMetadata, - TextPart, -) - - -class ProgrammableModel: - """A configurable model implementation for testing. - - This class allows test cases to define custom responses that the model - should return, making it useful for testing expected behavior in various - scenarios. - - Attributes: - request_count: Total number of requests received. - responses: List of predefined responses to return. - chunks: Optional list of chunks to stream for each request. - last_request: The most recent request received (deep copy). - - Example: - ```python - ai = Genkit() - pm, action = define_programmable_model(ai) - - # Set up responses - pm.responses = [ - GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='Response 1'))])), - GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='Response 2'))])), - ] - - # First call returns "Response 1" - result1 = await ai.generate(model='programmableModel', prompt='test') - - # Second call returns "Response 2" - result2 = await ai.generate(model='programmableModel', prompt='test') - - # Inspect last request - assert pm.last_request.messages[0].content[0].root.text == 'test' - ``` - """ - - def __init__(self) -> None: - """Initialize a new ProgrammableModel instance.""" - self._request_idx: int = 0 - self.request_count: int = 0 - self.responses: list[GenerateResponse] = [] - self.chunks: list[list[GenerateResponseChunk]] | None = None - self.last_request: GenerateRequest | None = None - - def reset(self) -> None: - """Reset the model state for reuse in tests.""" - self._request_idx = 0 - self.request_count = 0 - self.responses = [] - self.chunks = None - self.last_request = None - - def model_fn(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - """Process a generation request and return a programmed response. - - This function returns pre-configured responses and streams - pre-configured chunks based on the current request index. - - Args: - request: The generation request to process. - ctx: The action run context for streaming chunks. - - Returns: - The pre-configured response for the current request. - - Raises: - IndexError: If more requests are made than responses configured. - """ - # Store deep copy of request for inspection (matches JS behavior) - self.last_request = deepcopy(request) - self.request_count += 1 - - response = self.responses[self._request_idx] - if self.chunks and self._request_idx < len(self.chunks): - for chunk in self.chunks[self._request_idx]: - ctx.send_chunk(chunk) - self._request_idx += 1 - return response - - -def define_programmable_model( - ai: Genkit, - name: str = 'programmableModel', -) -> tuple[ProgrammableModel, Action]: - """Define a configurable programmable model for testing. - - Creates a model that returns pre-configured responses, useful for - testing specific scenarios like multi-turn conversations, tool calls, - or error conditions. - - Args: - ai: The Genkit instance to register the model with. - name: The name for the model. Defaults to 'programmableModel'. - - Returns: - A tuple of (ProgrammableModel instance, registered Action). - - Example: - ```python - ai = Genkit() - pm, action = define_programmable_model(ai) - - # Configure response with tool call - pm.responses = [ - GenerateResponse( - message=Message( - role=Role.MODEL, - content=[ - Part( - root=ToolRequestPart( - tool_request=ToolRequest( - name='myTool', - input={'arg': 'value'}, - ), - ) - ) - ], - ) - ) - ] - ``` - """ - pm = ProgrammableModel() - - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return pm.model_fn(request, ctx) - - action = ai.define_model(name=name, fn=model_fn) - - return (pm, action) - - -class EchoModel: - """A model implementation that echoes back the input with metadata. - - This model is useful for testing as it returns a readable representation - of the input it received, including config, tools, and output schema. - - The echo format is: - [ECHO] role: "content", role: "content" config tool_choice output - - Attributes: - last_request: The most recent request received. - stream_countdown: If True, streams "3", "2", "1" chunks before response. - - Example: - ```python - ai = Genkit() - echo, action = define_echo_model(ai, stream_countdown=True) - - response = await ai.generate(model='echoModel', prompt='Hello world', config={'temperature': 0.5}) - # Response text: '[ECHO] user: "Hello world" {"temperature":0.5}' - # With streaming: sends chunks "3", "2", "1" before final response - ``` - """ - - def __init__(self, stream_countdown: bool = False) -> None: - """Initialize a new EchoModel instance. - - Args: - stream_countdown: If True, stream "3", "2", "1" chunks before response. - """ - self.last_request: GenerateRequest | None = None - self.stream_countdown: bool = stream_countdown - - def model_fn(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - """Process a generation request and echo it back in the response. - - Args: - request: The generation request to process. - ctx: The action run context for streaming. - - Returns: - A response containing an echo of the input request details. - """ - self.last_request = request - - # Build echo string from messages - merged_txt = '' - for m in request.messages: - merged_txt += f' {m.role}: ' + ','.join( - dump_json(p.root.text) if p.root.text is not None else '""' for p in m.content - ) - echo_resp = f'[ECHO]{merged_txt}' - - # Add config, tools, and output info - if request.config: - echo_resp += f' {dump_json(request.config)}' - if request.tools: - echo_resp += f' tools={",".join(t.name for t in request.tools)}' - if request.tool_choice is not None: - echo_resp += f' tool_choice={request.tool_choice}' - if request.output and dump_json(request.output) != '{}': - echo_resp += f' output={dump_json(request.output)}' - - # Stream countdown chunks if enabled (matches JS behavior) - if self.stream_countdown: - for i, countdown in enumerate(['3', '2', '1']): - ctx.send_chunk( - chunk=GenerateResponseChunk(role=Role.MODEL, index=i, content=[Part(root=TextPart(text=countdown))]) - ) - - # NOTE: Part is a RootModel requiring root=TextPart(...) syntax. - return GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=echo_resp))])) - - -def define_echo_model( - ai: Genkit, - name: str = 'echoModel', - stream_countdown: bool = False, -) -> tuple[EchoModel, Action]: - """Define an echo model that returns input back as output. - - Creates a model that echoes the request back in a readable format, - useful for testing request formatting and middleware behavior. - - Args: - ai: The Genkit instance to register the model with. - name: The name for the model. Defaults to 'echoModel'. - stream_countdown: If True, stream "3", "2", "1" before final response. - - Returns: - A tuple of (EchoModel instance, registered Action). - - Example: - ```python - ai = Genkit() - echo, action = define_echo_model(ai, stream_countdown=True) - - # Test that middleware properly formats requests - response = await ai.generate(model='echoModel', prompt='test') - assert '[ECHO]' in response.text - ``` - """ - echo = EchoModel(stream_countdown=stream_countdown) - - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return echo.model_fn(request, ctx) - - action = ai.define_model(name=name, fn=model_fn) - - return (echo, action) - - -class StaticResponseModel: - """A model that always returns the same static response. - - Useful for simple test cases where a fixed response is needed. - - Attributes: - response_message: The message to always return. - last_request: The most recent request received. - request_count: Total number of requests received. - """ - - def __init__(self, message: dict[str, Any]) -> None: - """Initialize with a static message to return. - - Args: - message: The message data to always return. - """ - self.response_message: Message = Message.model_validate(message) - self.last_request: GenerateRequest | None = None - self.request_count: int = 0 - - def model_fn(self, request: GenerateRequest, _ctx: ActionRunContext) -> GenerateResponse: - """Return the static response. - - Args: - request: The generation request (stored but not used). - _ctx: The action run context (unused). - - Returns: - GenerateResponse with the static message. - """ - self.last_request = request - self.request_count += 1 - return GenerateResponse(message=self.response_message) - - -def define_static_response_model( - ai: Genkit, - message: dict[str, Any], - name: str = 'staticModel', -) -> tuple[StaticResponseModel, Action]: - """Define a model that always returns the same response. - - Args: - ai: The Genkit instance to register the model with. - message: The message data to always return. - name: The name for the model. Defaults to 'staticModel'. - - Returns: - A tuple of (StaticResponseModel instance, registered Action). - - Example: - ```python - ai = Genkit() - static, action = define_static_response_model( - ai, - message={'role': 'model', 'content': [{'text': 'Hello!'}]}, - ) - - # Always returns "Hello!" - response = await ai.generate(model='staticModel', prompt='anything') - assert response.text == 'Hello!' - ``` - """ - static = StaticResponseModel(message) - - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return static.model_fn(request, ctx) - - action = ai.define_model(name=name, fn=model_fn) - - return (static, action) - - -class SkipTestError(Exception): - """Exception raised to skip a test case. - - This is used internally by the test suite to indicate that a test - should be skipped (e.g., because the model doesn't support the - feature being tested). - """ - - -def skip() -> None: - """Skip the current test case. - - Raises: - SkipTestError: Always raised to skip the test. - """ - raise SkipTestError() - - -class ModelTestError(TypedDict, total=False): - """Error information from a failed test.""" - - message: str - stack: str | None - - -class ModelTestResult(TypedDict, total=False): - """Result of testing a single model on a single test case.""" - - name: str - passed: bool - skipped: bool - error: ModelTestError - - -class TestCaseReport(TypedDict): - """Report for a single test case across all models.""" - - description: str - models: list[ModelTestResult] - - -TestReport = list[TestCaseReport] -"""Complete test report for all test cases and models.""" - - -class GablorkenInput(BaseModel): - """Input for the gablorken tool used in tool calling tests.""" - - value: float = Field(..., description='The value to calculate gablorken for') - - -async def test_models(ai: Genkit, models: list[str]) -> TestReport: - r"""Run a standard test suite against one or more models. - - This function runs a series of tests to validate model implementations, - checking for basic functionality, multimodal support, conversation history, - system prompts, structured output, and tool calling. - - Test Suite Overview - =================== - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Model Test Suite β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ β”‚ - β”‚ Test Case β”‚ Description β”‚ Auto-Skip β”‚ - β”‚ ───────────────────────┼────────────────────────────────┼─────────────│ - β”‚ basic hi β”‚ Simple text generation β”‚ Never β”‚ - β”‚ multimodal β”‚ Image input processing β”‚ No media β”‚ - β”‚ history β”‚ Multi-turn conversation β”‚ No multiturn β”‚ - β”‚ system prompt β”‚ System message handling β”‚ Never β”‚ - β”‚ structured output β”‚ JSON schema output β”‚ Never β”‚ - β”‚ tool calling β”‚ Function calling β”‚ No tools β”‚ - β”‚ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - Args: - ai: The Genkit instance with models to test. - models: List of model names to test (e.g., ['googleai/gemini-2.0-flash']). - - Returns: - A TestReport containing results for each test case and model. - - Example: - ```python - from genkit.ai import Genkit - from genkit.plugins.google_genai import GoogleAI - from genkit.testing import test_models - - ai = Genkit(plugins=[GoogleAI()]) - - # Test multiple models - report = await test_models( - ai, - [ - 'googleai/gemini-2.0-flash', - 'googleai/gemini-1.5-pro', - ], - ) - - # Print results - for test in report: - print(f'\\n{test["description"]}:') - for model in test['models']: - status = 'βœ“' if model['passed'] else ('⊘' if model.get('skipped') else 'βœ—') - print(f' {status} {model["name"]}') - if 'error' in model: - print(f' Error: {model["error"]["message"]}') - ``` - - Note: - - Tests are automatically skipped if the model doesn't support - the required capability (e.g., tools, media, multiturn). - - A 'gablorkenTool' is automatically registered for tool calling tests. - - The test uses a small base64-encoded test image for multimodal tests. - - See Also: - - JS implementation: js/ai/src/testing/model-tester.ts - """ - - # Register the gablorken tool for tool calling tests - # NOTE: Tool name is camelCase to match JS implementation for parity - @ai.tool(name='gablorkenTool') - def gablorken_tool(input: GablorkenInput) -> float: - """Calculate the gablorken of a value. Use when need to calculate a gablorken.""" - return (input.value**3) + 1.407 - - async def get_model_info(model_name: str) -> ModelInfo | None: - """Get ModelInfo for a model, or None if not available. - - Args: - model_name: The name of the model to look up. - - Returns: - ModelInfo if available, None otherwise. - """ - model_action = await ai.registry.resolve_action(ActionKind.MODEL, model_name) - if model_action and model_action.metadata: - info_obj = model_action.metadata.get('model') - if isinstance(info_obj, ModelInfo): - return info_obj - return None - - # Define test cases - async def test_basic_hi(model: str) -> None: - """Test basic text generation.""" - response = await ai.generate(model=model, prompt='just say "Hi", literally') - got = response.text.strip() - assert 'hi' in got.lower(), f'Expected "Hi" in response, got: {got}' - - async def test_multimodal(model: str) -> None: - """Test multimodal (image) input.""" - info = await get_model_info(model) - if not (info and info.supports and info.supports.media): - skip() - - # Small test image (plus sign) - test_image = ( - 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAIAAACQkWg2' - 'AAABhGlDQ1BJQ0MgcHJvZmlsZQAAKJF9kT1Iw0AcxV9TpSoVETOIOGSoulgQFXHU' - 'KhShQqgVWnUwufRDaNKQtLg4Cq4FBz8Wqw4uzro6uAqC4AeIs4OToouU+L+k0CLG' - 'g+N+vLv3uHsHCLUi0+22MUA3ylYyHpPSmRUp9IpOhCCiFyMKs81ZWU7Ad3zdI8DX' - 'uyjP8j/35+jWsjYDAhLxDDOtMvE68dRm2eS8TyyygqIRnxOPWnRB4keuqx6/cc67' - 'LPBM0Uol54hFYinfwmoLs4KlE08SRzTdoHwh7bHGeYuzXqywxj35C8NZY3mJ6zQH' - 'EccCFiFDgooKNlBEGVFaDVJsJGk/5uMfcP0yuVRybYCRYx4l6FBcP/gf/O7Wzk2M' - 'e0nhGND+4jgfQ0BoF6hXHef72HHqJ0DwGbgymv5SDZj+JL3a1CJHQM82cHHd1NQ9' - '4HIH6H8yFUtxpSBNIZcD3s/omzJA3y3Qter11tjH6QOQoq4SN8DBITCcp+w1n3d3' - 'tPb275lGfz9aC3Kd0jYiSQAAAAlwSFlzAAAuIwAALiMBeKU/dgAAAAd0SU1FB+gJ' - 'BxQRO1/5qB8AAAAZdEVYdENvbW1lbnQAQ3JlYXRlZCB3aXRoIEdJTVBXgQ4XAAAA' - 'sUlEQVQoz61SMQqEMBDcO5SYToUE/IBPyRMCftAH+INUviApUwYjNkKCVcTiQK7I' - 'HSw45czODrMswCOQUkopEQZjzDiOWemdZfu+b5oGYYgx1nWNMPwB2vACAK01Y4wQ' - '8qGqqirL8jzPlNI9t64r55wQUgBA27be+xDCfaJhGJxzSqnv3UKIn7ne+2VZEB2s' - 'tZRSRLN93+d5RiRs28Y5RySEEI7jyEpFlp2mqeu6Zx75ApQwPdsIcq0ZAAAAAElF' - 'TkSuQmCC' - ) - - response = await ai.generate( - model=model, - prompt=[ - Part(root=MediaPart(media=Media(url=test_image))), - Part(root=TextPart(text='what math operation is this? plus, minus, multiply or divide?')), - ], - ) - got = response.text.strip().lower() - assert 'plus' in got, f'Expected "plus" in response, got: {got}' - - async def test_history(model: str) -> None: - """Test conversation history (multi-turn).""" - info = await get_model_info(model) - if not (info and info.supports and info.supports.multiturn): - skip() - - response1 = await ai.generate(model=model, prompt='My name is Glorb') - response2 = await ai.generate( - model=model, - prompt="What's my name?", - messages=response1.messages, - ) - got = response2.text.strip() - assert 'Glorb' in got, f'Expected "Glorb" in response, got: {got}' - - async def test_system_prompt(model: str) -> None: - """Test system prompt handling.""" - response = await ai.generate( - model=model, - prompt='Hi', - messages=[ - Message.model_validate({ - 'role': 'system', - 'content': [{'text': 'If the user says "Hi", just say "Bye"'}], - }), - ], - ) - got = response.text.strip() - assert 'Bye' in got, f'Expected "Bye" in response, got: {got}' - - async def test_structured_output(model: str) -> None: - """Test structured JSON output.""" - - class PersonInfo(BaseModel): - name: str - occupation: str - - response = await ai.generate( - model=model, - prompt='extract data as json from: Jack was a Lumberjack', - output=Output(schema=PersonInfo), - ) - got = response.output - assert got is not None, 'Expected structured output' - # Output can be dict or model instance depending on parsing - if isinstance(got, BaseModel): - got = got.model_dump() - - assert isinstance(got, dict), f'Expected output to be a dict or BaseModel, got {type(got)}' - assert got.get('name') == 'Jack', f"Expected name='Jack', got: {got.get('name')}" - assert got.get('occupation') == 'Lumberjack', f"Expected occupation='Lumberjack', got: {got.get('occupation')}" - - async def test_tool_calling(model: str) -> None: - """Test tool/function calling.""" - info = await get_model_info(model) - if not (info and info.supports and info.supports.tools): - skip() - - response = await ai.generate( - model=model, - prompt='what is a gablorken of 2? use provided tool', - tools=['gablorkenTool'], - ) - got = response.text.strip() - # 2^3 + 1.407 = 9.407 - assert '9.407' in got, f'Expected "9.407" in response, got: {got}' - - # Map of test cases - tests: dict[str, Any] = { - 'basic hi': test_basic_hi, - 'multimodal': test_multimodal, - 'history': test_history, - 'system prompt': test_system_prompt, - 'structured output': test_structured_output, - 'tool calling': test_tool_calling, - } - - # Run tests with tracing - report: TestReport = [] - - with run_in_new_span(SpanMetadata(name='testModels'), labels={'genkit:type': 'testSuite'}): - for test_name, test_fn in tests.items(): - with run_in_new_span(SpanMetadata(name=test_name), labels={'genkit:type': 'testCase'}): - case_report: TestCaseReport = { - 'description': test_name, - 'models': [], - } - - for model in models: - model_result: ModelTestResult = { - 'name': model, - 'passed': True, # Optimistic - } - - try: - await test_fn(model) - except SkipTestError: - model_result['passed'] = False - model_result['skipped'] = True - except AssertionError as e: - model_result['passed'] = False - model_result['error'] = { - 'message': str(e), - 'stack': None, - } - except Exception as e: - model_result['passed'] = False - model_result['error'] = { - 'message': str(e), - 'stack': None, - } - - case_report['models'].append(model_result) - - report.append(case_report) - - return report - - -# Export all public symbols -__all__ = [ - 'EchoModel', - 'GablorkenInput', - 'ModelTestError', - 'ModelTestResult', - 'ProgrammableModel', - 'SkipTestError', - 'StaticResponseModel', - 'TestCaseReport', - 'TestReport', - 'define_echo_model', - 'define_programmable_model', - 'define_static_response_model', - 'skip', - 'test_models', -] diff --git a/py/packages/genkit/src/genkit/types/__init__.py b/py/packages/genkit/src/genkit/types/__init__.py deleted file mode 100644 index 9ef3abacb8..0000000000 --- a/py/packages/genkit/src/genkit/types/__init__.py +++ /dev/null @@ -1,193 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""User-facing types for the Genkit framework. - -This module re-exports all public types that users may need when building -Genkit applications. It provides a single import point for common types -like Message, Part, Document, and response types. - -Overview: - Types in Genkit are organized into categories based on their use case. - This module exports types that are commonly needed in user code. - -Type Categories: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Category β”‚ Key Types β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Message/Part β”‚ Message, Part, TextPart, MediaPart, Role β”‚ - β”‚ Document β”‚ Document, DocumentData β”‚ - β”‚ Generation β”‚ GenerateRequest, GenerateResponse, OutputConf β”‚ - β”‚ Tools β”‚ ToolRequest, ToolResponse, ToolDefinition β”‚ - β”‚ Embedding β”‚ Embedding, EmbedRequest, EmbedResponse β”‚ - β”‚ Retrieval β”‚ RetrieverRequest, RetrieverResponse β”‚ - β”‚ Evaluation β”‚ EvalRequest, EvalResponse, Score β”‚ - β”‚ Model Info β”‚ ModelInfo, Supports, Constrained β”‚ - β”‚ Errors β”‚ GenkitError, ToolInterruptError β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Importing types for a Genkit application: - - ```python - from genkit import Genkit - from genkit.types import Message, Part, TextPart, Document - - ai = Genkit(...) - - # Create a message manually - message = Message( - role='user', - content=[Part(root=TextPart(text='Hello!'))], - ) - - # Create a document - doc = Document.from_text('Some content', metadata={'source': 'file.txt'}) - ``` - - Plugin authors may need model info types: - - ```python - from genkit.types import ModelInfo, Supports - - model_info = ModelInfo( - label='My Model', - supports=Supports( - multiturn=True, - tools=True, - systemRole=True, - ), - ) - ``` - -See Also: - - genkit.core.typing: Source definitions for these types - - genkit.blocks.document: Document class implementation -""" - -from genkit.blocks.document import Document -from genkit.blocks.tools import ToolInterruptError -from genkit.core.action._action import ActionRunContext -from genkit.core.error import GenkitError -from genkit.core.typing import ( - # Eval types - BaseEvalDataPoint, - Constrained, - CustomPart, - DataPart, - # Document types - DocumentData, - # Embedding types - Embedding, - EmbedRequest, - EmbedResponse, - EvalFnResponse, - EvalRequest, - EvalResponse, - EvalStatusEnum, - FinishReason, - GenerateActionOptions, - # Generation types - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationUsage, - Media, - MediaPart, - # Message and Part types - Message, - Metadata, - # Model info (for plugin authors) - ModelInfo, - OutputConfig, - Part, - ReasoningPart, - # Retriever types - RetrieverRequest, - RetrieverResponse, - Role, - Score, - Stage, - Supports, - TextPart, - ToolChoice, - ToolDefinition, - # Tool types - ToolRequest, - ToolRequestPart, - ToolResponse, - ToolResponsePart, -) -from genkit.model_types import GenerationCommonConfig, get_effective_api_key, get_request_api_key - -__all__ = [ - # Errors - 'GenkitError', - 'ToolInterruptError', - # Action context - 'ActionRunContext', - # Message and Part types - 'CustomPart', - 'DataPart', - 'Media', - 'MediaPart', - 'Message', - 'Metadata', - 'Part', - 'ReasoningPart', - 'Role', - 'TextPart', - 'ToolRequestPart', - 'ToolResponsePart', - # Document types - 'Document', - 'DocumentData', - # Generation types - 'FinishReason', - 'GenerateActionOptions', - 'GenerateRequest', - 'GenerateResponse', - 'GenerateResponseChunk', - 'GenerationCommonConfig', - 'get_request_api_key', - 'get_effective_api_key', - 'GenerationUsage', - 'OutputConfig', - 'ToolChoice', - # Tool types - 'ToolDefinition', - 'ToolRequest', - 'ToolResponse', - # Embedding types - 'EmbedRequest', - 'EmbedResponse', - 'Embedding', - # Retriever types - 'RetrieverRequest', - 'RetrieverResponse', - # Eval types - 'BaseEvalDataPoint', - 'EvalFnResponse', - 'EvalRequest', - 'EvalResponse', - 'EvalStatusEnum', - 'Score', - # Model info (for plugin authors) - 'Constrained', - 'ModelInfo', - 'Stage', - 'Supports', -] diff --git a/py/packages/genkit/src/genkit/web/__init__.py b/py/packages/genkit/src/genkit/web/__init__.py deleted file mode 100644 index 1e9265743d..0000000000 --- a/py/packages/genkit/src/genkit/web/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Web framework for the Genkit framework.""" diff --git a/py/packages/genkit/src/genkit/web/requests.py b/py/packages/genkit/src/genkit/web/requests.py deleted file mode 100644 index 70bca81d86..0000000000 --- a/py/packages/genkit/src/genkit/web/requests.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Helper functions for reading JSON request bodies.""" - -from starlette.datastructures import QueryParams -from starlette.requests import Request - - -def is_streaming_requested(request: Request) -> bool: - """Check if streaming is requested. - - Streaming is requested if the query parameter 'stream' is set to 'true' or - if the Accept header is 'text/event-stream'. - - Args: - request: Starlette request object. - - Returns: - True if streaming is requested, False otherwise. - """ - by_header = request.headers.get('accept', '') == 'text/event-stream' - by_query = is_query_flag_enabled(request.query_params, 'stream') - return by_header or by_query - - -def is_query_flag_enabled(query_params: QueryParams, flag: str) -> bool: - """Check if a query flag is enabled. - - Args: - query_params: Dictionary containing parsed query parameters. - flag: Flag name to check. - - Returns: - True if the query flag is enabled, False otherwise. - """ - return query_params.get(flag, 'false') == 'true' diff --git a/py/packages/genkit/src/genkit/web/typing.py b/py/packages/genkit/src/genkit/web/typing.py deleted file mode 100644 index 248c0e9892..0000000000 --- a/py/packages/genkit/src/genkit/web/typing.py +++ /dev/null @@ -1,141 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""ASGI-compatible type definitions for Genkit web framework integration. - -This module provides Protocol-based type definitions that are compatible with -any ASGI framework (litestar, starlette, fastapi, Django, etc.) without -requiring framework-specific type unions. - -The ASGI specification defines interfaces using structural typing, so we use -typing.Protocol to match this approach. Any framework that follows the ASGI -spec will be compatible with these types. - -Supported frameworks: - - asgiref (Django) - - FastAPI - - Litestar - - Starlette - - Any other ASGI-compliant framework - -Example: - ```python - from genkit.web.typing import Scope, Receive, Send - - - async def app(scope: Scope, receive: Receive, send: Send) -> None: - # This signature works with any ASGI framework - ... - ``` - -See Also: - - ASGI Spec: https://asgi.readthedocs.io/ - - PEP 544 (Protocols): https://peps.python.org/pep-0544/ -""" - -from __future__ import annotations - -from collections.abc import Awaitable, Callable, MutableMapping -from typing import Any - -# These Protocol-based types follow the ASGI specification and are compatible -# with any ASGI framework. Using Protocols instead of Union types allows: -# -# 1. Any framework implementing the ASGI interface to work -# 2. No import-time dependencies on specific frameworks -# 3. Proper structural subtyping for type checkers - -# ASGI Scope - a dict-like object with connection info -# Using MutableMapping[str, Any] as the base provides structural compatibility -Scope = MutableMapping[str, Any] -"""ASGI scope - connection metadata dict. - -Contains at minimum: - - type: 'http', 'websocket', or 'lifespan' - - asgi: dict with 'version' key - -For HTTP connections, also includes: - - method: HTTP method (GET, POST, etc.) - - path: URL path - - headers: list of (name, value) byte tuples -""" - -# ASGI Receive Callable -# Async function that returns the next message from the client -Receive = Callable[[], Awaitable[MutableMapping[str, Any]]] -"""ASGI receive callable - async function to get next message from client.""" - -# ASGI Send Callable -# Async function that sends a message to the client -Send = Callable[[MutableMapping[str, Any]], Awaitable[None]] -"""ASGI send callable - async function to send message to client.""" - - -# Type alias for ASGI applications -# Using Any because each framework (Litestar, Starlette, FastAPI) has its own -# type definitions that aren't structurally compatible with our Protocol. -# At runtime, they all work correctly - this is purely a type checker limitation. -Application = Any -"""Type alias for ASGI application objects. - -Note: Uses Any because external frameworks (Litestar, Starlette, etc.) define -their own ASGI types that aren't structurally compatible with our Protocol. -At runtime, all ASGI-compliant apps work correctly. -""" - - -# Specialized Scope Types (for documentation, not enforced at runtime) -# These are aliases for Scope with documentation about expected keys. -# Type checkers treat them as Scope, but the docstrings help developers. - -HTTPScope = Scope -"""HTTP connection scope. - -Expected keys beyond base Scope: - - method: str (GET, POST, etc.) - - path: str - - query_string: bytes - - headers: list[tuple[bytes, bytes]] -""" - -WebSocketScope = Scope -"""WebSocket connection scope. - -Expected keys beyond base Scope: - - path: str - - query_string: bytes - - headers: list[tuple[bytes, bytes]] -""" - -LifespanScope = Scope -"""Lifespan scope for startup/shutdown. - -Expected keys beyond base Scope: - - type: 'lifespan' -""" - - -# Handler Type Aliases - -LifespanHandler = Callable[[LifespanScope, Receive, Send], Awaitable[None]] -"""ASGI lifespan handler - manages app startup and shutdown.""" - -StartupHandler = Callable[[], Awaitable[None]] -"""Simple startup/shutdown handler (0-argument async function). - -Used by Starlette/Litestar for on_startup/on_shutdown callbacks. -This is distinct from LifespanHandler which is the full ASGI protocol. -""" diff --git a/py/packages/genkit/tests/genkit/ai/ai_plugin_test.py b/py/packages/genkit/tests/genkit/ai/ai_plugin_test.py index 6e4fea94a6..d81f193f29 100644 --- a/py/packages/genkit/tests/genkit/ai/ai_plugin_test.py +++ b/py/packages/genkit/tests/genkit/ai/ai_plugin_test.py @@ -22,11 +22,11 @@ import pytest -from genkit.ai import Genkit, Plugin -from genkit.core.action import Action, ActionMetadata, ActionRunContext -from genkit.core.registry import ActionKind -from genkit.core.typing import FinishReason -from genkit.types import GenerateRequest, GenerateResponse, Message, Part, Role, TextPart +from genkit import Genkit, Message, ModelResponse, Part, Plugin, Role, TextPart +from genkit._core._action import Action, ActionMetadata, ActionRunContext +from genkit._core._model import ModelRequest +from genkit._core._registry import ActionKind +from genkit._core._typing import FinishReason class AsyncResolveOnlyPlugin(Plugin): @@ -46,8 +46,8 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: if name != f'{self.name}/lazy-model': return None - async def _generate(req: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse( + async def _generate(req: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse( message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='OK: lazy'))]), finish_reason=FinishReason.STOP, ) @@ -85,8 +85,8 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: if name != f'{self.name}/init-model': return None - async def _generate(req: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse( + async def _generate(req: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse( message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='OK: resolve'))]), finish_reason=FinishReason.STOP, ) @@ -111,7 +111,7 @@ async def list_actions(self) -> list[ActionMetadata]: async def test_async_resolve_is_awaited_via_generate() -> None: """Test that async resolve is awaited when calling generate.""" ai = Genkit(plugins=[AsyncResolveOnlyPlugin()]) - resp = await ai.generate('async-resolve-only/lazy-model', prompt='hello') + resp = await ai.generate(model='async-resolve-only/lazy-model', prompt='hello') assert resp.text == 'OK: lazy' @@ -119,5 +119,5 @@ async def test_async_resolve_is_awaited_via_generate() -> None: async def test_async_init_is_awaited_via_generate() -> None: """Test that async init is awaited when calling generate.""" ai = Genkit(plugins=[AsyncInitPlugin()]) - resp = await ai.generate('async-init-plugin/init-model', prompt='hello') + resp = await ai.generate(model='async-init-plugin/init-model', prompt='hello') assert resp.text == 'OK: resolve' diff --git a/py/packages/genkit/tests/genkit/ai/ai_registry_test.py b/py/packages/genkit/tests/genkit/ai/ai_registry_test.py index d3b8822bff..6206935632 100644 --- a/py/packages/genkit/tests/genkit/ai/ai_registry_test.py +++ b/py/packages/genkit/tests/genkit/ai/ai_registry_test.py @@ -15,53 +15,19 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Tests for the AI registry module. - -This module contains unit tests for the GenkitRegistry class and its associated -functionality, ensuring proper registration and management of Genkit resources. - -Test Coverage -============= - -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ Test Case β”‚ Description β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ get_func_description Tests β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ test_with_explicit_description β”‚ Explicit desc takes precedence β”‚ -β”‚ test_with_docstring β”‚ Docstring used when no explicit desc β”‚ -β”‚ test_without_docstring β”‚ Empty string when no docstring β”‚ -β”‚ test_with_none_docstring β”‚ Empty string when docstring is None β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ define_json_schema Tests β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ test_define_json_schema_basic β”‚ Register a basic JSON schema β”‚ -β”‚ test_define_json_schema_complex β”‚ Register complex nested schema β”‚ -β”‚ test_define_json_schema_returns β”‚ Method returns the schema β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ define_dynamic_action_provider Tests β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ test_define_dap_with_string β”‚ DAP with string config β”‚ -β”‚ test_define_dap_with_config β”‚ DAP with DapConfig object β”‚ -β”‚ test_define_dap_returns_provider β”‚ Method returns DynamicActionProvider β”‚ -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ -""" +"""Tests for the AI registry module.""" import unittest import pytest -from genkit.ai import Genkit -from genkit.ai._registry import get_func_description -from genkit.blocks.dap import DapCacheConfig, DapConfig, DapValue, DynamicActionProvider +from genkit import Genkit +from genkit._core._dap import DapValue, DynamicActionProvider +from genkit._core._flow import get_func_description class TestGetFuncDescription(unittest.TestCase): - """Test the get_func_description function.""" - def test_get_func_description_with_explicit_description(self) -> None: - """Test that explicit description takes precedence over docstring.""" - def test_func() -> None: """This docstring should be ignored.""" pass @@ -70,8 +36,6 @@ def test_func() -> None: self.assertEqual(description, 'Explicit description') def test_get_func_description_with_docstring(self) -> None: - """Test that docstring is used when no explicit description is provided.""" - def test_func() -> None: """This is the function's docstring.""" pass @@ -80,8 +44,6 @@ def test_func() -> None: self.assertEqual(description, "This is the function's docstring.") def test_get_func_description_without_docstring(self) -> None: - """Test that empty string is returned when no docstring is present.""" - def test_func() -> None: pass @@ -89,8 +51,6 @@ def test_func() -> None: self.assertEqual(description, '') def test_get_func_description_with_none_docstring(self) -> None: - """Test that empty string is returned when docstring is None.""" - def test_func() -> None: pass @@ -101,10 +61,7 @@ def test_func() -> None: class TestDefineJsonSchema: - """Tests for the define_json_schema method.""" - def test_define_json_schema_basic(self) -> None: - """Test registering a basic JSON schema.""" ai = Genkit() schema = ai.define_json_schema( @@ -124,7 +81,6 @@ def test_define_json_schema_basic(self) -> None: assert 'properties' in schema def test_define_json_schema_complex(self) -> None: - """Test registering a complex nested JSON schema.""" ai = Genkit() schema = ai.define_json_schema( @@ -152,7 +108,6 @@ def test_define_json_schema_complex(self) -> None: assert schema is not None assert schema['type'] == 'object' - # Type checker doesn't narrow dict[str, object] properly after isinstance properties: dict[str, object] = schema['properties'] # type: ignore[assignment] assert isinstance(properties, dict) assert 'ingredients' in properties @@ -161,7 +116,6 @@ def test_define_json_schema_complex(self) -> None: assert ingredients['type'] == 'array' def test_define_json_schema_returns_same_schema(self) -> None: - """Test that define_json_schema returns the schema for convenience.""" ai = Genkit() input_schema: dict[str, object] = { @@ -170,17 +124,12 @@ def test_define_json_schema_returns_same_schema(self) -> None: } returned_schema = ai.define_json_schema('StringSchema', input_schema) - - # Should return the same schema object assert returned_schema is input_schema class TestDefineDynamicActionProvider: - """Tests for the define_dynamic_action_provider method via Genkit.""" - @pytest.mark.asyncio async def test_define_dap_with_string_config(self) -> None: - """Test defining a DAP with a string name.""" ai = Genkit() async def dap_fn() -> DapValue: @@ -189,34 +138,26 @@ async def dap_fn() -> DapValue: dap = ai.define_dynamic_action_provider('my-dap', dap_fn) assert isinstance(dap, DynamicActionProvider) - assert dap.config.name == 'my-dap' @pytest.mark.asyncio - async def test_define_dap_with_config_object(self) -> None: - """Test defining a DAP with a DapConfig object.""" + async def test_define_dap_with_options(self) -> None: ai = Genkit() async def dap_fn() -> DapValue: return {} - config = DapConfig( - name='configured-dap', + dap = ai.define_dynamic_action_provider( + 'configured-dap', + dap_fn, description='A configured DAP', - cache_config=DapCacheConfig(ttl_millis=5000), + cache_ttl_millis=5000, metadata={'custom': 'value'}, ) - dap = ai.define_dynamic_action_provider(config, dap_fn) - assert isinstance(dap, DynamicActionProvider) - assert dap.config.name == 'configured-dap' - assert dap.config.description == 'A configured DAP' - assert dap.config.cache_config is not None - assert dap.config.cache_config.ttl_millis == 5000 @pytest.mark.asyncio async def test_define_dap_returns_provider(self) -> None: - """Test that define_dynamic_action_provider returns a DynamicActionProvider.""" ai = Genkit() async def dap_fn() -> DapValue: @@ -225,7 +166,6 @@ async def dap_fn() -> DapValue: result = ai.define_dynamic_action_provider('test-dap', dap_fn) assert isinstance(result, DynamicActionProvider) - # Should have required methods assert hasattr(result, 'get_action') assert hasattr(result, 'list_action_metadata') assert hasattr(result, 'invalidate_cache') diff --git a/py/packages/genkit/tests/genkit/blocks/dap_test.py b/py/packages/genkit/tests/genkit/ai/dap_test.py similarity index 62% rename from py/packages/genkit/tests/genkit/blocks/dap_test.py rename to py/packages/genkit/tests/genkit/ai/dap_test.py index 51e52e0d64..02b650cd19 100644 --- a/py/packages/genkit/tests/genkit/blocks/dap_test.py +++ b/py/packages/genkit/tests/genkit/ai/dap_test.py @@ -14,61 +14,29 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Tests for the Dynamic Action Provider (DAP) module. - -This module contains comprehensive tests for the DAP functionality, -ensuring parity with the JavaScript implementation in: - js/core/tests/dynamic-action-provider_test.ts - -Test Coverage -============= - -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ Test Case β”‚ Description β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ test_gets_specific_action β”‚ Get a specific action by type and name β”‚ -β”‚ test_lists_action_metadata β”‚ List metadata with wildcard pattern β”‚ -β”‚ test_caches_actions β”‚ Verify caching prevents redundant fetches β”‚ -β”‚ test_invalidates_cache β”‚ Cache invalidation forces fresh fetch β”‚ -β”‚ test_respects_cache_ttl β”‚ TTL expiration triggers refresh β”‚ -β”‚ test_lists_actions_with_prefixβ”‚ Prefix matching for action names β”‚ -β”‚ test_lists_actions_exact_matchβ”‚ Exact name matching β”‚ -β”‚ test_gets_action_metadata_rec β”‚ Reflection API metadata record format β”‚ -β”‚ test_handles_concurrent_reqs β”‚ Concurrent requests share single fetch β”‚ -β”‚ test_handles_fetch_errors β”‚ Error recovery and cache invalidation β”‚ -β”‚ test_identifies_dap β”‚ Type identification helper β”‚ -β”‚ test_transform_dap_value β”‚ Value to metadata transformation β”‚ -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ -""" +"""Tests for the Dynamic Action Provider (DAP) module.""" import asyncio import pytest -from genkit.blocks.dap import ( - DapCacheConfig, - DapConfig, +from genkit._core._action import Action, ActionKind +from genkit._core._dap import ( DapValue, DynamicActionProvider, define_dynamic_action_provider, is_dynamic_action_provider, - transform_dap_value, ) -from genkit.core.action import Action -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry +from genkit._core._registry import Registry @pytest.fixture def registry() -> Registry: - """Create a fresh registry for each test.""" return Registry() @pytest.fixture def tool1(registry: Registry) -> Action: - """Create tool1 action for testing.""" - async def tool1_fn(input: str) -> str: return 'tool1' @@ -82,8 +50,6 @@ async def tool1_fn(input: str) -> str: @pytest.fixture def tool2(registry: Registry) -> Action: - """Create tool2 action for testing.""" - async def tool2_fn(input: str) -> str: return 'tool2' @@ -97,8 +63,6 @@ async def tool2_fn(input: str) -> str: @pytest.fixture def other_tool(registry: Registry) -> Action: - """Create other-tool action for testing prefix matching.""" - async def other_tool_fn(input: str) -> str: return 'other' @@ -112,10 +76,6 @@ async def other_tool_fn(input: str) -> str: @pytest.mark.asyncio async def test_gets_specific_action(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test getting a specific action by type and name. - - Corresponds to JS test: 'gets a specific action' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -132,10 +92,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_lists_action_metadata(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test listing action metadata with wildcard. - - Corresponds to JS test: 'lists action metadata' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -154,10 +110,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_caches_actions(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test that actions are cached across multiple calls. - - Corresponds to JS test: 'caches the actions' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -183,10 +135,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_invalidates_cache(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test that cache invalidation forces a fresh fetch. - - Corresponds to JS test: 'invalidates the cache' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -207,10 +155,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_respects_cache_ttl(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test that cache respects TTL configuration. - - Corresponds to JS test: 'respects cache ttl' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -218,8 +162,7 @@ async def dap_fn() -> DapValue: call_count += 1 return {'tool': [tool1, tool2]} - config = DapConfig(name='my-dap', cache_config=DapCacheConfig(ttl_millis=10)) - dap = define_dynamic_action_provider(registry, config, dap_fn) + dap = define_dynamic_action_provider(registry, 'my-dap', dap_fn, cache_ttl_millis=10) await dap.get_action('tool', 'tool1') assert call_count == 1 @@ -233,10 +176,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_lists_actions_with_prefix(registry: Registry, tool1: Action, tool2: Action, other_tool: Action) -> None: - """Test listing actions with prefix pattern matching. - - Corresponds to JS test: 'lists actions with prefix' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -255,10 +194,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_lists_actions_exact_match(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test listing actions with exact name matching. - - Corresponds to JS test: 'lists actions with exact match' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -276,10 +211,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_gets_action_metadata_record(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test getting action metadata record for reflection API. - - Corresponds to JS test: 'gets action metadata record' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -304,21 +235,16 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_handles_concurrent_requests(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test that concurrent requests share a single fetch operation. - - Corresponds to JS test: 'handles concurrent requests' - """ call_count = 0 async def dap_fn() -> DapValue: nonlocal call_count call_count += 1 - await asyncio.sleep(0.01) # Simulate async work + await asyncio.sleep(0.01) return {'tool': [tool1, tool2]} dap = define_dynamic_action_provider(registry, 'my-dap', dap_fn) - # Run two requests concurrently results = await asyncio.gather( dap.list_action_metadata('tool', '*'), dap.list_action_metadata('tool', '*'), @@ -329,16 +255,11 @@ async def dap_fn() -> DapValue: assert len(metadata2) == 2 assert metadata1[0] == tool1.metadata assert metadata2[0] == tool1.metadata - # Only one fetch should have occurred assert call_count == 1 @pytest.mark.asyncio async def test_handles_fetch_errors(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test error handling and cache invalidation on fetch failure. - - Corresponds to JS test: 'handles fetch errors' - """ call_count = 0 async def dap_fn() -> DapValue: @@ -350,12 +271,10 @@ async def dap_fn() -> DapValue: dap = define_dynamic_action_provider(registry, 'my-dap', dap_fn) - # First call should raise with pytest.raises(RuntimeError, match='Fetch failed'): await dap.list_action_metadata('tool', '*') assert call_count == 1 - # Second call should succeed (cache was invalidated) metadata = await dap.list_action_metadata('tool', '*') assert len(metadata) == 2 assert call_count == 2 @@ -363,74 +282,16 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_identifies_dap(registry: Registry, tool1: Action) -> None: - """Test the is_dynamic_action_provider helper function. - - Corresponds to JS test: 'identifies dynamic action providers' - """ - async def dap_fn() -> DapValue: return {} dap = define_dynamic_action_provider(registry, 'my-dap', dap_fn) assert is_dynamic_action_provider(dap) is True - - # Regular action should not be identified as DAP assert is_dynamic_action_provider(tool1) is False -def test_transform_dap_value(tool1: Action, tool2: Action) -> None: - """Test the transform_dap_value utility function.""" - value: DapValue = {'tool': [tool1, tool2]} - - metadata = transform_dap_value(value) - - assert 'tool' in metadata - assert len(metadata['tool']) == 2 - assert metadata['tool'][0] == tool1.metadata - assert metadata['tool'][1] == tool2.metadata - - -def test_dap_config_string_normalization(registry: Registry) -> None: - """Test that string config is normalized to DapConfig. - - The define_dynamic_action_provider should accept either a string - or a DapConfig object. - """ - - async def dap_fn() -> DapValue: - return {} - - # String config - dap = define_dynamic_action_provider(registry, 'my-dap', dap_fn) - assert isinstance(dap, DynamicActionProvider) - assert dap.config.name == 'my-dap' - - -def test_dap_config_with_full_options(registry: Registry) -> None: - """Test DapConfig with all options specified.""" - - async def dap_fn() -> DapValue: - return {} - - config = DapConfig( - name='full-config-dap', - description='A DAP with all options', - cache_config=DapCacheConfig(ttl_millis=5000), - metadata={'custom': 'value'}, - ) - - dap = define_dynamic_action_provider(registry, config, dap_fn) - assert dap.config.name == 'full-config-dap' - assert dap.config.description == 'A DAP with all options' - assert dap.config.cache_config is not None - assert dap.config.cache_config.ttl_millis == 5000 - assert dap.config.metadata == {'custom': 'value'} - - @pytest.mark.asyncio async def test_get_action_returns_none_for_unknown_type(registry: Registry, tool1: Action) -> None: - """Test that get_action returns None for unknown action types.""" - async def dap_fn() -> DapValue: return {'tool': [tool1]} @@ -442,8 +303,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_get_action_returns_none_for_unknown_name(registry: Registry, tool1: Action) -> None: - """Test that get_action returns None for unknown action names.""" - async def dap_fn() -> DapValue: return {'tool': [tool1]} @@ -455,8 +314,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_list_action_metadata_returns_empty_for_unknown_type(registry: Registry, tool1: Action) -> None: - """Test that list_action_metadata returns empty list for unknown types.""" - async def dap_fn() -> DapValue: return {'tool': [tool1]} @@ -468,7 +325,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_negative_ttl_disables_caching(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test that negative TTL disables caching (always fetch fresh).""" call_count = 0 async def dap_fn() -> DapValue: @@ -476,8 +332,7 @@ async def dap_fn() -> DapValue: call_count += 1 return {'tool': [tool1, tool2]} - config = DapConfig(name='my-dap', cache_config=DapCacheConfig(ttl_millis=-1)) - dap = define_dynamic_action_provider(registry, config, dap_fn) + dap = define_dynamic_action_provider(registry, 'my-dap', dap_fn, cache_ttl_millis=-1) await dap.get_action('tool', 'tool1') assert call_count == 1 @@ -489,7 +344,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_zero_ttl_uses_default(registry: Registry, tool1: Action, tool2: Action) -> None: - """Test that zero TTL uses the default (3000ms).""" call_count = 0 async def dap_fn() -> DapValue: @@ -497,8 +351,7 @@ async def dap_fn() -> DapValue: call_count += 1 return {'tool': [tool1, tool2]} - config = DapConfig(name='my-dap', cache_config=DapCacheConfig(ttl_millis=0)) - dap = define_dynamic_action_provider(registry, config, dap_fn) + dap = define_dynamic_action_provider(registry, 'my-dap', dap_fn, cache_ttl_millis=0) await dap.get_action('tool', 'tool1') assert call_count == 1 @@ -510,8 +363,6 @@ async def dap_fn() -> DapValue: @pytest.mark.asyncio async def test_get_action_metadata_record_raises_on_missing_name(registry: Registry) -> None: - """Test that get_action_metadata_record raises error when action has no name.""" - async def nameless_fn(input: str) -> str: return 'nameless' @@ -519,9 +370,8 @@ async def nameless_fn(input: str) -> str: name='nameless', kind=ActionKind.TOOL, fn=nameless_fn, - metadata={}, # No 'name' in metadata + metadata={}, ) - # Clear the name attribute to simulate a nameless action nameless_action._name = '' async def dap_fn() -> DapValue: @@ -531,3 +381,18 @@ async def dap_fn() -> DapValue: with pytest.raises(ValueError, match='name required'): await dap.get_action_metadata_record('dap/my-dap') + + +def test_define_dap_with_full_options(registry: Registry) -> None: + async def dap_fn() -> DapValue: + return {} + + dap = define_dynamic_action_provider( + registry, + 'full-config-dap', + dap_fn, + description='A DAP with all options', + cache_ttl_millis=5000, + metadata={'custom': 'value'}, + ) + assert isinstance(dap, DynamicActionProvider) diff --git a/py/packages/genkit/tests/genkit/blocks/document_test.py b/py/packages/genkit/tests/genkit/ai/document_test.py similarity index 55% rename from py/packages/genkit/tests/genkit/blocks/document_test.py rename to py/packages/genkit/tests/genkit/ai/document_test.py index 58f3279edf..ca2e640e63 100644 --- a/py/packages/genkit/tests/genkit/blocks/document_test.py +++ b/py/packages/genkit/tests/genkit/ai/document_test.py @@ -18,11 +18,9 @@ from typing import cast -from genkit.blocks.document import Document -from genkit.core.typing import ( - DocumentData, +from genkit import Document +from genkit._core._typing import ( DocumentPart, - Embedding, Media, MediaPart, TextPart, @@ -44,25 +42,18 @@ def test_makes_deep_copy() -> None: assert doc.metadata['foo'] == 'bar' -def test_from_document_data() -> None: - """Test creating a Document from DocumentData.""" - doc = Document.from_document_data(DocumentData(content=[DocumentPart(root=TextPart(text='some text'))])) - - assert doc.text() == 'some text' - - def test_simple_text_document() -> None: """Test creating a simple text Document.""" doc = Document.from_text('sample text') - assert doc.text() == 'sample text' + assert doc.text == 'sample text' def test_media_document() -> None: """Test creating a media Document.""" doc = Document.from_media(url='data:one') - assert doc.media() == [ + assert doc.media == [ Media(url='data:one'), ] @@ -74,9 +65,9 @@ def test_from_data_text_document() -> None: metadata = {'embedMetadata': {'embeddingType': 'text'}} doc = Document.from_data(data, data_type, metadata) - assert doc.text() == data + assert doc.text == data assert doc.metadata == metadata - assert doc.data_type() == data_type + assert doc.data_type == data_type def test_from_data_media_document() -> None: @@ -86,98 +77,58 @@ def test_from_data_media_document() -> None: metadata = {'embedMetadata': {'embeddingType': 'image'}} doc = Document.from_data(data, data_type, metadata) - assert doc.media() == [ + assert doc.media == [ Media(url=data, content_type=data_type), ] assert doc.metadata == metadata - assert doc.data_type() == data_type + assert doc.data_type == data_type def test_concatenates_text() -> None: - """Test that text() concatenates multiple text parts.""" + """Test that text concatenates multiple text parts.""" content = [DocumentPart(root=TextPart(text='hello')), DocumentPart(root=TextPart(text='world'))] doc = Document(content=content) - assert doc.text() == 'helloworld' + assert doc.text == 'helloworld' def test_multiple_media_document() -> None: - """Test that media() returns all media parts.""" + """Test that media returns all media parts.""" content = [ DocumentPart(root=MediaPart(media=Media(url='data:one'))), DocumentPart(root=MediaPart(media=Media(url='data:two'))), ] doc = Document(content=content) - assert doc.media() == [ + assert doc.media == [ Media(url='data:one'), Media(url='data:two'), ] def test_data_with_text() -> None: - """Test data() with a text document.""" + """Test data with a text document.""" doc = Document.from_text('hello') - assert doc.data() == 'hello' + assert doc.data == 'hello' def test_data_with_media() -> None: - """Test data() with a media document.""" + """Test data with a media document.""" doc = Document.from_media(url='gs://somebucket/someimage.png', content_type='image/png') - assert doc.data() == 'gs://somebucket/someimage.png' + assert doc.data == 'gs://somebucket/someimage.png' def test_data_type_with_text() -> None: - """Test data_type() with a text document.""" + """Test data_type with a text document.""" doc = Document.from_text('hello') - assert doc.data_type() == 'text' + assert doc.data_type == 'text' def test_data_type_with_media() -> None: - """Test data_type() with a media document.""" + """Test data_type with a media document.""" doc = Document.from_media(url='gs://somebucket/someimage.png', content_type='image/png') - assert doc.data_type() == 'image/png' - - -def test_get_embedding_documents() -> None: - """Test getting embedding documents for a single embedding.""" - doc = Document.from_text('foo') - embeddings: list[Embedding] = [Embedding(embedding=[0.1, 0.2, 0.3])] - docs = doc.get_embedding_documents(embeddings) - - assert docs == [doc] - - -def test_get_embedding_documents_multiple_embeddings() -> None: - """Test getting embedding documents for multiple embeddings.""" - url = 'gs://somebucket/somevideo.mp4' - content_type = 'video/mp4' - metadata = {'start': 0, 'end': 60} - doc = Document.from_media(url, content_type, metadata) - embeddings: list[Embedding] = [] - - for start in range(0, 60, 15): - embeddings.append(make_test_embedding(start)) - docs = doc.get_embedding_documents(embeddings) - - assert len(docs) == len(embeddings) - - for i in range(len(docs)): - assert docs[i].content == doc.content - metadata = docs[i].metadata or {} - assert metadata.get('embedMetadata') == embeddings[i].metadata - orig_metadata = dict(metadata) - orig_metadata.pop('embedMetadata', None) - assert orig_metadata, doc.metadata - - -def make_test_embedding(start: int) -> Embedding: - """Helper to create a test embedding with specific metadata.""" - return Embedding( - embedding=[0.1, 0.2, 0.3], - metadata={'embeddingType': 'video', 'start': start, 'end': start + 15}, - ) + assert doc.data_type == 'image/png' diff --git a/py/packages/genkit/tests/genkit/blocks/embedding_test.py b/py/packages/genkit/tests/genkit/ai/embedding_test.py similarity index 95% rename from py/packages/genkit/tests/genkit/blocks/embedding_test.py rename to py/packages/genkit/tests/genkit/ai/embedding_test.py index 2fdb1d73ac..786e242b7d 100644 --- a/py/packages/genkit/tests/genkit/blocks/embedding_test.py +++ b/py/packages/genkit/tests/genkit/ai/embedding_test.py @@ -23,18 +23,16 @@ import pytest from pydantic import BaseModel -from genkit.ai._aio import Genkit -from genkit.blocks.document import Document -from genkit.blocks.embedding import ( +from genkit import Document, Genkit +from genkit._ai._embedding import ( EmbedderOptions, EmbedderSupports, create_embedder_ref, embedder_action_metadata, ) -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionResponse -from genkit.core.schema import to_json_schema -from genkit.core.typing import Embedding, EmbedRequest, EmbedResponse +from genkit._core._action import Action, ActionMetadata, ActionResponse +from genkit._core._schema import to_json_schema +from genkit._core._typing import Embedding, EmbedRequest, EmbedResponse def test_embedder_action_metadata() -> None: @@ -166,7 +164,7 @@ async def mock_arun_side_effect(request: object, *args: object, **kwargs: object embed_response = await fn(request) return ActionResponse(response=embed_response, trace_id='mock_trace_id') - mock_action.arun = AsyncMock(side_effect=mock_arun_side_effect) + mock_action.run = AsyncMock(side_effect=mock_arun_side_effect) self.actions[kind, name] = mock_action return mock_action @@ -227,9 +225,9 @@ async def fake_embedder_fn(request: EmbedRequest) -> EmbedResponse: embed_action = await registry.resolve_action('embedder', 'my-plugin/my-embedder') assert embed_action is not None - embed_action.arun.assert_called_once() + embed_action.run.assert_called_once() - called_request = embed_action.arun.call_args[0][0] + called_request = embed_action.run.call_args[0][0] assert isinstance(called_request, EmbedRequest) assert called_request.input == [content] # Check if config from EmbedderRef and options are merged correctly @@ -263,7 +261,7 @@ async def fake_embedder_fn(request: EmbedRequest) -> EmbedResponse: assert response[0].embedding == [4.0, 5.0, 6.0] embed_action = await registry.resolve_action('embedder', 'another-embedder') - called_request = embed_action.arun.call_args[0][0] + called_request = embed_action.run.call_args[0][0] assert called_request.options == {'custom_setting': 'high'} @@ -303,7 +301,7 @@ async def fake_embedder_fn(request: EmbedRequest) -> EmbedResponse: assert response[1].embedding == [2.0, 2.1] embed_action = await registry.resolve_action('embedder', 'multi-embedder') - called_request = embed_action.arun.call_args[0][0] + called_request = embed_action.run.call_args[0][0] assert called_request.input == [Document.from_text('text1'), Document.from_text('text2')] diff --git a/py/packages/genkit/tests/genkit/blocks/formats/array_test.py b/py/packages/genkit/tests/genkit/ai/formats/array_test.py similarity index 70% rename from py/packages/genkit/tests/genkit/blocks/formats/array_test.py rename to py/packages/genkit/tests/genkit/ai/formats/array_test.py index 66895d2027..f847340a6b 100644 --- a/py/packages/genkit/tests/genkit/blocks/formats/array_test.py +++ b/py/packages/genkit/tests/genkit/ai/formats/array_test.py @@ -7,10 +7,10 @@ import pytest -from genkit.blocks.formats.array import ArrayFormat -from genkit.blocks.model import GenerateResponseChunkWrapper, MessageWrapper -from genkit.core.error import GenkitError -from genkit.core.typing import GenerateResponseChunk, Message, Part, TextPart +from genkit import Message, ModelResponseChunk +from genkit._ai._formats._array import ArrayFormat +from genkit._core._error import GenkitError +from genkit._core._typing import Part, TextPart class TestArrayFormatStreaming: @@ -22,18 +22,18 @@ def test_emits_complete_array_items_as_they_arrive(self) -> None: fmt = array_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) # Chunk 1: [{"id": 1, - chunk1 = GenerateResponseChunk(content=[Part(root=TextPart(text='[{"id": 1,'))]) - result1 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk1, index=0, previous_chunks=[])) + chunk1 = ModelResponseChunk(content=[Part(root=TextPart(text='[{"id": 1,'))]) + result1 = fmt.parse_chunk(ModelResponseChunk(chunk1, index=0, previous_chunks=[])) assert result1 == [] # Chunk 2: "name": "first"} - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text='"name": "first"}'))]) - result2 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk2, index=0, previous_chunks=[chunk1])) + chunk2 = ModelResponseChunk(content=[Part(root=TextPart(text='"name": "first"}'))]) + result2 = fmt.parse_chunk(ModelResponseChunk(chunk2, index=0, previous_chunks=[chunk1])) assert result2 == [{'id': 1, 'name': 'first'}] # Chunk 3: , {"id": 2, "name": "second"}] - chunk3 = GenerateResponseChunk(content=[Part(root=TextPart(text=', {"id": 2, "name": "second"}]'))]) - result3 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk3, index=0, previous_chunks=[chunk1, chunk2])) + chunk3 = ModelResponseChunk(content=[Part(root=TextPart(text=', {"id": 2, "name": "second"}]'))]) + result3 = fmt.parse_chunk(ModelResponseChunk(chunk3, index=0, previous_chunks=[chunk1, chunk2])) assert result3 == [{'id': 2, 'name': 'second'}] def test_handles_single_item_arrays(self) -> None: @@ -41,8 +41,8 @@ def test_handles_single_item_arrays(self) -> None: array_fmt = ArrayFormat() fmt = array_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) - chunk = GenerateResponseChunk(content=[Part(root=TextPart(text='[{"id": 1, "name": "single"}]'))]) - result = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk, index=0, previous_chunks=[])) + chunk = ModelResponseChunk(content=[Part(root=TextPart(text='[{"id": 1, "name": "single"}]'))]) + result = fmt.parse_chunk(ModelResponseChunk(chunk, index=0, previous_chunks=[])) assert result == [{'id': 1, 'name': 'single'}] def test_handles_preamble_with_code_fence(self) -> None: @@ -51,15 +51,15 @@ def test_handles_preamble_with_code_fence(self) -> None: fmt = array_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) # Chunk 1: preamble with code fence start - chunk1 = GenerateResponseChunk( + chunk1 = ModelResponseChunk( content=[Part(root=TextPart(text='Here is the array you requested:\n\n```json\n['))] ) - result1 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk1, index=0, previous_chunks=[])) + result1 = fmt.parse_chunk(ModelResponseChunk(chunk1, index=0, previous_chunks=[])) assert result1 == [] # Chunk 2: the actual data - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 1, "name": "item"}]\n```'))]) - result2 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk2, index=0, previous_chunks=[chunk1])) + chunk2 = ModelResponseChunk(content=[Part(root=TextPart(text='{"id": 1, "name": "item"}]\n```'))]) + result2 = fmt.parse_chunk(ModelResponseChunk(chunk2, index=0, previous_chunks=[chunk1])) assert result2 == [{'id': 1, 'name': 'item'}] @@ -72,7 +72,7 @@ def test_parses_complete_array_response(self) -> None: fmt = array_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) result = fmt.parse_message( - MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='[{"id": 1, "name": "test"}]'))])) + Message(Message(role='model', content=[Part(root=TextPart(text='[{"id": 1, "name": "test"}]'))])) ) assert result == [{'id': 1, 'name': 'test'}] @@ -81,7 +81,7 @@ def test_parses_empty_array(self) -> None: array_fmt = ArrayFormat() fmt = array_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='[]'))]))) + result = fmt.parse_message(Message(Message(role='model', content=[Part(root=TextPart(text='[]'))]))) assert result == [] def test_parses_array_with_preamble_and_code_fence(self) -> None: @@ -90,7 +90,7 @@ def test_parses_array_with_preamble_and_code_fence(self) -> None: fmt = array_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) result = fmt.parse_message( - MessageWrapper( + Message( Message( role='model', content=[Part(root=TextPart(text='Here is the array:\n\n```json\n[{"id": 1}]\n```'))] ) diff --git a/py/packages/genkit/tests/genkit/blocks/formats/enum_test.py b/py/packages/genkit/tests/genkit/ai/formats/enum_test.py similarity index 78% rename from py/packages/genkit/tests/genkit/blocks/formats/enum_test.py rename to py/packages/genkit/tests/genkit/ai/formats/enum_test.py index 822277be7f..2f83528260 100644 --- a/py/packages/genkit/tests/genkit/blocks/formats/enum_test.py +++ b/py/packages/genkit/tests/genkit/ai/formats/enum_test.py @@ -7,10 +7,10 @@ import pytest -from genkit.blocks.formats.enum import EnumFormat -from genkit.blocks.model import GenerateResponseChunkWrapper, MessageWrapper -from genkit.core.error import GenkitError -from genkit.core.typing import GenerateResponseChunk, Message, Part, TextPart +from genkit import Message, ModelResponseChunk +from genkit._ai._formats._enum import EnumFormat +from genkit._core._error import GenkitError +from genkit._core._typing import Part, TextPart class TestEnumFormatMessage: @@ -21,7 +21,7 @@ def test_parses_simple_enum_value(self) -> None: enum_fmt = EnumFormat() fmt = enum_fmt.handle({'type': 'string', 'enum': ['VALUE1', 'VALUE2']}) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='VALUE1'))]))) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text='VALUE1'))])) assert result == 'VALUE1' def test_trims_whitespace(self) -> None: @@ -29,9 +29,7 @@ def test_trims_whitespace(self) -> None: enum_fmt = EnumFormat() fmt = enum_fmt.handle({'type': 'string', 'enum': ['VALUE1', 'VALUE2']}) - result = fmt.parse_message( - MessageWrapper(Message(role='model', content=[Part(root=TextPart(text=' VALUE2\n'))])) - ) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text=' VALUE2\n'))])) assert result == 'VALUE2' def test_removes_double_quotes(self) -> None: @@ -39,7 +37,7 @@ def test_removes_double_quotes(self) -> None: enum_fmt = EnumFormat() fmt = enum_fmt.handle({'type': 'string', 'enum': ['foo', 'bar']}) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='"foo"'))]))) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text='"foo"'))])) assert result == 'foo' def test_removes_single_quotes(self) -> None: @@ -47,7 +45,7 @@ def test_removes_single_quotes(self) -> None: enum_fmt = EnumFormat() fmt = enum_fmt.handle({'type': 'string', 'enum': ['foo', 'bar']}) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text="'bar'"))]))) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text="'bar'"))])) assert result == 'bar' def test_handles_unquoted_value(self) -> None: @@ -55,7 +53,7 @@ def test_handles_unquoted_value(self) -> None: enum_fmt = EnumFormat() fmt = enum_fmt.handle({'type': 'string', 'enum': ['foo', 'bar']}) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='bar'))]))) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text='bar'))])) assert result == 'bar' @@ -67,11 +65,11 @@ def test_parses_accumulated_text_from_chunks(self) -> None: enum_fmt = EnumFormat() fmt = enum_fmt.handle({'type': 'string', 'enum': ['foo', 'bar']}) - chunk1 = GenerateResponseChunk(content=[Part(root=TextPart(text='"f'))]) - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text='oo"'))]) + chunk1 = ModelResponseChunk(content=[Part(TextPart(text='"f'))]) + chunk2 = ModelResponseChunk(content=[Part(TextPart(text='oo"'))]) result = fmt.parse_chunk( - GenerateResponseChunkWrapper( + ModelResponseChunk( chunk2, index=0, previous_chunks=[chunk1], diff --git a/py/packages/genkit/tests/genkit/blocks/formats/formats_test.py b/py/packages/genkit/tests/genkit/ai/formats/formats_test.py similarity index 94% rename from py/packages/genkit/tests/genkit/blocks/formats/formats_test.py rename to py/packages/genkit/tests/genkit/ai/formats/formats_test.py index 708567aaf7..6f8da2802b 100644 --- a/py/packages/genkit/tests/genkit/blocks/formats/formats_test.py +++ b/py/packages/genkit/tests/genkit/ai/formats/formats_test.py @@ -5,12 +5,11 @@ """Tests for the formats module initialization and built-in formats.""" -from genkit.blocks.formats import ( +from genkit._ai._formats import ( ArrayFormat, EnumFormat, FormatDef, Formatter, - FormatterConfig, JsonFormat, JsonlFormat, TextFormat, @@ -113,10 +112,6 @@ def test_formatter_exported(self) -> None: """Test that Formatter class is exported.""" assert Formatter is not None - def test_formatter_config_exported(self) -> None: - """Test that FormatterConfig class is exported.""" - assert FormatterConfig is not None - def test_all_format_classes_exported(self) -> None: """Test that all format classes are exported.""" assert ArrayFormat is not None diff --git a/py/packages/genkit/tests/genkit/blocks/formats/json_test.py b/py/packages/genkit/tests/genkit/ai/formats/json_test.py similarity index 65% rename from py/packages/genkit/tests/genkit/blocks/formats/json_test.py rename to py/packages/genkit/tests/genkit/ai/formats/json_test.py index 4f95959378..6e7383145f 100644 --- a/py/packages/genkit/tests/genkit/blocks/formats/json_test.py +++ b/py/packages/genkit/tests/genkit/ai/formats/json_test.py @@ -5,9 +5,9 @@ """Tests for the JSON format.""" -from genkit.blocks.formats import JsonFormat -from genkit.blocks.model import GenerateResponseChunkWrapper, MessageWrapper -from genkit.core.typing import GenerateResponseChunk, Message, Part, TextPart +from genkit import Message, ModelResponseChunk +from genkit._ai._formats import JsonFormat +from genkit._core._typing import Part, TextPart class TestJsonFormatStreaming: @@ -18,8 +18,8 @@ def test_parses_complete_json_object(self) -> None: json_fmt = JsonFormat() fmt = json_fmt.handle({'type': 'object'}) - chunk = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 1, "name": "test"}'))]) - result = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk, index=0, previous_chunks=[])) + chunk = ModelResponseChunk(content=[Part(TextPart(text='{"id": 1, "name": "test"}'))]) + result = fmt.parse_chunk(ModelResponseChunk(chunk, index=0, previous_chunks=[])) assert result == {'id': 1, 'name': 'test'} def test_handles_partial_json(self) -> None: @@ -28,13 +28,13 @@ def test_handles_partial_json(self) -> None: fmt = json_fmt.handle({'type': 'object'}) # Chunk 1: partial object - chunk1 = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 1'))]) - result1 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk1, index=0, previous_chunks=[])) + chunk1 = ModelResponseChunk(content=[Part(TextPart(text='{"id": 1'))]) + result1 = fmt.parse_chunk(ModelResponseChunk(chunk1, index=0, previous_chunks=[])) assert result1 == {'id': 1} # Chunk 2: complete object - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text=', "name": "test"}'))]) - result2 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk2, index=0, previous_chunks=[chunk1])) + chunk2 = ModelResponseChunk(content=[Part(TextPart(text=', "name": "test"}'))]) + result2 = fmt.parse_chunk(ModelResponseChunk(chunk2, index=0, previous_chunks=[chunk1])) assert result2 == {'id': 1, 'name': 'test'} def test_handles_preamble_with_code_fence(self) -> None: @@ -43,13 +43,13 @@ def test_handles_preamble_with_code_fence(self) -> None: fmt = json_fmt.handle({'type': 'object'}) # Chunk 1: preamble - chunk1 = GenerateResponseChunk(content=[Part(root=TextPart(text='Here is the JSON:\n\n```json\n'))]) - result1 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk1, index=0, previous_chunks=[])) + chunk1 = ModelResponseChunk(content=[Part(TextPart(text='Here is the JSON:\n\n```json\n'))]) + result1 = fmt.parse_chunk(ModelResponseChunk(chunk1, index=0, previous_chunks=[])) assert result1 is None # Chunk 2: actual data - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 1}\n```'))]) - result2 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk2, index=0, previous_chunks=[chunk1])) + chunk2 = ModelResponseChunk(content=[Part(TextPart(text='{"id": 1}\n```'))]) + result2 = fmt.parse_chunk(ModelResponseChunk(chunk2, index=0, previous_chunks=[chunk1])) assert result2 == {'id': 1} @@ -61,9 +61,7 @@ def test_parses_complete_json_response(self) -> None: json_fmt = JsonFormat() fmt = json_fmt.handle({'type': 'object'}) - result = fmt.parse_message( - MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='{"id": 1, "name": "test"}'))])) - ) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text='{"id": 1, "name": "test"}'))])) assert result == {'id': 1, 'name': 'test'} def test_handles_empty_response(self) -> None: @@ -71,7 +69,7 @@ def test_handles_empty_response(self) -> None: json_fmt = JsonFormat() fmt = json_fmt.handle({'type': 'object'}) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text=''))]))) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text=''))])) assert result is None def test_parses_json_with_preamble_and_code_fence(self) -> None: @@ -80,11 +78,7 @@ def test_parses_json_with_preamble_and_code_fence(self) -> None: fmt = json_fmt.handle({'type': 'object'}) result = fmt.parse_message( - MessageWrapper( - Message( - role='model', content=[Part(root=TextPart(text='Here is the JSON:\n\n```json\n{"id": 1}\n```'))] - ) - ) + Message(role='model', content=[Part(TextPart(text='Here is the JSON:\n\n```json\n{"id": 1}\n```'))]) ) assert result == {'id': 1} @@ -93,9 +87,7 @@ def test_parses_partial_json_message(self) -> None: json_fmt = JsonFormat() fmt = json_fmt.handle({'type': 'object'}) - result = fmt.parse_message( - MessageWrapper(Message(role='user', content=[Part(root=TextPart(text='{"foo": "bar"'))])) - ) + result = fmt.parse_message(Message(role='user', content=[Part(TextPart(text='{"foo": "bar"'))])) assert result == {'foo': 'bar'} def test_parses_complex_nested_json(self) -> None: @@ -104,14 +96,12 @@ def test_parses_complex_nested_json(self) -> None: fmt = json_fmt.handle({'type': 'object'}) result = fmt.parse_chunk( - GenerateResponseChunkWrapper( - GenerateResponseChunk(content=[Part(root=TextPart(text='", "baz": [1,2'))]), + ModelResponseChunk( + ModelResponseChunk(content=[Part(TextPart(text='", "baz": [1,2'))]), index=0, previous_chunks=[ - GenerateResponseChunk( - content=[Part(root=TextPart(text='{"bar":')), Part(root=TextPart(text='"ba'))] - ), - GenerateResponseChunk(content=[Part(root=TextPart(text='z'))]), + ModelResponseChunk(content=[Part(TextPart(text='{"bar":')), Part(TextPart(text='"ba'))]), + ModelResponseChunk(content=[Part(TextPart(text='z'))]), ], ) ) diff --git a/py/packages/genkit/tests/genkit/blocks/formats/jsonl_test.py b/py/packages/genkit/tests/genkit/ai/formats/jsonl_test.py similarity index 65% rename from py/packages/genkit/tests/genkit/blocks/formats/jsonl_test.py rename to py/packages/genkit/tests/genkit/ai/formats/jsonl_test.py index f3d53ec500..6a0df13034 100644 --- a/py/packages/genkit/tests/genkit/blocks/formats/jsonl_test.py +++ b/py/packages/genkit/tests/genkit/ai/formats/jsonl_test.py @@ -7,10 +7,10 @@ import pytest -from genkit.blocks.formats.jsonl import JsonlFormat -from genkit.blocks.model import GenerateResponseChunkWrapper, MessageWrapper -from genkit.core.error import GenkitError -from genkit.core.typing import GenerateResponseChunk, Message, Part, TextPart +from genkit import Message, ModelResponseChunk +from genkit._ai._formats._jsonl import JsonlFormat +from genkit._core._error import GenkitError +from genkit._core._typing import Part, TextPart class TestJsonlFormatStreaming: @@ -22,18 +22,18 @@ def test_emits_complete_json_objects_as_they_arrive(self) -> None: fmt = jsonl_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) # Chunk 1: first complete object - chunk1 = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 1, "name": "first"}\n'))]) - result1 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk1, index=0, previous_chunks=[])) + chunk1 = ModelResponseChunk(content=[Part(TextPart(text='{"id": 1, "name": "first"}\n'))]) + result1 = fmt.parse_chunk(ModelResponseChunk(chunk1, index=0, previous_chunks=[])) assert result1 == [{'id': 1, 'name': 'first'}] # Chunk 2: second object complete, third starts - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 2, "name": "second"}\n{"id": 3'))]) - result2 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk2, index=0, previous_chunks=[chunk1])) + chunk2 = ModelResponseChunk(content=[Part(TextPart(text='{"id": 2, "name": "second"}\n{"id": 3'))]) + result2 = fmt.parse_chunk(ModelResponseChunk(chunk2, index=0, previous_chunks=[chunk1])) assert result2 == [{'id': 2, 'name': 'second'}] # Chunk 3: third object completes - chunk3 = GenerateResponseChunk(content=[Part(root=TextPart(text=', "name": "third"}\n'))]) - result3 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk3, index=0, previous_chunks=[chunk1, chunk2])) + chunk3 = ModelResponseChunk(content=[Part(TextPart(text=', "name": "third"}\n'))]) + result3 = fmt.parse_chunk(ModelResponseChunk(chunk3, index=0, previous_chunks=[chunk1, chunk2])) assert result3 == [{'id': 3, 'name': 'third'}] def test_handles_single_object(self) -> None: @@ -41,8 +41,8 @@ def test_handles_single_object(self) -> None: jsonl_fmt = JsonlFormat() fmt = jsonl_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) - chunk = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 1, "name": "single"}\n'))]) - result = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk, index=0, previous_chunks=[])) + chunk = ModelResponseChunk(content=[Part(TextPart(text='{"id": 1, "name": "single"}\n'))]) + result = fmt.parse_chunk(ModelResponseChunk(chunk, index=0, previous_chunks=[])) assert result == [{'id': 1, 'name': 'single'}] def test_handles_preamble_with_code_fence(self) -> None: @@ -51,13 +51,13 @@ def test_handles_preamble_with_code_fence(self) -> None: fmt = jsonl_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) # Chunk 1: preamble - chunk1 = GenerateResponseChunk(content=[Part(root=TextPart(text='Here are the objects:\n\n```\n'))]) - result1 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk1, index=0, previous_chunks=[])) + chunk1 = ModelResponseChunk(content=[Part(TextPart(text='Here are the objects:\n\n```\n'))]) + result1 = fmt.parse_chunk(ModelResponseChunk(chunk1, index=0, previous_chunks=[])) assert result1 == [] # Chunk 2: actual data - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text='{"id": 1, "name": "item"}\n```'))]) - result2 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk2, index=0, previous_chunks=[chunk1])) + chunk2 = ModelResponseChunk(content=[Part(TextPart(text='{"id": 1, "name": "item"}\n```'))]) + result2 = fmt.parse_chunk(ModelResponseChunk(chunk2, index=0, previous_chunks=[chunk1])) assert result2 == [{'id': 1, 'name': 'item'}] def test_ignores_non_object_lines(self) -> None: @@ -65,10 +65,10 @@ def test_ignores_non_object_lines(self) -> None: jsonl_fmt = JsonlFormat() fmt = jsonl_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) - chunk = GenerateResponseChunk( - content=[Part(root=TextPart(text='First object:\n{"id": 1}\nSecond object:\n{"id": 2}\n'))] + chunk = ModelResponseChunk( + content=[Part(TextPart(text='First object:\n{"id": 1}\nSecond object:\n{"id": 2}\n'))] ) - result = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk, index=0, previous_chunks=[])) + result = fmt.parse_chunk(ModelResponseChunk(chunk, index=0, previous_chunks=[])) assert result == [{'id': 1}, {'id': 2}] @@ -81,9 +81,7 @@ def test_parses_complete_jsonl_response(self) -> None: fmt = jsonl_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) result = fmt.parse_message( - MessageWrapper( - Message(role='model', content=[Part(root=TextPart(text='{"id": 1, "name": "test"}\n{"id": 2}\n'))]) - ) + Message(role='model', content=[Part(TextPart(text='{"id": 1, "name": "test"}\n{"id": 2}\n'))]) ) assert result == [{'id': 1, 'name': 'test'}, {'id': 2}] @@ -92,7 +90,7 @@ def test_handles_empty_response(self) -> None: jsonl_fmt = JsonlFormat() fmt = jsonl_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text=''))]))) + result = fmt.parse_message(Message(role='model', content=[Part(TextPart(text=''))])) assert result == [] def test_parses_jsonl_with_preamble_and_code_fence(self) -> None: @@ -101,11 +99,9 @@ def test_parses_jsonl_with_preamble_and_code_fence(self) -> None: fmt = jsonl_fmt.handle({'type': 'array', 'items': {'type': 'object'}}) result = fmt.parse_message( - MessageWrapper( - Message( - role='model', - content=[Part(root=TextPart(text='Here are the objects:\n\n```\n{"id": 1}\n{"id": 2}\n```'))], - ) + Message( + role='model', + content=[Part(TextPart(text='Here are the objects:\n\n```\n{"id": 1}\n{"id": 2}\n```'))], ) ) assert result == [{'id': 1}, {'id': 2}] diff --git a/py/packages/genkit/tests/genkit/blocks/formats/text_test.py b/py/packages/genkit/tests/genkit/ai/formats/text_test.py similarity index 68% rename from py/packages/genkit/tests/genkit/blocks/formats/text_test.py rename to py/packages/genkit/tests/genkit/ai/formats/text_test.py index c4c16d9c86..45b01b313e 100644 --- a/py/packages/genkit/tests/genkit/blocks/formats/text_test.py +++ b/py/packages/genkit/tests/genkit/ai/formats/text_test.py @@ -5,9 +5,9 @@ """Tests for the Text format.""" -from genkit.blocks.formats.text import TextFormat -from genkit.blocks.model import GenerateResponseChunkWrapper, MessageWrapper -from genkit.core.typing import GenerateResponseChunk, Message, Part, TextPart +from genkit import Message, ModelResponseChunk +from genkit._ai._formats._text import TextFormat +from genkit._core._typing import Part, TextPart class TestTextFormatStreaming: @@ -19,13 +19,13 @@ def test_emits_text_chunks_as_they_arrive(self) -> None: fmt = text_fmt.handle(None) # Chunk 1: "Hello" - chunk1 = GenerateResponseChunk(content=[Part(root=TextPart(text='Hello'))]) - result1 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk1, index=0, previous_chunks=[])) + chunk1 = ModelResponseChunk(content=[Part(root=TextPart(text='Hello'))]) + result1 = fmt.parse_chunk(ModelResponseChunk(chunk1, index=0, previous_chunks=[])) assert result1 == 'Hello' # Chunk 2: " world" - should return only this chunk's text, not accumulated - chunk2 = GenerateResponseChunk(content=[Part(root=TextPart(text=' world'))]) - result2 = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk2, index=0, previous_chunks=[chunk1])) + chunk2 = ModelResponseChunk(content=[Part(root=TextPart(text=' world'))]) + result2 = fmt.parse_chunk(ModelResponseChunk(chunk2, index=0, previous_chunks=[chunk1])) assert result2 == ' world' def test_handles_empty_chunks(self) -> None: @@ -33,8 +33,8 @@ def test_handles_empty_chunks(self) -> None: text_fmt = TextFormat() fmt = text_fmt.handle(None) - chunk = GenerateResponseChunk(content=[Part(root=TextPart(text=''))]) - result = fmt.parse_chunk(GenerateResponseChunkWrapper(chunk, index=0, previous_chunks=[])) + chunk = ModelResponseChunk(content=[Part(root=TextPart(text=''))]) + result = fmt.parse_chunk(ModelResponseChunk(chunk, index=0, previous_chunks=[])) assert result == '' @@ -46,9 +46,7 @@ def test_parses_complete_text_response(self) -> None: text_fmt = TextFormat() fmt = text_fmt.handle(None) - result = fmt.parse_message( - MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='Hello world'))])) - ) + result = fmt.parse_message(Message(Message(role='model', content=[Part(root=TextPart(text='Hello world'))]))) assert result == 'Hello world' def test_handles_empty_response(self) -> None: @@ -56,7 +54,7 @@ def test_handles_empty_response(self) -> None: text_fmt = TextFormat() fmt = text_fmt.handle(None) - result = fmt.parse_message(MessageWrapper(Message(role='model', content=[Part(root=TextPart(text=''))]))) + result = fmt.parse_message(Message(Message(role='model', content=[Part(root=TextPart(text=''))]))) assert result == '' def test_handles_multiline_text(self) -> None: @@ -65,7 +63,7 @@ def test_handles_multiline_text(self) -> None: fmt = text_fmt.handle(None) result = fmt.parse_message( - MessageWrapper(Message(role='model', content=[Part(root=TextPart(text='Line 1\nLine 2\nLine 3'))])) + Message(Message(role='model', content=[Part(root=TextPart(text='Line 1\nLine 2\nLine 3'))])) ) assert result == 'Line 1\nLine 2\nLine 3' diff --git a/py/packages/genkit/tests/genkit/ai/generate_operation_test.py b/py/packages/genkit/tests/genkit/ai/generate_operation_test.py index 10d61bd64a..d3cdd75f5e 100644 --- a/py/packages/genkit/tests/genkit/ai/generate_operation_test.py +++ b/py/packages/genkit/tests/genkit/ai/generate_operation_test.py @@ -47,13 +47,11 @@ import pytest -from genkit.ai import Genkit -from genkit.core.action import ActionRunContext -from genkit.core.error import GenkitError -from genkit.core.typing import ( - GenerateRequest, - GenerateResponse, - Message, +from genkit import Genkit, Message, ModelResponse +from genkit._core._action import ActionRunContext +from genkit._core._error import GenkitError +from genkit._core._model import ModelRequest +from genkit._core._typing import ( ModelInfo, Operation, Part, @@ -95,8 +93,8 @@ async def test_generate_operation_model_no_long_running_support(ai: Genkit) -> N """ # Define a standard model without long_running support - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse( + async def model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Hello'))], @@ -127,8 +125,8 @@ async def test_generate_operation_model_no_supports_info(ai: Genkit) -> None: """Test that models without supports info are rejected.""" # Define a model without any ModelInfo - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse( + async def model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Hello'))], @@ -151,9 +149,9 @@ async def test_generate_operation_no_operation_returned(ai: Genkit) -> None: """ # Define a model that claims to support long_running but doesn't return an operation - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: # Return a normal response without an operation - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Hello'))], @@ -186,8 +184,8 @@ async def test_generate_operation_success_with_lro_model(ai: Genkit) -> None: ) # Define a model that supports long_running and returns an operation - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse( + async def model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Started'))], @@ -221,8 +219,8 @@ async def test_generate_operation_with_default_model(ai: Genkit) -> None: done=False, ) - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse( + async def model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Started'))], @@ -262,13 +260,13 @@ def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateRespons @pytest.mark.asyncio async def test_generate_operation_passes_all_options(ai: Genkit) -> None: """Test that generate_operation passes all options to generate().""" - captured_request: GenerateRequest | None = None + captured_request: ModelRequest | None = None expected_operation = Operation(id='opt-test-789', done=False) - def model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: nonlocal captured_request captured_request = request - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Started'))], diff --git a/py/packages/genkit/tests/genkit/blocks/generate_test.py b/py/packages/genkit/tests/genkit/ai/generate_test.py similarity index 65% rename from py/packages/genkit/tests/genkit/blocks/generate_test.py rename to py/packages/genkit/tests/genkit/ai/generate_test.py index 17fb0a4d55..be417a6878 100644 --- a/py/packages/genkit/tests/genkit/blocks/generate_test.py +++ b/py/packages/genkit/tests/genkit/ai/generate_test.py @@ -5,42 +5,55 @@ """Tests for the action module.""" +import json import pathlib -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Sequence from typing import Any, cast import pytest import yaml -from pydantic import TypeAdapter - -from genkit.ai import ActionKind, Genkit -from genkit.blocks.generate import generate_action -from genkit.blocks.model import ( - ModelMiddlewareNext, - text_from_content, - text_from_message, +from pydantic import BaseModel, TypeAdapter + +from genkit import ActionKind, Document, Genkit, Message, ModelResponse, ModelResponseChunk +from genkit._ai._generate import generate_action +from genkit._ai._model import text_from_content, text_from_message +from genkit._ai._testing import ( + ProgrammableModel, + define_echo_model, + define_programmable_model, ) -from genkit.codec import dump_dict, dump_json -from genkit.core.action import ActionRunContext -from genkit.core.typing import ( - DocumentData, +from genkit._core._action import ActionRunContext +from genkit._core._model import ModelRequest +from genkit._core._typing import ( DocumentPart, FinishReason, GenerateActionOptions, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - Message, Metadata, Part, Role, TextPart, ) -from genkit.testing import ( - ProgrammableModel, - define_echo_model, - define_programmable_model, -) + + +def _to_dict(obj: object) -> object: + """Convert object to dict for test comparisons.""" + if isinstance(obj, BaseModel): + return obj.model_dump() + if isinstance(obj, list): + return [_to_dict(item) for item in obj] + if isinstance(obj, dict): + return {k: _to_dict(v) for k, v in obj.items()} + return obj + + +def _to_json(obj: object, indent: int | None = None) -> str: + """Local test helper: serialize to JSON for assertion error messages. + + Uses model_dump_json for BaseModel, json.dumps for dicts/other. + """ + if isinstance(obj, BaseModel): + return obj.model_dump_json(indent=indent) + return json.dumps(obj, indent=indent) @pytest.fixture @@ -51,7 +64,7 @@ def setup_test() -> tuple[Genkit, ProgrammableModel]: pm, _ = define_programmable_model(ai) @ai.tool(name='testTool') - def test_tool() -> object: + async def test_tool() -> object: """description""" # noqa: D403, D415 return 'tool called' @@ -66,9 +79,9 @@ async def test_simple_text_generate_request( ai, pm = setup_test pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, - message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='bye'))]), + message=Message(role=Role.MODEL, content=[Part(TextPart(text='bye'))]), ) ) @@ -79,7 +92,7 @@ async def test_simple_text_generate_request( messages=[ Message( role=Role.USER, - content=[Part(root=TextPart(text='hi'))], + content=[Part(TextPart(text='hi'))], ), ], ), @@ -96,9 +109,9 @@ async def test_simulates_doc_grounding( ai, pm = setup_test pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, - message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='bye'))]), + message=Message(role=Role.MODEL, content=[Part(TextPart(text='bye'))]), ) ) @@ -109,10 +122,10 @@ async def test_simulates_doc_grounding( messages=[ Message( role=Role.USER, - content=[Part(root=TextPart(text='hi'))], + content=[Part(TextPart(text='hi'))], ), ], - docs=[DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))])], + docs=[Document(content=[DocumentPart(TextPart(text='doc content 1'))])], ), ) @@ -121,7 +134,7 @@ async def test_simulates_doc_grounding( assert response.request.messages[0] == Message( role=Role.USER, content=[ - Part(root=TextPart(text='hi')), + Part(TextPart(text='hi')), Part( root=TextPart( text='\n\nUse the following information to complete your task:' + '\n\n- [0]: doc content 1\n\n', @@ -141,31 +154,31 @@ async def test_generate_applies_middleware( define_echo_model(ai) async def pre_middle( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next: ModelMiddlewareNext, - ) -> GenerateResponse: + next: Callable[..., Awaitable[ModelResponse]], + ) -> ModelResponse: txt = ''.join(text_from_message(m) for m in req.messages) return await next( - GenerateRequest( + ModelRequest( messages=[ - Message(role=Role.USER, content=[Part(root=TextPart(text=f'PRE {txt}'))]), + Message(role=Role.USER, content=[Part(TextPart(text=f'PRE {txt}'))]), ], ), ctx, ) async def post_middle( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next: ModelMiddlewareNext, - ) -> GenerateResponse: - resp: GenerateResponse = await next(req, ctx) + next: Callable[..., Awaitable[ModelResponse]], + ) -> ModelResponse: + resp: ModelResponse = await next(req, ctx) assert resp.message is not None txt = text_from_message(resp.message) - return GenerateResponse( + return ModelResponse( finish_reason=resp.finish_reason, - message=Message(role=Role.USER, content=[Part(root=TextPart(text=f'{txt} POST'))]), + message=Message(role=Role.USER, content=[Part(TextPart(text=f'{txt} POST'))]), ) response = await generate_action( @@ -175,7 +188,7 @@ async def post_middle( messages=[ Message( role=Role.USER, - content=[Part(root=TextPart(text='hi'))], + content=[Part(TextPart(text='hi'))], ), ], ), @@ -194,16 +207,16 @@ async def test_generate_middleware_next_fn_args_optional( define_echo_model(ai) async def post_middle( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next: ModelMiddlewareNext, - ) -> GenerateResponse: - resp: GenerateResponse = await next(req, ctx) + next: Callable[..., Awaitable[ModelResponse]], + ) -> ModelResponse: + resp: ModelResponse = await next(req, ctx) assert resp.message is not None txt = text_from_message(resp.message) - return GenerateResponse( + return ModelResponse( finish_reason=resp.finish_reason, - message=Message(role=Role.USER, content=[Part(root=TextPart(text=f'{txt} POST'))]), + message=Message(role=Role.USER, content=[Part(TextPart(text=f'{txt} POST'))]), ) response = await generate_action( @@ -213,7 +226,7 @@ async def post_middle( messages=[ Message( role=Role.USER, - content=[Part(root=TextPart(text='hi'))], + content=[Part(TextPart(text='hi'))], ), ], ), @@ -232,24 +245,24 @@ async def test_generate_middleware_can_modify_context( define_echo_model(ai) async def add_context( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next: ModelMiddlewareNext, - ) -> GenerateResponse: + next: Callable[..., Awaitable[ModelResponse]], + ) -> ModelResponse: return await next(req, ActionRunContext(context={**ctx.context, 'banana': True})) async def inject_context( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next: ModelMiddlewareNext, - ) -> GenerateResponse: + next: Callable[..., Awaitable[ModelResponse]], + ) -> ModelResponse: txt = ''.join(text_from_message(m) for m in req.messages) return await next( - GenerateRequest( + ModelRequest( messages=[ Message( role=Role.USER, - content=[Part(root=TextPart(text=f'{txt} {ctx.context}'))], + content=[Part(TextPart(text=f'{txt} {ctx.context}'))], ), ], ), @@ -263,7 +276,7 @@ async def inject_context( messages=[ Message( role=Role.USER, - content=[Part(root=TextPart(text='hi'))], + content=[Part(TextPart(text='hi'))], ), ], ), @@ -282,52 +295,56 @@ async def test_generate_middleware_can_modify_stream( ai, pm = setup_test pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, - message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='bye'))]), + message=Message(role=Role.MODEL, content=[Part(TextPart(text='bye'))]), ) ) pm.chunks = [ [ - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='1'))]), - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='2'))]), - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='3'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(TextPart(text='1'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(TextPart(text='2'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(TextPart(text='3'))]), ] ] async def modify_stream( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next: ModelMiddlewareNext, - ) -> GenerateResponse: - ctx.send_chunk( - GenerateResponseChunk( - role=Role.MODEL, - content=[Part(root=TextPart(text='something extra before'))], - ) - ) - - def chunk_handler(chunk: object) -> None: - assert isinstance(chunk, GenerateResponseChunk) - ctx.send_chunk( - GenerateResponseChunk( + on_chunk: Callable[[ModelResponseChunk], None] | None, + next: Callable[..., Awaitable[ModelResponse]], + ) -> ModelResponse: + # 4-param streaming middleware signature + if on_chunk: + on_chunk( + ModelResponseChunk( role=Role.MODEL, - content=[Part(root=TextPart(text=f'intercepted: {text_from_content(chunk.content)}'))], + content=[Part(TextPart(text='something extra before'))], ) ) - resp = await next(req, ActionRunContext(context=ctx.context, on_chunk=chunk_handler)) - ctx.send_chunk( - GenerateResponseChunk( - role=Role.MODEL, - content=[Part(root=TextPart(text='something extra after'))], + def chunk_handler(chunk: ModelResponseChunk) -> None: + if on_chunk: + on_chunk( + ModelResponseChunk( + role=Role.MODEL, + content=[Part(TextPart(text=f'intercepted: {text_from_content(chunk.content)}'))], + ) + ) + + resp = await next(req, ctx, chunk_handler) + if on_chunk: + on_chunk( + ModelResponseChunk( + role=Role.MODEL, + content=[Part(TextPart(text='something extra after'))], + ) ) - ) return resp got_chunks = [] - def collect_chunks(c: GenerateResponseChunk) -> None: + def collect_chunks(c: ModelResponseChunk) -> None: got_chunks.append(text_from_content(c.content)) response = await generate_action( @@ -337,7 +354,7 @@ def collect_chunks(c: GenerateResponseChunk) -> None: messages=[ Message( role=Role.USER, - content=[Part(root=TextPart(text='hi'))], + content=[Part(TextPart(text='hi'))], ), ], ), @@ -379,12 +396,12 @@ async def test_generate_action_spec(spec: dict[str, Any]) -> None: pm, _ = define_programmable_model(ai) @ai.tool(name='testTool') - def test_tool() -> object: + async def test_tool() -> object: """description""" # noqa: D403, D415 return 'tool called' if 'modelResponses' in spec: - pm.responses = [TypeAdapter(GenerateResponse).validate_python(resp) for resp in spec['modelResponses']] + pm.responses = [TypeAdapter(ModelResponse).validate_python(resp) for resp in spec['modelResponses']] if 'streamChunks' in spec: pm.chunks = [] @@ -392,29 +409,29 @@ def test_tool() -> object: converted = [] if stream_chunks: for chunk in stream_chunks: - converted.append(TypeAdapter(GenerateResponseChunk).validate_python(chunk)) + converted.append(TypeAdapter(ModelResponseChunk).validate_python(chunk)) pm.chunks.append(converted) action = await ai.registry.resolve_action(kind=ActionKind.UTIL, name='generate') assert action is not None response = None - chunks: list[GenerateResponseChunk] | None = None + chunks: list[ModelResponseChunk] | None = None if spec.get('stream'): chunks = [] captured_chunks = chunks # Capture list reference for closure - def on_chunk(chunk: GenerateResponseChunk) -> None: + def on_chunk(chunk: ModelResponseChunk) -> None: captured_chunks.append(chunk) - action_response = await action.arun( + action_response = await action.run( ai.registry, TypeAdapter(GenerateActionOptions).validate_python(spec['input']), # type: ignore[arg-type] on_chunk=on_chunk, # type: ignore[misc] ) response = action_response.response else: - action_response = await action.arun( + action_response = await action.run( TypeAdapter(GenerateActionOptions).validate_python(spec['input']), ) response = action_response.response @@ -425,15 +442,15 @@ def on_chunk(chunk: GenerateResponseChunk) -> None: assert isinstance(got, list) and isinstance(want, list) if not is_equal_lists(got, want): raise AssertionError( - f'{dump_json(got, indent=2)}\n\nis not equal to expected:\n\n{dump_json(want, indent=2)}' + f'{_to_json(got, indent=2)}\n\nis not equal to expected:\n\n{_to_json(want, indent=2)}' ) if 'expectResponse' in spec: - got = clean_schema(dump_dict(response)) + got = clean_schema(_to_dict(response)) want = clean_schema(spec['expectResponse']) if got != want: raise AssertionError( - f'{dump_json(got, indent=2)}\n\nis not equal to expected:\n\n{dump_json(want, indent=2)}' + f'{_to_json(got, indent=2)}\n\nis not equal to expected:\n\n{_to_json(want, indent=2)}' ) @@ -442,7 +459,7 @@ def is_equal_lists(a: Sequence[object], b: Sequence[object]) -> bool: if len(a) != len(b): return False - return all(dump_dict(a[i]) == dump_dict(b[i]) for i in range(len(a))) + return all(_to_dict(a[i]) == _to_dict(b[i]) for i in range(len(a))) primitives = (bool, str, int, float, type(None)) diff --git a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py index 090e5eec0f..c47a4e6cd0 100644 --- a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py +++ b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py @@ -5,7 +5,6 @@ """Tests for the Genkit extra API methods.""" -from typing import Any from unittest import mock from unittest.mock import AsyncMock, MagicMock @@ -13,13 +12,9 @@ from opentelemetry import trace as trace_api from opentelemetry.sdk.trace import TracerProvider -from genkit.ai import Genkit -from genkit.ai._registry import SimpleRetrieverOptions -from genkit.core.action import Action -from genkit.core.action._action import _action_context -from genkit.core.action.types import ActionKind -from genkit.core.typing import DocumentPart, Operation -from genkit.types import DocumentData, RetrieverRequest, RetrieverResponse, TextPart +from genkit import Genkit +from genkit._core._action import Action, ActionKind, _action_context +from genkit._core._typing import Operation @pytest.mark.asyncio @@ -27,28 +22,22 @@ async def test_genkit_run() -> None: """Test Genkit.run method.""" ai = Genkit() - def sync_fn() -> str: - return 'hello' - async def async_fn() -> str: return 'world' - res1 = await ai.run('test1', sync_fn) - assert res1 == 'hello' - - res2 = await ai.run('test2', async_fn) - assert res2 == 'world' + res1 = await ai.run(name='test1', fn=async_fn) + assert res1 == 'world' # Test with metadata - res3 = await ai.run('test3', sync_fn, metadata={'foo': 'bar'}) - assert res3 == 'hello' + res2 = await ai.run(name='test2', fn=async_fn, metadata={'foo': 'bar'}) + assert res2 == 'world' - # Test with input overload - async def multiply(x: int) -> int: - return x * 2 + # Test that sync functions raise TypeError + def sync_fn() -> str: + return 'hello' - res4 = await ai.run('multiply', 10, multiply) - assert res4 == 20 + with pytest.raises(TypeError, match='fn must be a coroutine function'): + await ai.run(name='test3', fn=sync_fn) # type: ignore[arg-type] @pytest.mark.asyncio @@ -56,10 +45,10 @@ async def test_genkit_dynamic_tool() -> None: """Test Genkit.dynamic_tool method.""" ai = Genkit() - def my_tool(x: int) -> int: + async def my_tool(x: int) -> int: return x + 1 - tool = ai.dynamic_tool('my_tool', my_tool, description='increment x') + tool = ai.dynamic_tool(name='my_tool', fn=my_tool, description='increment x') assert isinstance(tool, Action) assert tool.kind == ActionKind.TOOL @@ -69,7 +58,7 @@ def my_tool(x: int) -> int: assert tool.metadata.get('dynamic') is True # Execution - resp = await tool.arun(5) + resp = await tool.run(5) assert resp.response == 6 @@ -86,7 +75,7 @@ async def test_genkit_check_operation() -> None: # Patch lookup_background_action to return our mock with mock.patch( - 'genkit.blocks.background_model.lookup_background_action', + 'genkit._core._background.lookup_background_action', new=AsyncMock(return_value=mock_background_action), ) as mock_lookup: updated_op = await ai.check_operation(op) @@ -117,56 +106,6 @@ async def test_genkit_check_operation_not_found() -> None: await ai.check_operation(op) -@pytest.mark.asyncio -async def test_define_simple_retriever_legacy() -> None: - """Test define_simple_retriever with legacy handler signature.""" - ai = Genkit() - - def simple_retriever(query: DocumentData, options: Any) -> list[DocumentData]: # noqa: ANN401 - # Returns list[DocumentData] directly - - text_part: DocumentPart = DocumentPart(root=TextPart(text='doc1')) - return [DocumentData(content=[text_part])] - - retriever_action = ai.define_simple_retriever('simple', simple_retriever) - - assert retriever_action.kind == ActionKind.RETRIEVER - - # Test execution - req = RetrieverRequest(query=DocumentData(content=[])) - resp_wrapper = await retriever_action.arun(req) - response = resp_wrapper.response - - assert isinstance(response, RetrieverResponse) - assert len(response.documents) == 1 - assert response.documents[0].content[0].root.text == 'doc1' - - -@pytest.mark.asyncio -async def test_define_simple_retriever_mapped() -> None: - """Test define_simple_retriever with mapping options.""" - ai = Genkit() - - def db_handler(query: DocumentData, options: Any) -> list[dict[str, Any]]: # noqa: ANN401 - return [ - {'id': '1', 'text': 'hello', 'extra': 'data'}, - {'id': '2', 'text': 'world', 'extra': 'more'}, - ] - - options = SimpleRetrieverOptions(name='mapped', content='text', metadata=['extra']) - - retriever_action = ai.define_simple_retriever(options, db_handler) - - req = RetrieverRequest(query=DocumentData(content=[])) - resp_wrapper = await retriever_action.arun(req) - response = resp_wrapper.response - - assert len(response.documents) == 2 - assert response.documents[0].content[0].root.text == 'hello' - assert response.documents[0].metadata == {'extra': 'data'} - assert 'id' not in response.documents[0].metadata - - @pytest.mark.asyncio async def test_current_context() -> None: """Test Genkit.current_context method.""" diff --git a/py/packages/genkit/tests/genkit/blocks/message_utils_test.py b/py/packages/genkit/tests/genkit/ai/message_utils_test.py similarity index 97% rename from py/packages/genkit/tests/genkit/blocks/message_utils_test.py rename to py/packages/genkit/tests/genkit/ai/message_utils_test.py index 9366d06bc0..9ebcc6b607 100644 --- a/py/packages/genkit/tests/genkit/blocks/message_utils_test.py +++ b/py/packages/genkit/tests/genkit/ai/message_utils_test.py @@ -5,9 +5,9 @@ """Tests for the message utils.""" -from genkit.blocks.messages import inject_instructions -from genkit.core.typing import ( - Message, +from genkit import Message +from genkit._ai._messages import inject_instructions +from genkit._core._typing import ( Metadata, Part, Role, diff --git a/py/packages/genkit/tests/genkit/blocks/middleware_test.py b/py/packages/genkit/tests/genkit/ai/middleware_test.py similarity index 76% rename from py/packages/genkit/tests/genkit/blocks/middleware_test.py rename to py/packages/genkit/tests/genkit/ai/middleware_test.py index 060385ff7d..37d919d977 100644 --- a/py/packages/genkit/tests/genkit/blocks/middleware_test.py +++ b/py/packages/genkit/tests/genkit/ai/middleware_test.py @@ -21,14 +21,12 @@ import pytest -from genkit.blocks.middleware import augment_with_context -from genkit.core.action import ActionRunContext -from genkit.core.typing import ( - DocumentData, +from genkit import Document, Message, ModelResponse +from genkit._ai._middleware import augment_with_context +from genkit._core._action import ActionRunContext +from genkit._core._model import ModelRequest +from genkit._core._typing import ( DocumentPart, - GenerateRequest, - GenerateResponse, - Message, Metadata, Part, Role, @@ -36,14 +34,14 @@ ) -async def run_augmenter(req: GenerateRequest) -> GenerateRequest: +async def run_augmenter(req: ModelRequest) -> ModelRequest: """Helper to run the augment_with_context middleware.""" augmenter = augment_with_context() req_future = asyncio.Future() - async def next(req: GenerateRequest, _: ActionRunContext) -> GenerateResponse: + async def next(req: ModelRequest, _: ActionRunContext) -> ModelResponse: req_future.set_result(req) - return GenerateResponse(message=Message(role=Role.USER, content=[Part(root=TextPart(text='hi'))])) + return ModelResponse(message=Message(role=Role.USER, content=[Part(root=TextPart(text='hi'))])) await augmenter(req, ActionRunContext(), next) @@ -53,7 +51,7 @@ async def next(req: GenerateRequest, _: ActionRunContext) -> GenerateResponse: @pytest.mark.asyncio async def test_augment_with_context_ignores_no_docs() -> None: """Test simple prompt rendering.""" - req = GenerateRequest( + req = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='hi'))]), ], @@ -67,19 +65,19 @@ async def test_augment_with_context_ignores_no_docs() -> None: @pytest.mark.asyncio async def test_augment_with_context_adds_docs_as_context() -> None: """Test simple prompt rendering.""" - req = GenerateRequest( + req = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='hi'))]), ], docs=[ - DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='doc content 2'))]), + Document(content=[DocumentPart(root=TextPart(text='doc content 1'))]), + Document(content=[DocumentPart(root=TextPart(text='doc content 2'))]), ], ) transformed_req = await run_augmenter(req) - assert transformed_req == GenerateRequest( + assert transformed_req == ModelRequest( messages=[ Message( role=Role.USER, @@ -98,8 +96,8 @@ async def test_augment_with_context_adds_docs_as_context() -> None: ) ], docs=[ - DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='doc content 2'))]), + Document(content=[DocumentPart(root=TextPart(text='doc content 1'))]), + Document(content=[DocumentPart(root=TextPart(text='doc content 2'))]), ], ) @@ -107,7 +105,7 @@ async def test_augment_with_context_adds_docs_as_context() -> None: @pytest.mark.asyncio async def test_augment_with_context_should_not_modify_non_pending_part() -> None: """Test simple prompt rendering.""" - req = GenerateRequest( + req = ModelRequest( messages=[ Message( role=Role.USER, @@ -123,7 +121,7 @@ async def test_augment_with_context_should_not_modify_non_pending_part() -> None ), ], docs=[ - DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))]), + Document(content=[DocumentPart(root=TextPart(text='doc content 1'))]), ], ) @@ -135,7 +133,7 @@ async def test_augment_with_context_should_not_modify_non_pending_part() -> None @pytest.mark.asyncio async def test_augment_with_context_with_purpose_part() -> None: """Test simple prompt rendering.""" - req = GenerateRequest( + req = ModelRequest( messages=[ Message( role=Role.USER, @@ -151,13 +149,13 @@ async def test_augment_with_context_with_purpose_part() -> None: ), ], docs=[ - DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))]), + Document(content=[DocumentPart(root=TextPart(text='doc content 1'))]), ], ) transformed_req = await run_augmenter(req) - assert transformed_req == GenerateRequest( + assert transformed_req == ModelRequest( messages=[ Message( role=Role.USER, @@ -175,6 +173,6 @@ async def test_augment_with_context_with_purpose_part() -> None: ) ], docs=[ - DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))]), + Document(content=[DocumentPart(root=TextPart(text='doc content 1'))]), ], ) diff --git a/py/packages/genkit/tests/genkit/ai/model_test.py b/py/packages/genkit/tests/genkit/ai/model_test.py new file mode 100644 index 0000000000..6699516918 --- /dev/null +++ b/py/packages/genkit/tests/genkit/ai/model_test.py @@ -0,0 +1,379 @@ +#!/usr/bin/env python3 +# +# Copyright 2025 Google LLC +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the action module.""" + +import pytest + +from genkit import Message, ModelRequest, ModelResponse, ModelResponseChunk, ModelUsage +from genkit._ai._model import text_from_content +from genkit._core._action import ActionMetadata +from genkit._core._typing import ( + DocumentPart, + Media, + MediaPart, + Metadata, + Part, + TextPart, + ToolRequest, + ToolRequestPart, +) +from genkit.model import get_basic_usage_stats, model_action_metadata + + +def test_message_wrapper_text() -> None: + """Test text property of Message.""" + wrapper = Message( + Message( + role='model', + content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))], + ), + ) + + assert wrapper.text == 'hello world' + + +def test_response_wrapper_text() -> None: + """Test text property of ModelResponse.""" + wrapper = ModelResponse( + message=Message( + role='model', + content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))], + ), + ) + wrapper.request = ModelRequest(messages=[]) + + assert wrapper.text == 'hello world' + + +def test_response_wrapper_output() -> None: + """Test output property of ModelResponse.""" + wrapper = ModelResponse( + message=Message( + role='model', + content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"bar'))], + ), + ) + wrapper.request = ModelRequest(messages=[]) + + assert wrapper.output == {'foo': 'bar'} + + +def test_response_wrapper_messages() -> None: + """Test messages property of ModelResponse.""" + wrapper = ModelResponse( + message=Message( + role='model', + content=[Part(root=TextPart(text='baz'))], + ) + ) + wrapper.request = ModelRequest( + messages=[ + Message( + role='user', + content=[Part(root=TextPart(text='foo'))], + ), + Message( + role='tool', + content=[Part(root=TextPart(text='bar'))], + ), + ], + ) + + assert wrapper.messages == [ + Message( + role='user', + content=[Part(root=TextPart(text='foo'))], + ), + Message( + role='tool', + content=[Part(root=TextPart(text='bar'))], + ), + Message( + role='model', + content=[Part(root=TextPart(text='baz'))], + ), + ] + + +def test_response_wrapper_output_uses_parser() -> None: + """Test that ModelResponse uses the provided message_parser.""" + wrapper = ModelResponse( + message=Message( + role='model', + content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"bar'))], + ), + ) + wrapper.request = ModelRequest(messages=[]) + wrapper._message_parser = lambda x: 'banana' + + assert wrapper.output == 'banana' + + +def test_chunk_wrapper_text() -> None: + """Test text property of ModelResponseChunk.""" + wrapper = ModelResponseChunk( + chunk=ModelResponseChunk(content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))]), + index=0, + previous_chunks=[], + ) + + assert wrapper.text == 'hello world' + + +def test_chunk_wrapper_accumulated_text() -> None: + """Test accumulated_text property of ModelResponseChunk.""" + wrapper = ModelResponseChunk( + ModelResponseChunk(content=[Part(root=TextPart(text=' PS: aliens'))]), + index=0, + previous_chunks=[ + ModelResponseChunk(content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' '))]), + ModelResponseChunk(content=[Part(root=TextPart(text='world!'))]), + ], + ) + + assert wrapper.accumulated_text == 'hello world! PS: aliens' + + +def test_chunk_wrapper_output() -> None: + """Test output property of ModelResponseChunk.""" + wrapper = ModelResponseChunk( + ModelResponseChunk(content=[Part(root=TextPart(text=', "baz":[1,2,'))]), + index=0, + previous_chunks=[ + ModelResponseChunk(content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"ba'))]), + ModelResponseChunk(content=[Part(root=TextPart(text='r"'))]), + ], + ) + + assert wrapper.output == {'foo': 'bar', 'baz': [1, 2]} + + +def test_chunk_wrapper_output_uses_parser() -> None: + """Test that ModelResponseChunk uses the provided chunk_parser.""" + wrapper = ModelResponseChunk( + ModelResponseChunk(content=[Part(root=TextPart(text=', "baz":[1,2,'))]), + index=0, + previous_chunks=[ + ModelResponseChunk(content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"ba'))]), + ModelResponseChunk(content=[Part(root=TextPart(text='r"'))]), + ], + chunk_parser=lambda x: 'banana', + ) + + assert wrapper.output == 'banana' + + +@pytest.mark.parametrize( + 'test_input,test_response,expected_output', + ( + [ + [], + Message(role='model', content=[]), + ModelUsage( + input_images=0, + input_videos=0, + input_characters=0, + input_audio_files=0, + output_audio_files=0, + output_characters=0, + output_images=0, + output_videos=0, + ), + ], + [ + [ + Message( + role='user', + content=[ + Part(root=TextPart(text='1')), + Part(root=TextPart(text='2')), + ], + ), + Message( + role='user', + content=[ + Part(root=MediaPart(media=Media(content_type='image', url=''))), + Part(root=MediaPart(media=Media(url='data:image'))), + Part(root=MediaPart(media=Media(content_type='audio', url=''))), + Part(root=MediaPart(media=Media(url='data:audio'))), + Part(root=MediaPart(media=Media(content_type='video', url=''))), + Part(root=MediaPart(media=Media(url='data:video'))), + ], + ), + ], + Message( + role='model', + content=[ + Part(root=TextPart(text='3')), + Part(root=MediaPart(media=Media(content_type='image', url=''))), + Part(root=MediaPart(media=Media(url='data:image'))), + Part(root=MediaPart(media=Media(content_type='audio', url=''))), + Part(root=MediaPart(media=Media(url='data:audio'))), + Part(root=MediaPart(media=Media(content_type='video', url=''))), + Part(root=MediaPart(media=Media(url='data:video'))), + ], + ), + ModelUsage( + input_images=2, + input_videos=2, + input_characters=2, + input_audio_files=2, + output_audio_files=2, + output_characters=1, + output_images=2, + output_videos=2, + ), + ], + ), +) +def test_get_basic_usage_stats( + test_input: list[Message], + test_response: Message, + expected_output: ModelUsage, +) -> None: + """Test get_basic_usage_stats utility.""" + assert get_basic_usage_stats(input_=test_input, response=test_response) == expected_output + + +def test_response_wrapper_tool_requests() -> None: + """Test tool_requests property of ModelResponse.""" + wrapper = ModelResponse( + message=Message( + role='model', + content=[Part(root=TextPart(text='bar'))], + ) + ) + wrapper.request = ModelRequest( + messages=[ + Message( + role='user', + content=[Part(root=TextPart(text='foo'))], + ), + ], + ) + + assert wrapper.tool_requests == [] + + wrapper = ModelResponse( + message=Message( + role='model', + content=[ + Part(root=ToolRequestPart(tool_request=ToolRequest(name='tool', input={'abc': 3}))), + Part(root=TextPart(text='bar')), + ], + ) + ) + wrapper.request = ModelRequest( + messages=[ + Message( + role='user', + content=[Part(root=TextPart(text='foo'))], + ), + ], + ) + + assert wrapper.tool_requests == [ToolRequestPart(tool_request=ToolRequest(name='tool', input={'abc': 3}))] + + +def test_response_wrapper_interrupts() -> None: + """Test interrupts property of ModelResponse.""" + wrapper = ModelResponse( + message=Message( + role='model', + content=[Part(root=TextPart(text='bar'))], + ) + ) + wrapper.request = ModelRequest( + messages=[ + Message( + role='user', + content=[Part(root=TextPart(text='foo'))], + ), + ], + ) + + assert wrapper.interrupts == [] + + wrapper = ModelResponse( + message=Message( + role='model', + content=[ + Part(root=ToolRequestPart(tool_request=ToolRequest(name='tool1', input={'abc': 3}))), + Part( + root=ToolRequestPart( + tool_request=ToolRequest(name='tool2', input={'bcd': 4}), + metadata=Metadata(root={'interrupt': {'banana': 'yes'}}), + ) + ), + Part(root=TextPart(text='bar')), + ], + ) + ) + wrapper.request = ModelRequest( + messages=[ + Message( + role='user', + content=[Part(root=TextPart(text='foo'))], + ), + ], + ) + + assert wrapper.interrupts == [ + ToolRequestPart( + tool_request=ToolRequest(name='tool2', input={'bcd': 4}), + metadata=Metadata(root={'interrupt': {'banana': 'yes'}}), + ) + ] + + +def test_model_action_metadata() -> None: + """Test for model_action_metadata.""" + action_metadata = model_action_metadata( + name='test_model', + info={'label': 'test_label'}, + config_schema=None, + ) + + assert isinstance(action_metadata, ActionMetadata) + assert action_metadata.input_json_schema is not None + assert action_metadata.output_json_schema is not None + assert action_metadata.metadata == {'model': {'customOptions': None, 'label': 'test_label'}} + + +def test_text_from_content_with_parts() -> None: + """Test text_from_content with list of Part objects.""" + content = [Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))] + assert text_from_content(content) == 'hello world' + + +def test_text_from_content_with_document_parts() -> None: + """Test text_from_content with list of DocumentPart objects.""" + content = [DocumentPart(root=TextPart(text='doc1')), DocumentPart(root=TextPart(text=' doc2'))] + assert text_from_content(content) == 'doc1 doc2' + + +def test_text_from_content_with_mixed_parts() -> None: + """Test text_from_content with mixed Part and DocumentPart objects.""" + content = [ + Part(root=TextPart(text='part')), + DocumentPart(root=TextPart(text=' text')), + ] + assert text_from_content(content) == 'part text' + + +def test_text_from_content_with_empty_list() -> None: + """Test text_from_content with empty list.""" + assert text_from_content([]) == '' + + +def test_text_from_content_with_none_text() -> None: + """Test text_from_content handles parts without text content.""" + content = [ + Part(root=TextPart(text='hello')), + Part(root=MediaPart(media=Media(url='http://example.com/image.png'))), + Part(root=TextPart(text=' world')), + ] + assert text_from_content(content) == 'hello world' diff --git a/py/packages/genkit/tests/genkit/blocks/prompt_test.py b/py/packages/genkit/tests/genkit/ai/prompt_test.py similarity index 86% rename from py/packages/genkit/tests/genkit/blocks/prompt_test.py rename to py/packages/genkit/tests/genkit/ai/prompt_test.py index 01b5e54538..ec823f7de6 100644 --- a/py/packages/genkit/tests/genkit/blocks/prompt_test.py +++ b/py/packages/genkit/tests/genkit/ai/prompt_test.py @@ -25,28 +25,23 @@ import pytest from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.blocks.prompt import load_prompt_folder, lookup_prompt, prompt -from genkit.core.action.types import ActionKind -from genkit.core.typing import ( - DocumentData, - DocumentPart, +from genkit import Genkit, Message, ModelResponse +from genkit._ai._prompt import load_prompt_folder, lookup_prompt, prompt +from genkit._ai._testing import ( + EchoModel, + ProgrammableModel, + define_echo_model, + define_programmable_model, +) +from genkit._core._action import ActionKind +from genkit._core._model import ModelConfig, ModelRequest +from genkit._core._typing import ( GenerateActionOptions, - GenerateRequest, - GenerateResponse, - GenerationCommonConfig, - Message, Part, Role, TextPart, ToolChoice, ) -from genkit.testing import ( - EchoModel, - ProgrammableModel, - define_echo_model, - define_programmable_model, -) def setup_test() -> tuple[Genkit, EchoModel, ProgrammableModel]: @@ -64,7 +59,7 @@ async def test_simple_prompt() -> None: """Test simple prompt rendering.""" ai, *_ = setup_test() - want_txt = '[ECHO] user: "hi" {"temperature":11}' + want_txt = '[ECHO] user: "hi" {"temperature":11.0}' my_prompt = ai.define_prompt(prompt='hi', config={'temperature': 11}) @@ -72,7 +67,7 @@ async def test_simple_prompt() -> None: assert response.text == want_txt - # New API: stream returns GenerateStreamResponse with .response property + # New API: stream returns ModelStreamResponse with .response property result = my_prompt.stream() assert (await result.response).text == want_txt @@ -87,7 +82,7 @@ async def test_simple_prompt_with_override_config() -> None: ai, *_ = setup_test() # Config is MERGED: prompt config (banana: true) + opts config (temperature: 12) - want_txt = '[ECHO] user: "hi" {"banana":true,"temperature":12}' + want_txt = '[ECHO] user: "hi" {"temperature":12.0,"banana":true}' my_prompt = ai.define_prompt(prompt='hi', config={'banana': True}) @@ -115,7 +110,7 @@ async def test_prompt_with_system() -> None: assert response.text == want_txt - # New API: stream returns GenerateStreamResponse + # New API: stream returns ModelStreamResponse result = my_prompt.stream() assert (await result.response).text == want_txt @@ -133,7 +128,7 @@ class ToolInput(BaseModel): value: int | None = Field(default=None, description='value field') @ai.tool(name='testTool') - def test_tool(input: ToolInput) -> str: + async def test_tool(input: ToolInput) -> str: """The tool.""" return 'abc' @@ -160,63 +155,12 @@ def test_tool(input: ToolInput) -> str: assert response.text == want_txt - # New API: stream returns GenerateStreamResponse + # New API: stream returns ModelStreamResponse result = my_prompt.stream() assert (await result.response).text == want_txt -@pytest.mark.asyncio -async def test_prompt_with_resolvers() -> None: - """Test that the rendering works with resolvers.""" - ai, *_ = setup_test() - - async def system_resolver(input: dict[str, Any], context: object) -> str: - return f'system {input["name"]}' - - def prompt_resolver(input: dict[str, Any], context: object) -> str: - return f'prompt {input["name"]}' - - async def messages_resolver(input: dict[str, Any], context: object) -> list[Message]: - return [Message(role=Role.USER, content=[Part(root=TextPart(text=f'msg {input["name"]}'))])] - - my_prompt = ai.define_prompt( - system=system_resolver, - prompt=prompt_resolver, - messages=messages_resolver, - ) - - want_txt = '[ECHO] system: "system world" user: "msg world" user: "prompt world"' - - response = await my_prompt(input={'name': 'world'}) - - assert response.text == want_txt - - -@pytest.mark.asyncio -async def test_prompt_with_docs_resolver() -> None: - """Test that the rendering works with docs resolver.""" - ai, _, pm = setup_test() - - pm.responses = [GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='ok'))]))] - - async def docs_resolver(input: dict[str, Any], context: object) -> list[DocumentData]: - return [DocumentData(content=[DocumentPart(root=TextPart(text=f'doc {input["name"]}'))])] - - my_prompt = ai.define_prompt( - model='programmableModel', - prompt='hi', - docs=docs_resolver, - ) - - await my_prompt(input={'name': 'world'}) - - # Check that PM received the docs - assert pm.last_request is not None - assert pm.last_request.docs is not None - assert pm.last_request.docs[0].content[0].root.text == 'doc world' - - test_cases_parse_partial_json = [ ( 'renders system prompt', @@ -233,10 +177,10 @@ async def docs_resolver(input: dict[str, Any], context: object) -> list[Document 'metadata': {'state': {'name': 'bar'}}, }, {'name': 'foo'}, - GenerationCommonConfig.model_validate({'temperature': 11}), + ModelConfig.model_validate({'temperature': 11}), {}, # Config is MERGED: prompt config (banana: ripe) + opts config (temperature: 11) - """[ECHO] system: "hello foo (bar)" {"banana":"ripe","temperature":11.0}""", + """[ECHO] system: "hello foo (bar)" {"temperature":11.0,"banana":"ripe"}""", ), ( 'renders user prompt', @@ -253,10 +197,10 @@ async def docs_resolver(input: dict[str, Any], context: object) -> list[Document 'metadata': {'state': {'name': 'bar_system'}}, }, {'name': 'foo'}, - GenerationCommonConfig.model_validate({'temperature': 11}), + ModelConfig.model_validate({'temperature': 11}), {}, # Config is MERGED: prompt config (banana: ripe) + opts config (temperature: 11) - """[ECHO] user: "hello foo (bar_system)" {"banana":"ripe","temperature":11.0}""", + """[ECHO] user: "hello foo (bar_system)" {"temperature":11.0,"banana":"ripe"}""", ), ( 'renders user prompt with context', @@ -273,10 +217,10 @@ async def docs_resolver(input: dict[str, Any], context: object) -> list[Document 'metadata': {'state': {'name': 'bar'}}, }, {'name': 'foo'}, - GenerationCommonConfig.model_validate({'temperature': 11}), + ModelConfig.model_validate({'temperature': 11}), {'auth': {'email': 'a@b.c'}}, # Config is MERGED: prompt config (banana: ripe) + opts config (temperature: 11) - """[ECHO] user: "hello foo (bar, a@b.c)" {"banana":"ripe","temperature":11.0}""", + """[ECHO] user: "hello foo (bar, a@b.c)" {"temperature":11.0,"banana":"ripe"}""", ), ] @@ -291,7 +235,7 @@ async def test_prompt_rendering_dotprompt( test_case: str, prompt: dict[str, Any], input: dict[str, Any], - input_option: GenerationCommonConfig, + input_option: ModelConfig, context: dict[str, Any], want_rendered: str, ) -> None: @@ -467,7 +411,7 @@ class ToolInput(BaseModel): value: int = Field(description='A value') @ai.tool(name='myTool') - def my_tool(input: ToolInput) -> int: + async def my_tool(input: ToolInput) -> int: return input.value * 2 my_prompt = ai.define_prompt( @@ -559,9 +503,7 @@ async def test_opts_can_override_model() -> None: """Test that opts.model can override the prompt's default model.""" ai, _, pm = setup_test() - pm.responses = [ - GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='pm response'))])) - ] + pm.responses = [ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='pm response'))]))] my_prompt = ai.define_prompt( model='echoModel', @@ -602,7 +544,7 @@ async def test_opts_can_append_messages() -> None: @pytest.mark.asyncio async def test_generate_stream_response_api() -> None: - """Test that GenerateStreamResponse provides both stream and response.""" + """Test that ModelStreamResponse provides both stream and response.""" ai, *_ = setup_test() my_prompt = ai.define_prompt( @@ -612,7 +554,7 @@ async def test_generate_stream_response_api() -> None: # Get stream response result = my_prompt.stream() - # Verify it has the expected properties (matching JS GenerateStreamResponse) + # Verify it has the expected properties (matching JS ModelStreamResponse) assert hasattr(result, 'stream') assert hasattr(result, 'response') @@ -688,7 +630,7 @@ async def test_file_based_prompt_registers_two_actions() -> None: @pytest.mark.asyncio async def test_prompt_and_executable_prompt_return_types() -> None: - """PROMPT action returns GenerateRequest, EXECUTABLE_PROMPT returns GenerateActionOptions.""" + """PROMPT action returns ModelRequest, EXECUTABLE_PROMPT returns GenerateActionOptions.""" ai, *_ = setup_test() # Test with file-based prompt (which creates both actions) @@ -709,10 +651,10 @@ async def test_prompt_and_executable_prompt_return_types() -> None: assert prompt_action is not None assert executable_prompt_action is not None - prompt_result = await prompt_action.arun(input={'name': 'World'}) - assert isinstance(prompt_result.response, GenerateRequest) + prompt_result = await prompt_action.run(input={'name': 'World'}) + assert isinstance(prompt_result.response, ModelRequest) - exec_result = await executable_prompt_action.arun(input={'name': 'World'}) + exec_result = await executable_prompt_action.run(input={'name': 'World'}) assert isinstance(exec_result.response, GenerateActionOptions) @@ -797,7 +739,7 @@ async def test_automatic_prompt_loading_default_none() -> None: @pytest.mark.asyncio async def test_automatic_prompt_loading_defaults_mock() -> None: """Test that Genkit defaults to ./prompts when prompt_dir is not specified and dir exists.""" - with patch('genkit.ai._aio.load_prompt_folder') as mock_load, patch('genkit.ai._aio.Path') as mock_path: + with patch('genkit._ai._aio.load_prompt_folder') as mock_load, patch('genkit._ai._aio.Path') as mock_path: # Setup mock to simulate ./prompts existing mock_path_instance = MagicMock() mock_path_instance.is_dir.return_value = True @@ -810,7 +752,7 @@ async def test_automatic_prompt_loading_defaults_mock() -> None: @pytest.mark.asyncio async def test_automatic_prompt_loading_defaults_missing() -> None: """Test that Genkit skips loading when ./prompts is missing.""" - with patch('genkit.ai._aio.load_prompt_folder') as mock_load, patch('genkit.ai._aio.Path') as mock_path: + with patch('genkit._ai._aio.load_prompt_folder') as mock_load, patch('genkit._ai._aio.Path') as mock_path: # Setup mock to simulate ./prompts missing mock_path_instance = MagicMock() mock_path_instance.is_dir.return_value = False diff --git a/py/packages/genkit/tests/genkit/ai/resource_integration_test.py b/py/packages/genkit/tests/genkit/ai/resource_integration_test.py index 922ba0ca73..ff8cc37aeb 100644 --- a/py/packages/genkit/tests/genkit/ai/resource_integration_test.py +++ b/py/packages/genkit/tests/genkit/ai/resource_integration_test.py @@ -21,15 +21,14 @@ import pytest -from genkit.blocks.generate import generate_action -from genkit.blocks.resource import ResourceInput, ResourceOutput, define_resource, resource -from genkit.core.action import ActionRunContext -from genkit.core.registry import ActionKind, Registry -from genkit.core.typing import ( +from genkit import Message, ModelResponse +from genkit._ai._generate import generate_action +from genkit._ai._resource import ResourceInput, ResourceOutput, define_resource, resource +from genkit._core._action import ActionRunContext +from genkit._core._model import ModelRequest +from genkit._core._registry import ActionKind, Registry +from genkit._core._typing import ( GenerateActionOptions, - GenerateRequest, - GenerateResponse, - Message, Part, Resource1, ResourcePart, @@ -50,13 +49,13 @@ async def my_resource(input: ResourceInput, ctx: ActionRunContext) -> ResourceOu define_resource(registry, {'uri': 'test://foo'}, my_resource) # 2. Register a mock model - async def mock_model(input: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def mock_model(input: ModelRequest, ctx: ActionRunContext) -> ModelResponse: # Verify docs are EMPTY (not auto-populated) assert not input.docs # Access via root because DocumentPart is a RootModel # Verify the message content was hydrated (replaced resource part with text part) assert input.messages[0].content[0].root.text == 'Resource content for test://foo' - return GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='Done'))])) + return ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='Done'))])) registry.register_action(ActionKind.MODEL, 'mock-model', mock_model) @@ -78,7 +77,7 @@ async def test_dynamic_action_provider_resource() -> None: registry = Registry() # Register a dynamic provider that handles any "dynamic://*" uri - def provider_fn(input: dict[str, object], ctx: ActionRunContext) -> object: + async def provider_fn(input: dict[str, object], ctx: ActionRunContext) -> object: kind = cast(ActionKind, input['kind']) name = cast(str, input['name']) if kind == ActionKind.RESOURCE and name.startswith('dynamic://'): @@ -96,12 +95,12 @@ async def dyn_res_fn(input: ResourceInput, ctx: ActionRunContext) -> ResourceOut # Register mock model # Register mock model - async def mock_model(input: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def mock_model(input: ModelRequest, ctx: ActionRunContext) -> ModelResponse: # Verify docs are empty assert not input.docs # Verify dynamic hydration assert input.messages[0].content[0].root.text == 'Dynamic content for dynamic://bar' - return GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='Done'))])) + return ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='Done'))])) registry.register_action(ActionKind.MODEL, 'mock-model', mock_model) diff --git a/py/packages/genkit/tests/genkit/ai/resource_test.py b/py/packages/genkit/tests/genkit/ai/resource_test.py index 68e9cb8fc7..457f301eb4 100644 --- a/py/packages/genkit/tests/genkit/ai/resource_test.py +++ b/py/packages/genkit/tests/genkit/ai/resource_test.py @@ -25,7 +25,7 @@ import pytest -from genkit.blocks.resource import ( +from genkit._ai._resource import ( ResourceInput, define_resource, find_matching_resource, @@ -33,10 +33,9 @@ resolve_resources, resource, ) -from genkit.core.action import ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry -from genkit.core.typing import Metadata, Part, TextPart +from genkit._core._action import ActionKind, ActionRunContext +from genkit._core._registry import Registry +from genkit._core._typing import Metadata, Part, TextPart @pytest.mark.asyncio @@ -181,7 +180,7 @@ async def fn(input: ResourceInput, ctx: ActionRunContext) -> dict[str, object]: res = define_resource(registry, {'template': 'file://{id}'}, fn) - output = await res.arun({'uri': 'file://dir'}) + output = await res.run({'uri': 'file://dir'}) # output is ActionResponse # content is in output.response['content'] because wrapped_fn ensures serialization diff --git a/py/packages/genkit/tests/genkit/aio/util_test.py b/py/packages/genkit/tests/genkit/aio/util_test.py deleted file mode 100644 index af854140c3..0000000000 --- a/py/packages/genkit/tests/genkit/aio/util_test.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for aio utility functions.""" - -import pytest - -from genkit.aio._util import ensure_async - - -@pytest.mark.asyncio -async def test_ensure_async_with_sync_function() -> None: - """Test that sync functions are wrapped correctly.""" - - def sync_fn(x: int) -> int: - return x * 2 - - async_fn = ensure_async(sync_fn) - result = await async_fn(5) - assert result == 10 - - -@pytest.mark.asyncio -async def test_ensure_async_with_async_function() -> None: - """Test that async functions are returned as-is.""" - - async def async_fn(x: int) -> int: - return x * 2 - - wrapped = ensure_async(async_fn) - assert wrapped is async_fn # Should be the same function - result = await wrapped(5) - assert result == 10 - - -@pytest.mark.asyncio -async def test_ensure_async_with_lambda_returning_coroutine() -> None: - """Test that lambdas returning coroutines are handled correctly. - - This is the key fix: when a sync function (lambda) returns a coroutine, - ensure_async should await it. - """ - - async def async_operation(x: int) -> int: - return x * 2 - - # Lambda that returns a coroutine (common pattern with ai.run()) - lambda_fn = lambda: async_operation(5) # noqa: E731 - - async_fn = ensure_async(lambda_fn) - result = await async_fn() - assert result == 10 - - -@pytest.mark.asyncio -async def test_ensure_async_with_lambda_returning_value() -> None: - """Test that lambdas returning regular values work correctly.""" - lambda_fn = lambda x: x * 2 # noqa: E731 - - async_fn = ensure_async(lambda_fn) - result = await async_fn(5) - assert result == 10 - - -@pytest.mark.asyncio -async def test_ensure_async_with_nested_coroutine_pattern() -> None: - """Test the nested pattern used in ai.run() with recursive async functions.""" - - async def recursive_fn(depth: int) -> str: - if depth <= 0: - return 'done' - # Simulate what ai.run() does with a lambda - wrapped = ensure_async(lambda: recursive_fn(depth - 1)) - return f'level-{depth}:' + await wrapped() - - result = await recursive_fn(3) - assert result == 'level-3:level-2:level-1:done' diff --git a/py/packages/genkit/tests/genkit/blocks/model_test.py b/py/packages/genkit/tests/genkit/blocks/model_test.py deleted file mode 100644 index 7240d442ac..0000000000 --- a/py/packages/genkit/tests/genkit/blocks/model_test.py +++ /dev/null @@ -1,536 +0,0 @@ -#!/usr/bin/env python3 -# -# Copyright 2025 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for the action module.""" - -import pytest - -from genkit.blocks.model import ( - GenerateResponseChunkWrapper, - GenerateResponseWrapper, - MessageWrapper, - PartCounts, - get_basic_usage_stats, - get_part_counts, - model_action_metadata, - text_from_content, -) -from genkit.core.action import ActionMetadata -from genkit.core.typing import ( - Candidate, - DocumentPart, - FinishReason, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationUsage, - Media, - MediaPart, - Message, - Metadata, - Part, - TextPart, - ToolRequest, - ToolRequestPart, -) - - -def test_message_wrapper_text() -> None: - """Test text property of MessageWrapper.""" - wrapper = MessageWrapper( - Message( - role='model', - content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))], - ), - ) - - assert wrapper.text == 'hello world' - - -def test_response_wrapper_text() -> None: - """Test text property of GenerateResponseWrapper.""" - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))], - ) - ), - request=GenerateRequest( - messages=[], # doesn't matter for now - ), - ) - - assert wrapper.text == 'hello world' - - -def test_response_wrapper_output() -> None: - """Test output property of GenerateResponseWrapper.""" - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"bar'))], - ) - ), - request=GenerateRequest( - messages=[], # doesn't matter for now - ), - ) - - assert wrapper.output == {'foo': 'bar'} - - -def test_response_wrapper_messages() -> None: - """Test messages property of GenerateResponseWrapper.""" - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[Part(root=TextPart(text='baz'))], - ) - ), - request=GenerateRequest( - messages=[ - Message( - role='user', - content=[Part(root=TextPart(text='foo'))], - ), - Message( - role='tool', - content=[Part(root=TextPart(text='bar'))], - ), - ], - ), - ) - - assert wrapper.messages == [ - Message( - role='user', - content=[Part(root=TextPart(text='foo'))], - ), - Message( - role='tool', - content=[Part(root=TextPart(text='bar'))], - ), - Message( - role='model', - content=[Part(root=TextPart(text='baz'))], - ), - ] - - -def test_response_wrapper_output_uses_parser() -> None: - """Test that GenerateResponseWrapper uses the provided message_parser.""" - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"bar'))], - ) - ), - request=GenerateRequest( - messages=[], # doesn't matter for now - ), - message_parser=lambda x: 'banana', - ) - - assert wrapper.output == 'banana' - - -def test_chunk_wrapper_text() -> None: - """Test text property of GenerateResponseChunkWrapper.""" - wrapper = GenerateResponseChunkWrapper( - chunk=GenerateResponseChunk(content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))]), - index=0, - previous_chunks=[], - ) - - assert wrapper.text == 'hello world' - - -def test_chunk_wrapper_accumulated_text() -> None: - """Test accumulated_text property of GenerateResponseChunkWrapper.""" - wrapper = GenerateResponseChunkWrapper( - GenerateResponseChunk(content=[Part(root=TextPart(text=' PS: aliens'))]), - index=0, - previous_chunks=[ - GenerateResponseChunk(content=[Part(root=TextPart(text='hello')), Part(root=TextPart(text=' '))]), - GenerateResponseChunk(content=[Part(root=TextPart(text='world!'))]), - ], - ) - - assert wrapper.accumulated_text == 'hello world! PS: aliens' - - -def test_chunk_wrapper_output() -> None: - """Test output property of GenerateResponseChunkWrapper.""" - wrapper = GenerateResponseChunkWrapper( - GenerateResponseChunk(content=[Part(root=TextPart(text=', "baz":[1,2,'))]), - index=0, - previous_chunks=[ - GenerateResponseChunk(content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"ba'))]), - GenerateResponseChunk(content=[Part(root=TextPart(text='r"'))]), - ], - ) - - assert wrapper.output == {'foo': 'bar', 'baz': [1, 2]} - - -def test_chunk_wrapper_output_uses_parser() -> None: - """Test that GenerateResponseChunkWrapper uses the provided chunk_parser.""" - wrapper = GenerateResponseChunkWrapper( - GenerateResponseChunk(content=[Part(root=TextPart(text=', "baz":[1,2,'))]), - index=0, - previous_chunks=[ - GenerateResponseChunk(content=[Part(root=TextPart(text='{"foo":')), Part(root=TextPart(text='"ba'))]), - GenerateResponseChunk(content=[Part(root=TextPart(text='r"'))]), - ], - chunk_parser=lambda x: 'banana', - ) - - assert wrapper.output == 'banana' - - -@pytest.mark.parametrize( - 'test_parts,expected_part_counts', - ( - [[], PartCounts()], - [ - [ - Part(root=MediaPart(media=Media(content_type='image', url=''))), - Part(root=MediaPart(media=Media(url='data:image'))), - Part(root=MediaPart(media=Media(content_type='audio', url=''))), - Part(root=MediaPart(media=Media(url='data:audio'))), - Part(root=MediaPart(media=Media(content_type='video', url=''))), - Part(root=MediaPart(media=Media(url='data:video'))), - Part(root=TextPart(text='test')), - ], - PartCounts( - characters=len('test'), - audio=2, - videos=2, - images=2, - ), - ], - ), -) -def test_get_part_counts(test_parts: list[Part], expected_part_counts: PartCounts) -> None: - """Test get_part_counts utility.""" - assert get_part_counts(parts=test_parts) == expected_part_counts - - -@pytest.mark.parametrize( - 'test_input,test_response,expected_output', - ( - [ - [], - [], - GenerationUsage( - input_images=0, - input_videos=0, - input_characters=0, - input_audio_files=0, - output_audio_files=0, - output_characters=0, - output_images=0, - output_videos=0, - ), - ], - [ - [ - Message( - role='user', - content=[ - Part(root=TextPart(text='1')), - Part(root=TextPart(text='2')), - ], - ), - Message( - role='user', - content=[ - Part(root=MediaPart(media=Media(content_type='image', url=''))), - Part(root=MediaPart(media=Media(url='data:image'))), - Part(root=MediaPart(media=Media(content_type='audio', url=''))), - Part(root=MediaPart(media=Media(url='data:audio'))), - Part(root=MediaPart(media=Media(content_type='video', url=''))), - Part(root=MediaPart(media=Media(url='data:video'))), - ], - ), - ], - Message( - role='user', - content=[ - Part(root=TextPart(text='3')), - Part(root=MediaPart(media=Media(content_type='image', url=''))), - Part(root=MediaPart(media=Media(url='data:image'))), - Part(root=MediaPart(media=Media(content_type='audio', url=''))), - Part(root=MediaPart(media=Media(url='data:audio'))), - Part(root=MediaPart(media=Media(content_type='video', url=''))), - Part(root=MediaPart(media=Media(url='data:video'))), - ], - ), - GenerationUsage( - input_images=2, - input_videos=2, - input_characters=2, - input_audio_files=2, - output_audio_files=2, - output_characters=1, - output_images=2, - output_videos=2, - ), - ], - [ - [ - Message( - role='user', - content=[ - Part(root=TextPart(text='1')), - Part(root=TextPart(text='2')), - ], - ), - ], - [ - Candidate( - index=0, - finish_reason=FinishReason.STOP, - message=Message( - role='user', - content=[ - Part(root=TextPart(text='3')), - ], - ), - ), - Candidate( - index=1, - finish_reason=FinishReason.STOP, - message=Message( - role='user', - content=[ - Part(root=MediaPart(media=Media(content_type='image', url=''))), - ], - ), - ), - Candidate( - index=2, - finish_reason=FinishReason.STOP, - message=Message( - role='user', - content=[ - Part(root=MediaPart(media=Media(url='data:image'))), - ], - ), - ), - Candidate( - index=3, - finish_reason=FinishReason.STOP, - message=Message( - role='user', - content=[ - Part(root=MediaPart(media=Media(content_type='audio', url=''))), - ], - ), - ), - Candidate( - index=4, - finish_reason=FinishReason.STOP, - message=Message( - role='user', - content=[ - Part(root=MediaPart(media=Media(url='data:audio'))), - ], - ), - ), - Candidate( - index=5, - finish_reason=FinishReason.STOP, - message=Message( - role='user', - content=[ - Part(root=MediaPart(media=Media(content_type='video', url=''))), - ], - ), - ), - Candidate( - index=6, - finish_reason=FinishReason.STOP, - message=Message( - role='user', - content=[ - Part(root=MediaPart(media=Media(url='data:video'))), - ], - ), - ), - ], - GenerationUsage( - input_images=0, - input_videos=0, - input_characters=2, - input_audio_files=0, - output_audio_files=2, - output_characters=1, - output_images=2, - output_videos=2, - ), - ], - ), -) -def test_get_basic_usage_stats( - test_input: list[Message], - test_response: Message | list[Candidate], - expected_output: GenerationUsage, -) -> None: - """Test get_basic_usage_stats utility.""" - assert get_basic_usage_stats(input_=test_input, response=test_response) == expected_output - - -def test_response_wrapper_tool_requests() -> None: - """Test tool_requests property of GenerateResponseWrapper.""" - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[Part(root=TextPart(text='bar'))], - ) - ), - request=GenerateRequest( - messages=[ - Message( - role='user', - content=[Part(root=TextPart(text='foo'))], - ), - ], - ), - ) - - assert wrapper.tool_requests == [] - - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[ - Part(root=ToolRequestPart(tool_request=ToolRequest(name='tool', input={'abc': 3}))), - Part(root=TextPart(text='bar')), - ], - ) - ), - request=GenerateRequest( - messages=[ - Message( - role='user', - content=[Part(root=TextPart(text='foo'))], - ), - ], - ), - ) - - assert wrapper.tool_requests == [ToolRequestPart(tool_request=ToolRequest(name='tool', input={'abc': 3}))] - - -def test_response_wrapper_interrupts() -> None: - """Test interrupts property of GenerateResponseWrapper.""" - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[Part(root=TextPart(text='bar'))], - ) - ), - request=GenerateRequest( - messages=[ - Message( - role='user', - content=[Part(root=TextPart(text='foo'))], - ), - ], - ), - ) - - assert wrapper.interrupts == [] - - wrapper = GenerateResponseWrapper( - response=GenerateResponse( - message=Message( - role='model', - content=[ - Part(root=ToolRequestPart(tool_request=ToolRequest(name='tool1', input={'abc': 3}))), - Part( - root=ToolRequestPart( - tool_request=ToolRequest(name='tool2', input={'bcd': 4}), - metadata=Metadata(root={'interrupt': {'banana': 'yes'}}), - ) - ), - Part(root=TextPart(text='bar')), - ], - ) - ), - request=GenerateRequest( - messages=[ - Message( - role='user', - content=[Part(root=TextPart(text='foo'))], - ), - ], - ), - ) - - assert wrapper.interrupts == [ - ToolRequestPart( - tool_request=ToolRequest(name='tool2', input={'bcd': 4}), - metadata=Metadata(root={'interrupt': {'banana': 'yes'}}), - ) - ] - - -def test_model_action_metadata() -> None: - """Test for model_action_metadata.""" - action_metadata = model_action_metadata( - name='test_model', - info={'label': 'test_label'}, - config_schema=None, - ) - - assert isinstance(action_metadata, ActionMetadata) - assert action_metadata.input_json_schema is not None - assert action_metadata.output_json_schema is not None - assert action_metadata.metadata == {'model': {'customOptions': None, 'label': 'test_label'}} - - -def test_text_from_content_with_parts() -> None: - """Test text_from_content with list of Part objects.""" - content = [Part(root=TextPart(text='hello')), Part(root=TextPart(text=' world'))] - assert text_from_content(content) == 'hello world' - - -def test_text_from_content_with_document_parts() -> None: - """Test text_from_content with list of DocumentPart objects.""" - content = [DocumentPart(root=TextPart(text='doc1')), DocumentPart(root=TextPart(text=' doc2'))] - assert text_from_content(content) == 'doc1 doc2' - - -def test_text_from_content_with_mixed_parts() -> None: - """Test text_from_content with mixed Part and DocumentPart objects.""" - content = [ - Part(root=TextPart(text='part')), - DocumentPart(root=TextPart(text=' text')), - ] - assert text_from_content(content) == 'part text' - - -def test_text_from_content_with_empty_list() -> None: - """Test text_from_content with empty list.""" - assert text_from_content([]) == '' - - -def test_text_from_content_with_none_text() -> None: - """Test text_from_content handles parts without text content.""" - content = [ - Part(root=TextPart(text='hello')), - Part(root=MediaPart(media=Media(url='http://example.com/image.png'))), - Part(root=TextPart(text=' world')), - ] - assert text_from_content(content) == 'hello world' diff --git a/py/packages/genkit/tests/genkit/blocks/reranker_test.py b/py/packages/genkit/tests/genkit/blocks/reranker_test.py deleted file mode 100644 index 07f122b002..0000000000 --- a/py/packages/genkit/tests/genkit/blocks/reranker_test.py +++ /dev/null @@ -1,547 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for the reranker module. - -This module contains tests for the reranker functionality including -RankedDocument, define_reranker, and rerank functions. -""" - -from typing import Any, cast - -import pytest - -from genkit.blocks.reranker import ( - RankedDocument, - RerankerInfo, - RerankerOptions, - create_reranker_ref, - define_reranker, - rerank, - reranker_action_metadata, -) -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry -from genkit.core.typing import ( - DocumentData, - DocumentPart, - RankedDocumentData, - RankedDocumentMetadata, - RerankerResponse, - TextPart, -) -from genkit.types import Document - - -class TestRankedDocument: - """Tests for the RankedDocument class.""" - - def test_ranked_document_creation(self) -> None: - """Test creating a RankedDocument with content and score.""" - content = [DocumentPart(root=TextPart(text='Test content'))] - metadata = {'key': 'value'} - score = 0.95 - - doc = RankedDocument(content=content, metadata=metadata, score=score) - - assert doc.score == 0.95 - assert doc.text() == 'Test content' - assert doc.metadata == {'key': 'value', 'score': 0.95} - # Original metadata should not be modified - assert metadata == {'key': 'value'} - - def test_ranked_document_default_score(self) -> None: - """Test that RankedDocument has a default score of None.""" - content = [DocumentPart(root=TextPart(text='Test'))] - doc = RankedDocument(content=content) - - assert doc.score is None - - def test_ranked_document_from_data(self) -> None: - """Test creating RankedDocument from RankedDocumentData.""" - data = RankedDocumentData( - content=[DocumentPart(root=TextPart(text='Test content'))], - metadata=RankedDocumentMetadata(score=0.85), - ) - - doc = RankedDocument.from_ranked_document_data(data) - - assert doc.score == 0.85 - assert doc.text() == 'Test content' - - -class TestRerankerRef: - """Tests for RerankerRef and related helper functions.""" - - def test_create_reranker_ref_basic(self) -> None: - """Test creating a basic reranker reference.""" - ref = create_reranker_ref('test-reranker') - - assert ref.name == 'test-reranker' - assert ref.config is None - assert ref.version is None - assert ref.info is None - - def test_create_reranker_ref_with_options(self) -> None: - """Test creating a reranker reference with all options.""" - info = RerankerInfo(label='Test Reranker') - ref = create_reranker_ref( - name='test-reranker', - config={'k': 10}, - version='1.0.0', - info=info, - ) - - assert ref.name == 'test-reranker' - assert ref.config == {'k': 10} - assert ref.version == '1.0.0' - assert ref.info is not None - assert ref.info.label == 'Test Reranker' - - -class TestRerankerActionMetadata: - """Tests for reranker action metadata creation.""" - - def test_action_metadata_basic(self) -> None: - """Test creating basic action metadata.""" - metadata = reranker_action_metadata('test-reranker') - - assert metadata.kind == ActionKind.RERANKER - assert metadata.name == 'test-reranker' - assert metadata.metadata is not None - assert 'reranker' in metadata.metadata - - def test_action_metadata_with_options(self) -> None: - """Test creating action metadata with options.""" - options = RerankerOptions( - label='Custom Label', - config_schema={'type': 'object'}, - ) - metadata = reranker_action_metadata('test-reranker', options) - - assert metadata.metadata is not None - md = cast(dict[str, Any], metadata.metadata) - assert md['reranker']['label'] == 'Custom Label' - assert md['reranker']['customOptions'] == {'type': 'object'} - - -class TestDefineReranker: - """Tests for the define_reranker function.""" - - @pytest.fixture - def registry(self) -> Registry: - """Create a fresh registry for each test.""" - return Registry() - - @pytest.mark.asyncio - async def test_define_reranker_registers_action(self, registry: Registry) -> None: - """Test that define_reranker registers an action in the registry.""" - - async def simple_reranker( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - # Return documents in same order with scores - return RerankerResponse( - documents=[ - RankedDocumentData( - content=doc.content, - metadata=RankedDocumentMetadata(score=1.0 - i * 0.1), - ) - for i, doc in enumerate(documents) - ] - ) - - action = define_reranker(registry, 'test-reranker', simple_reranker) - - # Verify action was registered - lookup = await registry.resolve_action(ActionKind.RERANKER, 'test-reranker') - assert lookup is not None - assert action.name == 'test-reranker' - - @pytest.mark.asyncio - async def test_define_reranker_with_options(self, registry: Registry) -> None: - """Test define_reranker with custom options.""" - - async def reranker_fn( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - return RerankerResponse(documents=[]) - - options = RerankerOptions(label='My Reranker') - action = define_reranker(registry, 'my-reranker', reranker_fn, options) - - assert action is not None - - -class TestRerank: - """Tests for the rerank function.""" - - @pytest.fixture - def registry(self) -> Registry: - """Create a fresh registry for each test.""" - return Registry() - - @pytest.fixture - def sample_documents(self) -> list[DocumentData]: - """Create sample documents for testing.""" - return [ - DocumentData(content=[DocumentPart(root=TextPart(text='First document'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='Second document'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='Third document'))]), - ] - - @pytest.mark.asyncio - async def test_rerank_with_string_query( - self, - registry: Registry, - sample_documents: list[Document], - ) -> None: - """Test rerank with a string query.""" - - async def score_by_length( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - # Score documents by content length (longer = higher score) - scored = [] - for doc in documents: - length = len(doc.text()) - scored.append( - RankedDocumentData( - content=doc.content, - metadata=RankedDocumentMetadata(score=float(length)), - ) - ) - return RerankerResponse(documents=scored) - - define_reranker(registry, 'length-reranker', score_by_length) - - results = await rerank( - registry, - { - 'reranker': 'length-reranker', - 'query': 'test query', - 'documents': sample_documents, - }, - ) - - assert len(results) == 3 - assert all(isinstance(r, RankedDocument) for r in results) - - @pytest.mark.asyncio - async def test_rerank_with_reranker_ref( - self, - registry: Registry, - sample_documents: list[Document], - ) -> None: - """Test rerank with a RerankerRef.""" - - async def simple_reranker( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - return RerankerResponse( - documents=[ - RankedDocumentData( - content=doc.content, - metadata=RankedDocumentMetadata(score=0.5), - ) - for doc in documents - ] - ) - - define_reranker(registry, 'ref-reranker', simple_reranker) - ref = create_reranker_ref('ref-reranker') - - results = await rerank( - registry, - { - 'reranker': ref, - 'query': 'test', - 'documents': sample_documents, - }, - ) - - assert len(results) == 3 - assert all(doc.score == 0.5 for doc in results) - - @pytest.mark.asyncio - async def test_rerank_unknown_reranker_raises( - self, - registry: Registry, - sample_documents: list[Document], - ) -> None: - """Test that rerank raises ValueError for unknown reranker.""" - with pytest.raises(ValueError, match='Unable to resolve reranker'): - await rerank( - registry, - { - 'reranker': 'non-existent-reranker', - 'query': 'test', - 'documents': sample_documents, - }, - ) - - -class TestCustomRerankers: - """Tests for custom reranker implementations. - - These tests demonstrate how to create custom rerankers as shown - in the genkit.dev documentation: - https://genkit.dev/docs/rag/#rerankers-and-two-stage-retrieval - """ - - @pytest.fixture - def registry(self) -> Registry: - """Create a fresh registry for each test.""" - return Registry() - - @pytest.fixture - def sample_documents(self) -> list[DocumentData]: - """Create sample documents matching genkit.dev documentation example.""" - return [ - DocumentData(content=[DocumentPart(root=TextPart(text='pythagorean theorem'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='e=mc^2'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='pi'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='dinosaurs'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='quantum mechanics'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='pizza'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='harry potter'))]), - ] - - @pytest.mark.asyncio - async def test_custom_keyword_overlap_reranker( - self, - registry: Registry, - sample_documents: list[Document], - ) -> None: - """Test a custom reranker that scores by keyword overlap. - - This demonstrates the pattern shown in genkit.dev docs for - creating custom reranking logic. - """ - - async def keyword_overlap_reranker( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - """Reranker that scores documents by keyword overlap with query.""" - query_words = set(query.text().lower().split()) - scored = [] - - for doc in documents: - doc_words = set(doc.text().lower().split()) - overlap = len(query_words & doc_words) - score = overlap / max(len(query_words), 1) - scored.append((doc, score)) - - # Sort by score descending - scored.sort(key=lambda x: x[1], reverse=True) - - # Apply k limit if provided in options - k = options.get('k', len(scored)) if options else len(scored) - top_k = scored[:k] - - return RerankerResponse( - documents=[ - RankedDocumentData( - content=doc.content, - metadata=RankedDocumentMetadata(score=score), - ) - for doc, score in top_k - ] - ) - - define_reranker(registry, 'custom/keyword-overlap', keyword_overlap_reranker) - - # Query for 'quantum' should rank 'quantum mechanics' highest - results = await rerank( - registry, - { - 'reranker': 'custom/keyword-overlap', - 'query': 'quantum mechanics physics', - 'documents': sample_documents, - }, - ) - - assert len(results) == 7 - # 'quantum mechanics' should have the highest score (overlaps 2 words) - assert results[0].text() == 'quantum mechanics' - assert results[0].score is not None - assert results[0].score > 0 - - @pytest.mark.asyncio - async def test_custom_reranker_with_top_k_option( - self, - registry: Registry, - sample_documents: list[Document], - ) -> None: - """Test custom reranker with k option to limit results. - - Demonstrates using options to configure reranking behavior. - """ - - async def random_score_reranker( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - """Reranker that assigns incrementing scores and respects k option.""" - k = options.get('k', 3) if options else 3 - - scored_docs = [] - for i, doc in enumerate(documents): - # Score in reverse order so we have a predictable ranking - score = float(len(documents) - i) - scored_docs.append( - RankedDocumentData( - content=doc.content, - metadata=RankedDocumentMetadata(score=score), - ) - ) - - # Sort by score descending and limit to k - scored_docs.sort(key=lambda d: d.metadata.score, reverse=True) - return RerankerResponse(documents=scored_docs[:k]) - - define_reranker(registry, 'custom/with-k-option', random_score_reranker) - - results = await rerank( - registry, - { - 'reranker': 'custom/with-k-option', - 'query': 'test', - 'documents': sample_documents, - 'options': {'k': 3}, - }, - ) - - # Should only return top 3 results - assert len(results) == 3 - - @pytest.mark.asyncio - async def test_custom_reranker_preserves_document_content(self, registry: Registry) -> None: - """Test that custom reranker preserves original document content.""" - - async def identity_reranker( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - """Reranker that returns documents with their original content.""" - return RerankerResponse( - documents=[ - RankedDocumentData( - content=doc.content, - metadata=RankedDocumentMetadata(score=1.0), - ) - for doc in documents - ] - ) - - define_reranker(registry, 'custom/identity', identity_reranker) - - original_texts = ['Document A', 'Document B with more text', 'Doc C'] - documents = [DocumentData(content=[DocumentPart(root=TextPart(text=t))]) for t in original_texts] - - results = await rerank( - registry, - { - 'reranker': 'custom/identity', - 'query': 'test', - 'documents': documents, - }, - ) - - # Verify all original content is preserved - result_texts = [doc.text() for doc in results] - assert result_texts == original_texts - - @pytest.mark.asyncio - async def test_custom_reranker_two_stage_retrieval_pattern(self, registry: Registry) -> None: - """Test the two-stage retrieval pattern: retrieve then rerank. - - This demonstrates the typical RAG pattern where: - 1. Stage 1: Retrieve a broad set of candidates - 2. Stage 2: Rerank to find most relevant documents - """ - # Simulate stage 1 retrieval results (unranked) - retrieved_documents = [ - DocumentData(content=[DocumentPart(root=TextPart(text='Machine learning is a subset of AI'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='Pizza is a popular food'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='Deep learning uses neural networks'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='Cats are domestic animals'))]), - DocumentData(content=[DocumentPart(root=TextPart(text='AI transforms industries'))]), - ] - - async def relevance_reranker( - query: Document, - documents: list[Document], - options: dict[str, Any] | None, - ) -> RerankerResponse: - """Reranker that scores by word presence in query.""" - query_lower = query.text().lower() - scored = [] - - for doc in documents: - doc_text = doc.text().lower() - # Simple relevance: count query words in document - score = sum(1 for word in query_lower.split() if word in doc_text) - scored.append((doc, float(score))) - - scored.sort(key=lambda x: x[1], reverse=True) - - return RerankerResponse( - documents=[ - RankedDocumentData( - content=doc.content, - metadata=RankedDocumentMetadata(score=score), - ) - for doc, score in scored - ] - ) - - define_reranker(registry, 'custom/relevance', relevance_reranker) - - # Stage 2: Rerank with query about AI - reranked = await rerank( - registry, - { - 'reranker': 'custom/relevance', - 'query': 'artificial intelligence AI', - 'documents': retrieved_documents, - }, - ) - - # AI-related documents should rank higher than unrelated ones - # Get scores for AI and non-AI documents - ai_scores = [doc.score for doc in reranked if 'AI' in doc.text() or 'learning' in doc.text()] - non_ai_scores = [doc.score for doc in reranked if 'Pizza' in doc.text() or 'Cats' in doc.text()] - - # Filter out None values and validate lists are not empty before comparing - ai_scores_valid = [s for s in ai_scores if s is not None] - non_ai_scores_valid = [s for s in non_ai_scores if s is not None] - assert ai_scores_valid and non_ai_scores_valid, ( - 'Cannot compare scores if one of the lists is empty or contains only None' - ) - assert max(ai_scores_valid) > max(non_ai_scores_valid) diff --git a/py/packages/genkit/tests/genkit/blocks/retriever_test.py b/py/packages/genkit/tests/genkit/blocks/retriever_test.py deleted file mode 100644 index 40af79f73c..0000000000 --- a/py/packages/genkit/tests/genkit/blocks/retriever_test.py +++ /dev/null @@ -1,188 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Tests for Genkit retrievers and indexers.""" - -from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock - -import pytest -from pydantic import BaseModel - -from genkit.blocks.retriever import ( - IndexerOptions, - RetrieverOptions, - RetrieverResponse, - RetrieverSupports, - create_indexer_ref, - create_retriever_ref, - define_indexer, - define_retriever, - indexer_action_metadata, - retriever_action_metadata, -) -from genkit.core.action import ActionMetadata -from genkit.core.schema import to_json_schema - - -def test_retriever_action_metadata() -> None: - """Test for retriever_action_metadata with basic options.""" - options = RetrieverOptions(label='Test Retriever') - action_metadata = retriever_action_metadata( - name='test_retriever', - options=options, - ) - - assert isinstance(action_metadata, ActionMetadata) - assert action_metadata.input_json_schema is not None - assert action_metadata.output_json_schema is not None - assert action_metadata.metadata == { - 'retriever': { - 'label': options.label, - 'customOptions': None, - } - } - - -def test_retriever_action_metadata_with_supports_and_config_schema() -> None: - """Test for retriever_action_metadata with supports and config_schema.""" - - class CustomConfig(BaseModel): - k: int - - options = RetrieverOptions( - label='Advanced Retriever', - supports=RetrieverSupports(media=True), - config_schema=to_json_schema(CustomConfig), - ) - action_metadata = retriever_action_metadata( - name='advanced_retriever', - options=options, - ) - assert isinstance(action_metadata, ActionMetadata) - assert action_metadata.metadata is not None - metadata = cast(dict[str, Any], action_metadata.metadata) - assert metadata.get('retriever') is not None - retriever_meta = cast(dict[str, Any], metadata['retriever']) - assert retriever_meta['label'] == 'Advanced Retriever' - assert retriever_meta['supports'] == { - 'media': True, - } - assert retriever_meta['customOptions'] == { - 'title': 'CustomConfig', - 'type': 'object', - 'properties': { - 'k': {'title': 'K', 'type': 'integer'}, - }, - 'required': ['k'], - } - - -def test_retriever_action_metadata_no_options() -> None: - """Test retriever_action_metadata when no options are provided.""" - action_metadata = retriever_action_metadata(name='default_retriever') - assert isinstance(action_metadata, ActionMetadata) - assert action_metadata.metadata == {'retriever': {'customOptions': None}} - - -def test_create_retriever_ref_basic() -> None: - """Test basic creation of RetrieverRef.""" - ref = create_retriever_ref('my-retriever') - assert ref.name == 'my-retriever' - assert ref.config is None - assert ref.version is None - - -def test_create_retriever_ref_with_config() -> None: - """Test creation of RetrieverRef with configuration.""" - config = {'k': 5} - ref = create_retriever_ref('configured-retriever', config=config) - assert ref.name == 'configured-retriever' - assert ref.config == config - assert ref.version is None - - -def test_create_retriever_ref_with_version() -> None: - """Test creation of RetrieverRef with a version.""" - ref = create_retriever_ref('versioned-retriever', version='v1.0') - assert ref.name == 'versioned-retriever' - assert ref.config is None - assert ref.version == 'v1.0' - - -def test_create_retriever_ref_with_config_and_version() -> None: - """Test creation of RetrieverRef with both config and version.""" - config = {'k': 10} - ref = create_retriever_ref('full-retriever', config=config, version='beta') - assert ref.name == 'full-retriever' - assert ref.config == config - assert ref.version == 'beta' - - -@pytest.mark.asyncio -async def test_define_retriever() -> None: - """Test define_retriever registration.""" - registry = MagicMock() - fn = AsyncMock(return_value=RetrieverResponse(documents=[])) - - define_retriever(registry, 'test_retriever', fn) - - registry.register_action.assert_called_once() - call_args = registry.register_action.call_args - assert call_args.kwargs['kind'] == 'retriever' - assert call_args.kwargs['name'] == 'test_retriever' - - -@pytest.mark.asyncio -async def test_define_indexer() -> None: - """Test define_indexer registration.""" - registry = MagicMock() - fn = AsyncMock() - - define_indexer(registry, 'test_indexer', fn) - - registry.register_action.assert_called_once() - call_args = registry.register_action.call_args - assert call_args.kwargs['kind'] == 'indexer' - assert call_args.kwargs['name'] == 'test_indexer' - - -def test_indexer_action_metadata() -> None: - """Test for indexer_action_metadata with basic options.""" - options = IndexerOptions(label='Test Indexer') - action_metadata = indexer_action_metadata( - name='test_indexer', - options=options, - ) - - assert isinstance(action_metadata, ActionMetadata) - assert action_metadata.input_json_schema is not None - assert action_metadata.output_json_schema is not None - assert action_metadata.metadata == { - 'indexer': { - 'label': options.label, - 'customOptions': None, - } - } - - -def test_create_indexer_ref_basic() -> None: - """Test basic creation of IndexerRef.""" - ref = create_indexer_ref('my-indexer') - assert ref.name == 'my-indexer' - assert ref.config is None - assert ref.version is None diff --git a/py/packages/genkit/tests/genkit/codec_test.py b/py/packages/genkit/tests/genkit/codec_test.py deleted file mode 100644 index 97e61f870e..0000000000 --- a/py/packages/genkit/tests/genkit/codec_test.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for the codec module.""" - -from pydantic import BaseModel - -from genkit.codec import dump_dict, dump_json - - -def test_dump_json_basic() -> None: - """Test basic JSON serialization.""" - # Test dictionary - assert dump_json({'a': 1, 'b': 'test'}) == '{"a":1,"b":"test"}' - - # Test list - assert dump_json([1, 2, 3]) == '[1,2,3]' - - # Test nested structures - assert dump_json({'a': [1, 2], 'b': {'c': 3}}) == '{"a":[1,2],"b":{"c":3}}' - - -def test_dump_json_special_types() -> None: - """Test JSON serialization of special Python types.""" - # Test None - assert dump_json(None) == 'null' - - # Test boolean - assert dump_json(True) == 'true' - assert dump_json(False) == 'false' - - -def test_dump_json_numbers() -> None: - """Test JSON serialization of different number types.""" - # Test integers - assert dump_json(42) == '42' - - # Test floats - assert dump_json(3.14) == '3.14' - - # Test scientific notation - assert dump_json(1e-10) == '1e-10' - - -def test_dump_json_pydantic() -> None: - """Test JSON serialization of Pydantic models.""" - - class MyModel(BaseModel): - a: int - b: str - - assert dump_json(MyModel(a=1, b='test')) == '{"a":1,"b":"test"}' - - -def test_dump_dict_list_of_models() -> None: - """Test dump_dict with a list of Pydantic models.""" - - class MyModel(BaseModel): - a: int - - models = [MyModel(a=1), MyModel(a=2)] - # This was failing before the fix because dump_dict didn't recurse into lists - assert dump_dict(models) == [{'a': 1}, {'a': 2}] - - -def test_dump_dict_nested_models() -> None: - """Test dump_dict with nested structures containing Pydantic models.""" - - class MyModel(BaseModel): - a: int - - data = {'key': [MyModel(a=1)], 'other': MyModel(a=2)} - assert dump_dict(data) == {'key': [{'a': 1}], 'other': {'a': 2}} diff --git a/py/packages/genkit/tests/genkit/core/action_test.py b/py/packages/genkit/tests/genkit/core/action_test.py index 406857be7c..79610a4d04 100644 --- a/py/packages/genkit/tests/genkit/core/action_test.py +++ b/py/packages/genkit/tests/genkit/core/action_test.py @@ -5,20 +5,20 @@ """Tests for the action module.""" +import json from typing import cast import pytest -from genkit.codec import dump_json -from genkit.core.action import ( +from genkit._core._action import ( Action, + ActionKind, ActionRunContext, create_action_key, parse_action_key, parse_plugin_name_from_action_name, ) -from genkit.core.action.types import ActionKind -from genkit.core.error import GenkitError +from genkit._core._error import GenkitError def test_action_enum_behaves_like_str() -> None: @@ -34,8 +34,6 @@ def test_action_enum_behaves_like_str() -> None: assert ActionKind.FLOW == 'flow' assert ActionKind.MODEL == 'model' assert ActionKind.PROMPT == 'prompt' - assert ActionKind.RERANKER == 'reranker' - assert ActionKind.RETRIEVER == 'retriever' assert ActionKind.TOOL == 'tool' assert ActionKind.UTIL == 'util' @@ -79,108 +77,18 @@ def test_create_action_key() -> None: assert create_action_key(ActionKind.CUSTOM, 'foo') == '/custom/foo' assert create_action_key(ActionKind.MODEL, 'foo') == '/model/foo' assert create_action_key(ActionKind.PROMPT, 'foo') == '/prompt/foo' - assert create_action_key(ActionKind.RETRIEVER, 'foo') == '/retriever/foo' assert create_action_key(ActionKind.TOOL, 'foo') == '/tool/foo' assert create_action_key(ActionKind.UTIL, 'foo') == '/util/foo' -@pytest.mark.asyncio -async def test_define_sync_action() -> None: - """Define and run a sync action.""" +def test_sync_action_rejected() -> None: + """Sync functions are rejected - all actions must be async.""" def sync_foo() -> str: - """A sync action that returns 'syncFoo'.""" return 'syncFoo' - action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo) - - assert (await action.arun()).response == 'syncFoo' - assert sync_foo() == 'syncFoo' - - -@pytest.mark.asyncio -async def test_define_sync_action_with_input_without_type_annotation() -> None: - """Define and run a sync action with input without type annotation.""" - - def sync_foo(input: object) -> str: - """A sync action that returns 'syncFoo' with an input.""" - return f'syncFoo {input}' - - action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo) - - assert (await action.arun('foo')).response == 'syncFoo foo' - assert sync_foo('foo') == 'syncFoo foo' - - -@pytest.mark.asyncio -async def test_define_sync_action_with_input() -> None: - """Define and run a sync action with input.""" - - def sync_foo(input: str) -> str: - """A sync action that returns 'syncFoo' with an input.""" - return f'syncFoo {input}' - - action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo) - - assert (await action.arun('foo')).response == 'syncFoo foo' - assert sync_foo('foo') == 'syncFoo foo' - - -@pytest.mark.asyncio -async def test_define_sync_action_with_input_and_context() -> None: - """Define and run a sync action with input and context.""" - - def sync_foo(input: str, ctx: ActionRunContext) -> str: - """A sync action that returns 'syncFoo' with an input and context.""" - return f'syncFoo {input} {ctx.context["foo"]}' - - action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo) - - assert (await action.arun('foo', context={'foo': 'bar'})).response == 'syncFoo foo bar' - assert sync_foo('foo', ActionRunContext(context={'foo': 'bar'})) == 'syncFoo foo bar' - - -@pytest.mark.asyncio -async def test_define_sync_streaming_action() -> None: - """Define and run a sync streaming action.""" - - def sync_foo(input: str, ctx: ActionRunContext) -> int: - """A sync action that returns 'syncFoo' with streaming output.""" - ctx.send_chunk('1') - ctx.send_chunk('2') - return 3 - - action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo) - - chunks = [] - - def on_chunk(c: object) -> None: - chunks.append(c) - - assert (await action.arun('foo', context={'foo': 'bar'}, on_chunk=on_chunk)).response == 3 - assert chunks == ['1', '2'] - - -@pytest.mark.asyncio -async def test_define_streaming_action_and_stream_it() -> None: - """Define and stream a streaming action.""" - - def sync_foo(input: str, ctx: ActionRunContext) -> int: - """A sync action that returns 'syncFoo' with streaming output.""" - ctx.send_chunk('1') - ctx.send_chunk('2') - return 3 - - action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo) - - chunks = [] - - stream, response = action.stream('foo', context={'foo': 'bar'}) - async for chunk in stream: - chunks.append(chunk) - - assert (await response).response == 3 - assert chunks == ['1', '2'] + with pytest.raises(TypeError, match='Action handlers must be async functions'): + Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo) # type: ignore[arg-type] @pytest.mark.asyncio @@ -193,7 +101,7 @@ async def async_foo() -> str: action = Action(name='asyncFoo', kind=ActionKind.CUSTOM, fn=async_foo) - assert (await action.arun()).response == 'asyncFoo' + assert (await action.run()).response == 'asyncFoo' assert (await async_foo()) == 'asyncFoo' @@ -207,7 +115,7 @@ async def async_foo(input: str) -> str: action = Action(name='asyncFoo', kind=ActionKind.CUSTOM, fn=async_foo) - assert (await action.arun('foo')).response == 'asyncFoo foo' + assert (await action.run('foo')).response == 'asyncFoo foo' assert (await async_foo('foo')) == 'asyncFoo foo' @@ -221,28 +129,51 @@ async def async_foo(input: str, ctx: ActionRunContext) -> str: action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=async_foo) - assert (await action.arun('foo', context={'foo': 'bar'})).response == 'syncFoo foo bar' + assert (await action.run('foo', context={'foo': 'bar'})).response == 'syncFoo foo bar' assert (await async_foo('foo', ActionRunContext(context={'foo': 'bar'}))) == 'syncFoo foo bar' @pytest.mark.asyncio -async def test_define_async_streaming_action() -> None: - """Define and run an async streaming action.""" +async def test_streaming_action_with_callback() -> None: + """Streaming action with on_chunk callback.""" - async def async_foo(input: str, ctx: ActionRunContext) -> int: - """An async action that returns 'syncFoo' with streaming output.""" + async def foo( + input: str, + ctx: ActionRunContext, + ) -> int: ctx.send_chunk('1') ctx.send_chunk('2') return 3 - action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=async_foo) + action = Action(name='foo', kind=ActionKind.CUSTOM, fn=foo) - chunks = [] + chunks: list[object] = [] + result = await action.run('foo', on_chunk=chunks.append) + + assert result.response == 3 + assert chunks == ['1', '2'] - def on_chunk(c: object) -> None: - chunks.append(c) - assert (await action.arun('foo', context={'foo': 'bar'}, on_chunk=on_chunk)).response == 3 +@pytest.mark.asyncio +async def test_streaming_action_with_stream_method() -> None: + """Streaming action using the stream() method.""" + + async def foo( + input: str, + ctx: ActionRunContext, + ) -> int: + ctx.send_chunk('1') + ctx.send_chunk('2') + return 3 + + action = Action(name='foo', kind=ActionKind.CUSTOM, fn=foo) + + chunks: list[object] = [] + result = action.stream('foo') + async for chunk in result.stream: + chunks.append(chunk) + + assert await result.response == 3 assert chunks == ['1', '2'] @@ -258,55 +189,38 @@ async def test_propagates_context_via_contextvar() -> None: """Context is properly propagated via contextvar.""" async def foo(_: str | None, ctx: ActionRunContext) -> str: - return dump_json(ctx.context) + return json.dumps(ctx.context) foo_action = cast(Action[str | None, str], Action(name='foo', kind=ActionKind.CUSTOM, fn=foo)) async def bar() -> str: - return (await foo_action.arun()).response + return (await foo_action.run()).response bar_action = cast(Action[None, str], Action(name='bar', kind=ActionKind.CUSTOM, fn=bar)) async def baz() -> str: - return (await bar_action.arun()).response + return (await bar_action.run()).response baz_action = cast(Action[None, str], Action(name='baz', kind=ActionKind.CUSTOM, fn=baz)) - first = baz_action.arun(context={'foo': 'bar'}) - second = baz_action.arun(context={'bar': 'baz'}) - - assert (await second).response == '{"bar":"baz"}' - assert (await first).response == '{"foo":"bar"}' - - -@pytest.mark.asyncio -async def test_sync_action_raises_errors() -> None: - """Sync action raises error with necessary metadata.""" + first = baz_action.run(context={'foo': 'bar'}) + second = baz_action.run(context={'bar': 'baz'}) - def sync_foo(_: str | None, ctx: ActionRunContext) -> None: - raise Exception('oops') - - action = Action(name='fooAction', kind=ActionKind.CUSTOM, fn=sync_foo) - - with pytest.raises(GenkitError, match=r'.*Error while running action fooAction.*') as e: - await action.arun() - - assert 'stack' in e.value.details - assert 'trace_id' in e.value.details - assert str(e.value.cause) == 'oops' + assert (await second).response == '{"bar": "baz"}' + assert (await first).response == '{"foo": "bar"}' @pytest.mark.asyncio -async def test_async_action_raises_errors() -> None: - """Async action raises error with necessary metadata.""" +async def test_action_raises_errors() -> None: + """Action raises error with necessary metadata.""" - async def async_foo(_: str | None, ctx: ActionRunContext) -> None: + async def foo(_: str | None, ctx: ActionRunContext) -> None: raise Exception('oops') - action = Action(name='fooAction', kind=ActionKind.CUSTOM, fn=async_foo) + action = Action(name='fooAction', kind=ActionKind.CUSTOM, fn=foo) with pytest.raises(GenkitError, match=r'.*Error while running action fooAction.*') as e: - await action.arun() + await action.run() assert 'stack' in e.value.details assert 'trace_id' in e.value.details @@ -314,8 +228,8 @@ async def async_foo(_: str | None, ctx: ActionRunContext) -> None: @pytest.mark.asyncio -async def test_arun_raw_raises_on_none_input_when_input_required() -> None: - """arun_raw raises GenkitError when input is None but the action requires it.""" +async def test_run_raises_on_none_input_when_input_required() -> None: + """run() raises GenkitError when input is None but the action requires it.""" async def typed_fn(input: str) -> str: return f'got {input}' @@ -323,30 +237,30 @@ async def typed_fn(input: str) -> str: action = Action(name='typedAction', kind=ActionKind.CUSTOM, fn=typed_fn) with pytest.raises(GenkitError, match=r'.*requires input but none was provided.*'): - await action.arun_raw(raw_input=None) + await action.run(input=None) @pytest.mark.asyncio -async def test_arun_raw_succeeds_with_valid_input() -> None: - """arun_raw succeeds when valid input is provided.""" +async def test_run_succeeds_with_valid_input() -> None: + """run() succeeds when valid input is provided.""" async def typed_fn(input: str) -> str: return f'got {input}' action = Action(name='typedAction', kind=ActionKind.CUSTOM, fn=typed_fn) - result = await action.arun_raw(raw_input='hello') + result = await action.run(input='hello') assert result.response == 'got hello' @pytest.mark.asyncio -async def test_arun_raw_no_input_type_allows_none() -> None: - """arun_raw allows None input when action has no input type.""" +async def test_run_no_input_type_allows_none() -> None: + """run() allows None input when action has no input type.""" async def no_input_fn() -> str: return 'no input needed' action = Action(name='noInputAction', kind=ActionKind.CUSTOM, fn=no_input_fn) - result = await action.arun_raw(raw_input=None) + result = await action.run(input=None) assert result.response == 'no input needed' diff --git a/py/packages/genkit/tests/genkit/aio/channel_test.py b/py/packages/genkit/tests/genkit/core/channel_test.py similarity index 99% rename from py/packages/genkit/tests/genkit/aio/channel_test.py rename to py/packages/genkit/tests/genkit/core/channel_test.py index ff2703b74c..3981e0ccf8 100644 --- a/py/packages/genkit/tests/genkit/aio/channel_test.py +++ b/py/packages/genkit/tests/genkit/core/channel_test.py @@ -23,7 +23,7 @@ import pytest -from genkit.aio import Channel +from genkit._core._channel import Channel T = TypeVar('T') diff --git a/py/packages/genkit/tests/genkit/core/endpoints/reflection_test.py b/py/packages/genkit/tests/genkit/core/endpoints/reflection_test.py index 8f03750d61..d98659b16e 100644 --- a/py/packages/genkit/tests/genkit/core/endpoints/reflection_test.py +++ b/py/packages/genkit/tests/genkit/core/endpoints/reflection_test.py @@ -38,16 +38,15 @@ from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any, cast -from unittest.mock import ANY, AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock import pytest import pytest_asyncio from httpx import ASGITransport, AsyncClient -from genkit.core.action import ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.reflection import create_reflection_asgi_app -from genkit.core.registry import Registry +from genkit._core._action import ActionKind, ActionMetadata +from genkit._core._reflection import create_reflection_asgi_app +from genkit._core._registry import Registry @pytest.fixture @@ -142,7 +141,7 @@ async def test_run_action_standard(asgi_client: AsyncClient, mock_registry: Magi mock_output.span_id = 'test_span_id' async def side_effect( - raw_input: object, + input: object = None, on_chunk: object | None = None, context: object | None = None, on_trace_start: Callable[[str, str], None] | None = None, @@ -152,7 +151,7 @@ async def side_effect( on_trace_start('test_trace_id', 'test_span_id') return mock_output - mock_action.arun_raw.side_effect = side_effect + mock_action.run.side_effect = side_effect async def mock_resolve_action_by_key(key: str) -> AsyncMock: return mock_action @@ -169,7 +168,7 @@ async def mock_resolve_action_by_key(key: str) -> AsyncMock: assert response_data['telemetry']['spanId'] == 'test_span_id' assert response.headers['X-Genkit-Trace-Id'] == 'test_trace_id' assert response.headers['X-Genkit-Span-Id'] == 'test_span_id' - mock_action.arun_raw.assert_called_once_with(raw_input={'data': 'test'}, context={}, on_trace_start=ANY) + mock_action.run.assert_called_once_with(input={'data': 'test'}, context={}, on_trace_start=ANY, on_chunk=None) @pytest.mark.asyncio @@ -180,7 +179,7 @@ async def test_run_action_with_context(asgi_client: AsyncClient, mock_registry: mock_output.response = {'result': 'success'} mock_output.trace_id = 'test_trace_id' mock_output.span_id = 'test_span_id' - mock_action.arun_raw.return_value = mock_output + mock_action.run.return_value = mock_output async def mock_resolve_action_by_key(key: str) -> AsyncMock: return mock_action @@ -197,26 +196,24 @@ async def mock_resolve_action_by_key(key: str) -> AsyncMock: ) assert response.status_code == 200 - mock_action.arun_raw.assert_called_once_with( - raw_input={'data': 'test'}, + mock_action.run.assert_called_once_with( + input={'data': 'test'}, context={'user': 'test_user'}, on_trace_start=ANY, + on_chunk=None, ) @pytest.mark.asyncio -@patch('genkit.core.reflection.is_streaming_requested') async def test_run_action_streaming( - mock_is_streaming: MagicMock, asgi_client: AsyncClient, mock_registry: MagicMock, ) -> None: """Test that streaming actions work correctly.""" - mock_is_streaming.return_value = True mock_action = AsyncMock() async def mock_streaming( - raw_input: object, + input: object = None, on_chunk: object | None = None, context: object | None = None, on_trace_start: Callable[[str, str], None] | None = None, @@ -234,7 +231,7 @@ async def mock_streaming( mock_output.span_id = 'stream_span_id' return mock_output - mock_action.arun_raw.side_effect = mock_streaming + mock_action.run.side_effect = mock_streaming mock_registry.resolve_action_by_key.return_value = mock_action response = await asgi_client.post( @@ -245,4 +242,3 @@ async def mock_streaming( assert response.status_code == 200 assert response.headers['X-Genkit-Trace-Id'] == 'stream_trace_id' assert response.headers['X-Genkit-Span-Id'] == 'stream_span_id' - assert mock_is_streaming.called diff --git a/py/packages/genkit/tests/genkit/core/environment_test.py b/py/packages/genkit/tests/genkit/core/environment_test.py index e1b817afa6..2cab0df0f5 100644 --- a/py/packages/genkit/tests/genkit/core/environment_test.py +++ b/py/packages/genkit/tests/genkit/core/environment_test.py @@ -8,12 +8,11 @@ import os from unittest import mock -from genkit.core.environment import ( - EnvVar, +from genkit._core._environment import ( + GENKIT_ENV, GenkitEnvironment, get_current_environment, is_dev_environment, - is_prod_environment, ) @@ -28,33 +27,14 @@ def test_is_dev_environment() -> None: assert not is_dev_environment() # Test when GENKIT_ENV is set to 'dev' - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.DEV}): assert is_dev_environment() # Test when GENKIT_ENV is set to something else - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.PROD}): assert not is_dev_environment() -def test_is_prod_environment() -> None: - """Test the is_prod_environment function. - - Verifies that the is_prod_environment function correctly detects - production environments based on environment variables. - """ - # Test when GENKIT_ENV is not set - with mock.patch.dict(os.environ, clear=True): - assert is_prod_environment() - - # Test when GENKIT_ENV is set to 'prod' - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}): - assert is_prod_environment() - - # Test when GENKIT_ENV is set to something else - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}): - assert not is_prod_environment() - - def test_get_current_environment() -> None: """Test the get_current_environment function. @@ -66,13 +46,13 @@ def test_get_current_environment() -> None: assert get_current_environment() == GenkitEnvironment.PROD # Test when GENKIT_ENV is set to 'prod' - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.PROD}): assert get_current_environment() == GenkitEnvironment.PROD # Test when GENKIT_ENV is set to 'dev' - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.DEV}): assert get_current_environment() == GenkitEnvironment.DEV # Test when GENKIT_ENV is set to something else - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: 'invalid'}): + with mock.patch.dict(os.environ, {GENKIT_ENV: 'invalid'}): assert get_current_environment() == GenkitEnvironment.PROD diff --git a/py/packages/genkit/tests/genkit/core/error_test.py b/py/packages/genkit/tests/genkit/core/error_test.py index ed1d57db2a..c16467b6f4 100644 --- a/py/packages/genkit/tests/genkit/core/error_test.py +++ b/py/packages/genkit/tests/genkit/core/error_test.py @@ -16,22 +16,18 @@ """Unit tests for the error module.""" -from genkit.core.error import ( +from genkit._core._error import ( GenkitError, HttpErrorWireFormat, + PublicError, ReflectionError, - UnstableApiError, - UserFacingError, get_callable_json, - get_error_message, get_error_stack, get_http_status, ) -# New tests start here def test_genkit_error() -> None: - """Test that creating a GenkitError works.""" error = GenkitError( status='INVALID_ARGUMENT', message='Test message', @@ -45,13 +41,11 @@ def test_genkit_error() -> None: assert error.source == 'test_source' assert str(error) == 'test_source: INVALID_ARGUMENT: Test message' - # Test without source error_no_source = GenkitError(status='INTERNAL', message='Test message 2') assert str(error_no_source) == 'INTERNAL: Test message 2' def test_genkit_error_to_json() -> None: - """Test that GenkitError can be serialized to JSON.""" error = GenkitError(status='NOT_FOUND', message='Resource not found', details={'id': 123}) serializable = error.to_serializable() assert isinstance(serializable, ReflectionError) @@ -61,20 +55,8 @@ def test_genkit_error_to_json() -> None: assert serializable.details.model_dump()['id'] == 123 -def test_unstable_api_error() -> None: - """Test that creating an UnstableApiError works.""" - error = UnstableApiError(level='alpha', message='Test feature') - assert error.status == 'FAILED_PRECONDITION' - assert 'Test feature' in error.original_message - assert "This API requires 'alpha' stability level" in error.original_message - - error_no_message = UnstableApiError() - assert "This API requires 'beta' stability level" in error_no_message.original_message - - -def test_user_facing_error() -> None: - """Test creating a UserFacingError.""" - error = UserFacingError( +def test_public_error() -> None: + error = PublicError( status='UNAUTHENTICATED', message='Please log in', details={'extra_msg': 'Session expired'}, @@ -85,7 +67,6 @@ def test_user_facing_error() -> None: def test_get_http_status() -> None: - """Test that get_http_status returns the correct HTTP status code.""" genkit_error = GenkitError(status='PERMISSION_DENIED', message='No access') assert get_http_status(genkit_error) == 403 @@ -94,7 +75,6 @@ def test_get_http_status() -> None: def test_get_callable_json() -> None: - """Test that get_callable_json returns the correct JSON data.""" genkit_error = GenkitError(status='DATA_LOSS', message='Oops') json_data = get_callable_json(genkit_error) assert isinstance(json_data, HttpErrorWireFormat) @@ -108,21 +88,9 @@ def test_get_callable_json() -> None: assert json_data.message == 'Type error' -def test_get_error_message() -> None: - """Test that get_error_message returns the correct error message.""" - error_message = get_error_message(ValueError('Test Value Error')) - assert error_message == 'Test Value Error' - - error_message = get_error_message('Test String Error') - assert error_message == 'Test String Error' - - def test_get_error_stack() -> None: - """Test that get_error_stack returns the correct error stack.""" try: raise ValueError('Example Error') except ValueError as e: - # get_error_stack returns an empty string to keep Dev UI clean. - # While this may limit debuggability, it satisfies the required display format. tb = get_error_stack(e) assert tb == '' diff --git a/py/packages/genkit/tests/genkit/core/extract_test.py b/py/packages/genkit/tests/genkit/core/extract_test.py index e8ae56668b..f8c0842153 100644 --- a/py/packages/genkit/tests/genkit/core/extract_test.py +++ b/py/packages/genkit/tests/genkit/core/extract_test.py @@ -21,12 +21,12 @@ import pytest -from genkit.core.extract import extract_items, extract_json, parse_partial_json +from genkit._core._extract_json import extract_json, extract_json_array_from_text, parse_partial_json # TODO(#4356): consider extracting these tests into shared yaml spec. They are already # duplicated in js/ai/tests/extract_test.ts -test_cases_extract_items = [ +test_cases_extract_json_array_from_text = [ ( 'handles simple array in chunks', [ @@ -85,16 +85,16 @@ @pytest.mark.parametrize( 'name, steps', - test_cases_extract_items, - ids=[tc[0] for tc in test_cases_extract_items], + test_cases_extract_json_array_from_text, + ids=[tc[0] for tc in test_cases_extract_json_array_from_text], ) -def test_extract_items(name: str, steps: list[dict[str, Any]]) -> None: +def test_extract_json_array_from_text(name: str, steps: list[dict[str, Any]]) -> None: """Test extraction of incomplete json that can be fixed.""" text = '' cursor = 0 for step in steps: text += step['chunk'] - result = extract_items(text, cursor) + result = extract_json_array_from_text(text, cursor) assert result.items == step['want'] cursor = result.cursor diff --git a/py/packages/genkit/tests/genkit/core/http_client_test.py b/py/packages/genkit/tests/genkit/core/http_client_test.py index 59ffd46996..7cbaabe3bd 100644 --- a/py/packages/genkit/tests/genkit/core/http_client_test.py +++ b/py/packages/genkit/tests/genkit/core/http_client_test.py @@ -14,230 +14,169 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Tests for genkit.core.http_client module. - -This module tests the per-event-loop HTTP client caching functionality, -which is critical for avoiding "bound to different event loop" errors -when using httpx.AsyncClient in async code. - -Test Categories: - - Client caching: Verify clients are cached and reused per event loop - - Cache key isolation: Different cache keys get different clients - - Closed client handling: Closed clients are replaced - - Cleanup utilities: close_cached_clients and clear_client_cache work - - Configuration: Headers and timeouts are applied correctly -""" - -import asyncio +"""Tests for HTTP client caching.""" import httpx import pytest -from genkit.core.http_client import ( - _loop_clients, +from genkit._core._http_client import ( clear_client_cache, close_cached_clients, get_cached_client, ) -class TestGetCachedClient: - """Tests for get_cached_client function.""" - - @pytest.fixture(autouse=True) - def clear_cache(self) -> None: - """Clear the client cache before each test.""" - clear_client_cache() - - @pytest.mark.asyncio - async def test_returns_httpx_async_client(self) -> None: - """Verify get_cached_client returns an httpx.AsyncClient instance.""" - client = get_cached_client(cache_key='test') - - assert isinstance(client, httpx.AsyncClient) - - @pytest.mark.asyncio - async def test_client_cached_per_event_loop(self) -> None: - """Verify the same client is returned for the same cache key in same loop.""" - client1 = get_cached_client(cache_key='test') - client2 = get_cached_client(cache_key='test') - - assert client1 is client2 +@pytest.fixture(autouse=True) +def clear_cache() -> None: + clear_client_cache() - @pytest.mark.asyncio - async def test_different_cache_keys_get_different_clients(self) -> None: - """Verify different cache keys return different client instances.""" - client1 = get_cached_client(cache_key='plugin-a') - client2 = get_cached_client(cache_key='plugin-b') - assert client1 is not client2 +@pytest.mark.asyncio +async def test_returns_httpx_async_client() -> None: + client = get_cached_client(cache_key='test') + assert isinstance(client, httpx.AsyncClient) - @pytest.mark.asyncio - async def test_client_has_correct_headers(self) -> None: - """Verify client is configured with provided headers.""" - headers = { - 'Authorization': 'Bearer test-token', - 'X-Custom-Header': 'custom-value', - } - client = get_cached_client(cache_key='test', headers=headers) - # Check headers are set on the client - assert client.headers.get('Authorization') == 'Bearer test-token' - assert client.headers.get('X-Custom-Header') == 'custom-value' +@pytest.mark.asyncio +async def test_client_cached_per_event_loop() -> None: + client1 = get_cached_client(cache_key='test') + client2 = get_cached_client(cache_key='test') + assert client1 is client2 - @pytest.mark.asyncio - async def test_client_has_correct_timeout_float(self) -> None: - """Verify client is configured with float timeout.""" - client = get_cached_client(cache_key='test', timeout=30.0) - assert client.timeout.read == 30.0 - assert client.timeout.connect == 30.0 +@pytest.mark.asyncio +async def test_different_cache_keys_get_different_clients() -> None: + client1 = get_cached_client(cache_key='plugin-a') + client2 = get_cached_client(cache_key='plugin-b') + assert client1 is not client2 - @pytest.mark.asyncio - async def test_client_has_correct_timeout_object(self) -> None: - """Verify client is configured with httpx.Timeout object.""" - timeout = httpx.Timeout(60.0, connect=10.0) - client = get_cached_client(cache_key='test', timeout=timeout) - assert client.timeout.read == 60.0 - assert client.timeout.connect == 10.0 +@pytest.mark.asyncio +async def test_client_has_correct_headers() -> None: + headers = {'Authorization': 'Bearer test-token', 'X-Custom': 'value'} + client = get_cached_client(cache_key='test', headers=headers) + assert client.headers.get('Authorization') == 'Bearer test-token' + assert client.headers.get('X-Custom') == 'value' - @pytest.mark.asyncio - async def test_default_timeout_applied(self) -> None: - """Verify default timeout is applied when not specified.""" - client = get_cached_client(cache_key='test') - # Default is 60s total, 10s connect - assert client.timeout.read == 60.0 - assert client.timeout.connect == 10.0 +@pytest.mark.asyncio +async def test_client_has_correct_timeout_float() -> None: + client = get_cached_client(cache_key='test', timeout=30.0) + assert client.timeout.read == 30.0 + assert client.timeout.connect == 30.0 - @pytest.mark.asyncio - async def test_closed_client_gets_replaced(self) -> None: - """Verify a closed client is replaced with a new one.""" - client1 = get_cached_client(cache_key='test') - await client1.aclose() - assert client1.is_closed +@pytest.mark.asyncio +async def test_client_has_correct_timeout_object() -> None: + timeout = httpx.Timeout(60.0, connect=10.0) + client = get_cached_client(cache_key='test', timeout=timeout) + assert client.timeout.read == 60.0 + assert client.timeout.connect == 10.0 - client2 = get_cached_client(cache_key='test') - assert client2 is not client1 - assert not client2.is_closed +@pytest.mark.asyncio +async def test_default_timeout_applied() -> None: + client = get_cached_client(cache_key='test') + assert client.timeout.read == 60.0 + assert client.timeout.connect == 10.0 - @pytest.mark.asyncio - async def test_client_stored_in_cache(self) -> None: - """Verify client is stored in the module-level cache.""" - clear_client_cache() - _ = get_cached_client(cache_key='test') +@pytest.mark.asyncio +async def test_closed_client_gets_replaced() -> None: + client1 = get_cached_client(cache_key='test') + await client1.aclose() + assert client1.is_closed - loop = asyncio.get_running_loop() - assert loop in _loop_clients - assert 'test' in _loop_clients[loop] + client2 = get_cached_client(cache_key='test') + assert client2 is not client1 + assert not client2.is_closed - def test_raises_without_running_event_loop(self) -> None: - """Verify RuntimeError is raised when called outside async context.""" - with pytest.raises(RuntimeError, match='must be called from within an async context'): - get_cached_client(cache_key='test') +@pytest.mark.asyncio +async def test_client_stored_in_cache() -> None: + clear_client_cache() + client = get_cached_client(cache_key='test') + client2 = get_cached_client(cache_key='test') + assert client is client2 -class TestCloseCachedClients: - """Tests for close_cached_clients function.""" - @pytest.fixture(autouse=True) - def clear_cache(self) -> None: - """Clear the client cache before each test.""" - clear_client_cache() +def test_raises_without_running_event_loop() -> None: + with pytest.raises(RuntimeError, match='no running event loop'): + get_cached_client(cache_key='test') - @pytest.mark.asyncio - async def test_close_specific_client(self) -> None: - """Verify closing a specific client by cache key.""" - _ = get_cached_client(cache_key='to-close') - _ = get_cached_client(cache_key='keep') - await close_cached_clients('to-close') - - # The closed client should be removed from cache - loop = asyncio.get_running_loop() - assert 'to-close' not in _loop_clients[loop] - assert 'keep' in _loop_clients[loop] +@pytest.mark.asyncio +async def test_close_specific_client() -> None: + client_to_close = get_cached_client(cache_key='to-close') + client_keep = get_cached_client(cache_key='keep') - @pytest.mark.asyncio - async def test_close_all_clients_in_loop(self) -> None: - """Verify closing all clients in current event loop.""" - _ = get_cached_client(cache_key='client-a') - _ = get_cached_client(cache_key='client-b') + await close_cached_clients('to-close') - await close_cached_clients() + # 'to-close' should be gone; new fetch returns new client + client_after = get_cached_client(cache_key='to-close') + assert client_after is not client_to_close + # 'keep' should still be cached + assert get_cached_client(cache_key='keep') is client_keep - # All clients should be removed from this loop's cache - loop = asyncio.get_running_loop() - assert loop not in _loop_clients or len(_loop_clients[loop]) == 0 - - @pytest.mark.asyncio - async def test_close_nonexistent_key_is_noop(self) -> None: - """Verify closing a non-existent key doesn't raise.""" - _ = get_cached_client(cache_key='exists') - # Should not raise - await close_cached_clients('does-not-exist') +@pytest.mark.asyncio +async def test_close_all_clients_in_loop() -> None: + client_a = get_cached_client(cache_key='client-a') + client_b = get_cached_client(cache_key='client-b') - @pytest.mark.asyncio - async def test_close_when_no_clients_is_noop(self) -> None: - """Verify closing when no clients exist doesn't raise.""" - # Should not raise - await close_cached_clients() + await close_cached_clients() + # Cache should be empty; new fetches return new clients + new_a = get_cached_client(cache_key='client-a') + new_b = get_cached_client(cache_key='client-b') + assert new_a is not client_a + assert new_b is not client_b -class TestClearClientCache: - """Tests for clear_client_cache function.""" - @pytest.mark.asyncio - async def test_clear_removes_all_cached_clients(self) -> None: - """Verify clear_client_cache removes all clients from cache.""" - _ = get_cached_client(cache_key='client-a') - _ = get_cached_client(cache_key='client-b') +@pytest.mark.asyncio +async def test_close_nonexistent_key_is_noop() -> None: + _ = get_cached_client(cache_key='exists') + await close_cached_clients('does-not-exist') - clear_client_cache() - assert len(_loop_clients) == 0 +@pytest.mark.asyncio +async def test_close_when_no_clients_is_noop() -> None: + await close_cached_clients() - def test_clear_when_empty_is_noop(self) -> None: - """Verify clearing an empty cache doesn't raise.""" - clear_client_cache() - clear_client_cache() # Should not raise +@pytest.mark.asyncio +async def test_clear_removes_all_cached_clients() -> None: + client_a = get_cached_client(cache_key='client-a') + client_b = get_cached_client(cache_key='client-b') + clear_client_cache() + # Cache cleared; new fetches return new clients + new_a = get_cached_client(cache_key='client-a') + new_b = get_cached_client(cache_key='client-b') + assert new_a is not client_a + assert new_b is not client_b -class TestMultipleEventLoops: - """Tests for behavior across multiple event loops. - Note: These tests verify the cache key isolation but cannot fully test - different event loops in the same test due to pytest-asyncio limitations. - The WeakKeyDictionary behavior ensures proper isolation in production. - """ +def test_clear_when_empty_is_noop() -> None: + clear_client_cache() + clear_client_cache() - @pytest.fixture(autouse=True) - def clear_cache(self) -> None: - """Clear the client cache before each test.""" - clear_client_cache() - @pytest.mark.asyncio - async def test_cache_uses_current_event_loop_as_key(self) -> None: - """Verify the cache is keyed by the current event loop.""" - _ = get_cached_client(cache_key='test') +@pytest.mark.asyncio +async def test_cache_uses_current_event_loop_as_key() -> None: + client1 = get_cached_client(cache_key='test') + client2 = get_cached_client(cache_key='test') + assert client1 is client2 - loop = asyncio.get_running_loop() - assert loop in _loop_clients - @pytest.mark.asyncio - async def test_multiple_cache_keys_same_loop(self) -> None: - """Verify multiple cache keys can coexist in the same loop.""" - client_a = get_cached_client(cache_key='plugin-a') - client_b = get_cached_client(cache_key='plugin-b') - client_c = get_cached_client(cache_key='plugin-c') +@pytest.mark.asyncio +async def test_multiple_cache_keys_same_loop() -> None: + client_a = get_cached_client(cache_key='plugin-a') + client_b = get_cached_client(cache_key='plugin-b') + client_c = get_cached_client(cache_key='plugin-c') - loop = asyncio.get_running_loop() - assert len(_loop_clients[loop]) == 3 - assert _loop_clients[loop]['plugin-a'] is client_a - assert _loop_clients[loop]['plugin-b'] is client_b - assert _loop_clients[loop]['plugin-c'] is client_c + assert client_a is not client_b + assert client_b is not client_c + assert client_a is not client_c + assert get_cached_client(cache_key='plugin-a') is client_a + assert get_cached_client(cache_key='plugin-b') is client_b + assert get_cached_client(cache_key='plugin-c') is client_c diff --git a/py/packages/genkit/tests/genkit/core/latency_test.py b/py/packages/genkit/tests/genkit/core/latency_test.py index a759b61dbd..d647d0177a 100644 --- a/py/packages/genkit/tests/genkit/core/latency_test.py +++ b/py/packages/genkit/tests/genkit/core/latency_test.py @@ -6,14 +6,12 @@ """Tests for latency tracking in actions.""" import asyncio -import time from typing import cast import pytest from pydantic import BaseModel -from genkit.core.action import Action -from genkit.core.action.types import ActionKind +from genkit._core._action import Action, ActionKind class MockResponse(BaseModel): @@ -36,30 +34,13 @@ async def async_model_fn(input: str) -> MockResponse: # if we want to test sync wrapper or just await a task. action = cast(Action[str, MockResponse], Action(name='testModel', kind=ActionKind.MODEL, fn=async_model_fn)) - response = await action.arun('world') + response = await action.run('world') assert response.response.value == 'hello world' assert response.response.latency_ms is not None assert response.response.latency_ms >= 100 # Should be at least 100ms due to sleep -def test_sync_action_latency_ms_population() -> None: - """Verify that latency_ms is automatically populated for sync actions.""" - - def sync_model_fn(input: str) -> MockResponse: - time.sleep(0.1) - return MockResponse(value=f'sync hello {input}') - - action = cast(Action[str, MockResponse], Action(name='syncTestModel', kind=ActionKind.CUSTOM, fn=sync_model_fn)) - - # run() is sync - response = action.run('world') - - assert response.response.value == 'sync hello world' - assert response.response.latency_ms is not None - assert response.response.latency_ms >= 100 - - class ImmutableMockResponse(BaseModel): """A mock response object for testing latency tracking with frozen models.""" @@ -80,7 +61,7 @@ async def async_model_fn(input: str) -> ImmutableMockResponse: Action(name='testImmutableModel', kind=ActionKind.MODEL, fn=async_model_fn), ) - response = await action.arun('world') + response = await action.run('world') assert response.response.value == 'hello world' assert response.response.latency_ms is not None @@ -115,7 +96,7 @@ async def async_model_fn(input: str) -> ReadOnlyMockResponse: # will still try to use setattr if the field exists, or it won't have latency_ms in its fields. # Actually, Pydantic's model_copy only updates fields. latency_ms is a property here. - response = await action.arun('world') + response = await action.run('world') assert response.response.value == 'hello world' # Since it's a property without a setter AND not a Pydantic field, diff --git a/py/packages/genkit/tests/genkit/core/registry_test.py b/py/packages/genkit/tests/genkit/core/registry_test.py index 261e6dea6e..43cf4452f5 100644 --- a/py/packages/genkit/tests/genkit/core/registry_test.py +++ b/py/packages/genkit/tests/genkit/core/registry_test.py @@ -11,17 +11,20 @@ import pytest -from genkit.ai import Genkit, Plugin -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.registry import Registry +from genkit import Genkit, Plugin +from genkit._core._action import Action, ActionKind, ActionMetadata +from genkit._core._registry import Registry + + +async def _identity(x: object) -> object: + return x @pytest.mark.asyncio async def test_register_action_with_name_and_kind() -> None: """Ensure we can register an action with a name and kind.""" registry = Registry() - action = registry.register_action(name='test_action', kind=ActionKind.CUSTOM, fn=lambda x: x) + action = registry.register_action(name='test_action', kind=ActionKind.CUSTOM, fn=_identity) got = await registry.resolve_action(ActionKind.CUSTOM, 'test_action') assert got == action @@ -34,7 +37,7 @@ async def test_register_action_with_name_and_kind() -> None: async def test_resolve_action_by_key() -> None: """Ensure we can resolve an action by its key.""" registry = Registry() - action = registry.register_action(name='test_action', kind=ActionKind.CUSTOM, fn=lambda x: x) + action = registry.register_action(name='test_action', kind=ActionKind.CUSTOM, fn=_identity) got = await registry.resolve_action_by_key('/custom/test_action') assert got == action @@ -66,7 +69,7 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action: nonlocal resolver_calls resolver_calls.append([action_type, name]) - def model_fn() -> None: + async def model_fn() -> None: pass return Action(name=name, fn=model_fn, kind=action_type) @@ -119,10 +122,13 @@ async def self_resolving_factory() -> None: # _trigger_lazy_loading again. Without the guard, infinite recursion. await registry.resolve_action(ActionKind.CUSTOM, 'self_ref') + async def noop() -> None: + pass + action = registry.register_action( kind=ActionKind.CUSTOM, name='self_ref', - fn=lambda: None, + fn=noop, metadata={'lazy': True}, ) setattr(action, '_async_factory', self_resolving_factory) # noqa: B010 diff --git a/py/packages/genkit/tests/genkit/core/schema_test.py b/py/packages/genkit/tests/genkit/core/schema_test.py index 2a4662e331..07dcced8ef 100644 --- a/py/packages/genkit/tests/genkit/core/schema_test.py +++ b/py/packages/genkit/tests/genkit/core/schema_test.py @@ -21,7 +21,7 @@ import pytest from pydantic import BaseModel, Field -from genkit.core.schema import to_json_schema +from genkit._core._schema import to_json_schema def test_to_json_schema_pydantic_model() -> None: diff --git a/py/packages/genkit/tests/genkit/core/status_types_test.py b/py/packages/genkit/tests/genkit/core/status_types_test.py index ce90c560a0..ff0681479b 100644 --- a/py/packages/genkit/tests/genkit/core/status_types_test.py +++ b/py/packages/genkit/tests/genkit/core/status_types_test.py @@ -19,7 +19,7 @@ import pytest from pydantic import ValidationError -from genkit.core.status_types import Status, StatusCodes, http_status_code +from genkit._core._error import Status, StatusCodes, http_status_code def test_status_codes_values() -> None: diff --git a/py/packages/genkit/tests/genkit/core/trace/__init__.py b/py/packages/genkit/tests/genkit/core/trace/__init__.py index be6860617c..284d9abae8 100644 --- a/py/packages/genkit/tests/genkit/core/trace/__init__.py +++ b/py/packages/genkit/tests/genkit/core/trace/__init__.py @@ -14,4 +14,4 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Tests for the genkit.core.trace module.""" +"""Tests for the genkit._core._trace module.""" diff --git a/py/packages/genkit/tests/genkit/core/trace/adjusting_exporter_test.py b/py/packages/genkit/tests/genkit/core/trace/adjusting_exporter_test.py index af29400c09..8ea0996da0 100644 --- a/py/packages/genkit/tests/genkit/core/trace/adjusting_exporter_test.py +++ b/py/packages/genkit/tests/genkit/core/trace/adjusting_exporter_test.py @@ -36,7 +36,7 @@ from opentelemetry.trace import Status, StatusCode from opentelemetry.util.types import Attributes -from genkit.core.trace.adjusting_exporter import AdjustingTraceExporter, RedactedSpan +from genkit._core._trace._adjusting_exporter import AdjustingTraceExporter, RedactedSpan class MockSpanExporter(SpanExporter): diff --git a/py/packages/genkit/tests/genkit/core/trace/default_exporter_test.py b/py/packages/genkit/tests/genkit/core/trace/default_exporter_test.py index ae5b331eb9..3bc9dd3b7b 100644 --- a/py/packages/genkit/tests/genkit/core/trace/default_exporter_test.py +++ b/py/packages/genkit/tests/genkit/core/trace/default_exporter_test.py @@ -17,9 +17,8 @@ """Tests for the default telemetry exporter module. This module tests: - - TelemetryServerSpanExporter: Exports spans to a telemetry server + - TraceServerExporter: Exports spans to a telemetry server - extract_span_data: Extracts span data for export - - is_realtime_telemetry_enabled: Checks if realtime telemetry is enabled - create_span_processor: Creates appropriate span processor based on environment - init_telemetry_server_exporter: Initializes the telemetry server exporter """ @@ -30,100 +29,36 @@ from opentelemetry import trace as trace_api from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExportResult -from genkit.core.environment import EnvVar, GenkitEnvironment -from genkit.core.trace.default_exporter import ( - TelemetryServerSpanExporter, +from genkit._core._environment import GENKIT_ENV, GenkitEnvironment +from genkit._core._trace._default_exporter import ( + TraceServerExporter, create_span_processor, extract_span_data, init_telemetry_server_exporter, - is_realtime_telemetry_enabled, ) -from genkit.core.trace.realtime_processor import RealtimeSpanProcessor - -# ============================================================================= -# Tests for is_realtime_telemetry_enabled -# ============================================================================= - - -def test_is_realtime_telemetry_enabled_when_not_set() -> None: - """Test that realtime telemetry is disabled when env var is not set.""" - with mock.patch.dict(os.environ, clear=True): - assert is_realtime_telemetry_enabled() is False - - -def test_is_realtime_telemetry_enabled_when_true() -> None: - """Test that realtime telemetry is enabled when env var is 'true'.""" - with mock.patch.dict(os.environ, {'GENKIT_ENABLE_REALTIME_TELEMETRY': 'true'}): - assert is_realtime_telemetry_enabled() is True - - -def test_is_realtime_telemetry_enabled_when_true_uppercase() -> None: - """Test that realtime telemetry is enabled regardless of case.""" - with mock.patch.dict(os.environ, {'GENKIT_ENABLE_REALTIME_TELEMETRY': 'TRUE'}): - assert is_realtime_telemetry_enabled() is True - - -def test_is_realtime_telemetry_enabled_when_true_mixed_case() -> None: - """Test that realtime telemetry is enabled with mixed case.""" - with mock.patch.dict(os.environ, {'GENKIT_ENABLE_REALTIME_TELEMETRY': 'True'}): - assert is_realtime_telemetry_enabled() is True - - -def test_is_realtime_telemetry_enabled_when_false() -> None: - """Test that realtime telemetry is disabled when env var is 'false'.""" - with mock.patch.dict(os.environ, {'GENKIT_ENABLE_REALTIME_TELEMETRY': 'false'}): - assert is_realtime_telemetry_enabled() is False - - -def test_is_realtime_telemetry_enabled_when_invalid() -> None: - """Test that realtime telemetry is disabled with invalid value.""" - with mock.patch.dict(os.environ, {'GENKIT_ENABLE_REALTIME_TELEMETRY': 'invalid'}): - assert is_realtime_telemetry_enabled() is False - - -def test_is_realtime_telemetry_enabled_when_empty() -> None: - """Test that realtime telemetry is disabled when env var is empty.""" - with mock.patch.dict(os.environ, {'GENKIT_ENABLE_REALTIME_TELEMETRY': ''}): - assert is_realtime_telemetry_enabled() is False - +from genkit._core._trace._realtime_processor import RealtimeSpanProcessor # ============================================================================= # Tests for create_span_processor # ============================================================================= -def test_create_span_processor_returns_realtime_in_dev_with_env_var() -> None: - """Test that RealtimeSpanProcessor is returned in dev mode with env var set.""" +def test_create_span_processor_returns_realtime_in_dev() -> None: + """Test that RealtimeSpanProcessor is returned in dev mode.""" mock_exporter = MagicMock() with mock.patch.dict( os.environ, { - EnvVar.GENKIT_ENV: GenkitEnvironment.DEV, - 'GENKIT_ENABLE_REALTIME_TELEMETRY': 'true', + GENKIT_ENV: GenkitEnvironment.DEV, }, ): processor = create_span_processor(mock_exporter) assert isinstance(processor, RealtimeSpanProcessor) -def test_create_span_processor_returns_simple_in_dev_without_env_var() -> None: - """Test that SimpleSpanProcessor is returned in dev mode without env var.""" - mock_exporter = MagicMock() - - with mock.patch.dict( - os.environ, - { - EnvVar.GENKIT_ENV: GenkitEnvironment.DEV, - }, - clear=True, - ): - processor = create_span_processor(mock_exporter) - assert isinstance(processor, SimpleSpanProcessor) - - def test_create_span_processor_returns_batch_in_prod() -> None: """Test that BatchSpanProcessor is returned in production mode.""" mock_exporter = MagicMock() @@ -131,7 +66,7 @@ def test_create_span_processor_returns_batch_in_prod() -> None: with mock.patch.dict( os.environ, { - EnvVar.GENKIT_ENV: GenkitEnvironment.PROD, + GENKIT_ENV: GenkitEnvironment.PROD, }, ): processor = create_span_processor(mock_exporter) @@ -147,22 +82,6 @@ def test_create_span_processor_returns_batch_when_no_env_set() -> None: assert isinstance(processor, BatchSpanProcessor) -def test_create_span_processor_ignores_realtime_env_in_prod() -> None: - """Test that realtime env var is ignored in production mode (matches JS behavior).""" - mock_exporter = MagicMock() - - with mock.patch.dict( - os.environ, - { - EnvVar.GENKIT_ENV: GenkitEnvironment.PROD, - 'GENKIT_ENABLE_REALTIME_TELEMETRY': 'true', - }, - ): - processor = create_span_processor(mock_exporter) - # Should still return BatchSpanProcessor in prod, not RealtimeSpanProcessor - assert isinstance(processor, BatchSpanProcessor) - - # ============================================================================= # Tests for init_telemetry_server_exporter # ============================================================================= @@ -173,7 +92,7 @@ def test_init_telemetry_server_exporter_returns_exporter_when_url_set() -> None: with mock.patch.dict(os.environ, {'GENKIT_TELEMETRY_SERVER': 'http://localhost:4000'}): exporter = init_telemetry_server_exporter() assert exporter is not None - assert isinstance(exporter, TelemetryServerSpanExporter) + assert isinstance(exporter, TraceServerExporter) assert exporter.telemetry_server_url == 'http://localhost:4000' @@ -185,21 +104,21 @@ def test_init_telemetry_server_exporter_returns_none_when_url_not_set() -> None: # ============================================================================= -# Tests for TelemetryServerSpanExporter +# Tests for TraceServerExporter # ============================================================================= def test_telemetry_server_exporter_init_default_endpoint() -> None: - """Test TelemetryServerSpanExporter initialization with default endpoint.""" - exporter = TelemetryServerSpanExporter(telemetry_server_url='http://localhost:4000') + """Test TraceServerExporter initialization with default endpoint.""" + exporter = TraceServerExporter(telemetry_server_url='http://localhost:4000') assert exporter.telemetry_server_url == 'http://localhost:4000' assert exporter.telemetry_server_endpoint == '/api/traces' def test_telemetry_server_exporter_init_custom_endpoint() -> None: - """Test TelemetryServerSpanExporter initialization with custom endpoint.""" - exporter = TelemetryServerSpanExporter( + """Test TraceServerExporter initialization with custom endpoint.""" + exporter = TraceServerExporter( telemetry_server_url='http://localhost:4000', telemetry_server_endpoint='/custom/traces', ) @@ -210,7 +129,7 @@ def test_telemetry_server_exporter_init_custom_endpoint() -> None: def test_telemetry_server_exporter_force_flush_returns_true() -> None: """Test that force_flush always returns True (no buffering).""" - exporter = TelemetryServerSpanExporter(telemetry_server_url='http://localhost:4000') + exporter = TraceServerExporter(telemetry_server_url='http://localhost:4000') result = exporter.force_flush() assert result is True @@ -218,13 +137,13 @@ def test_telemetry_server_exporter_force_flush_returns_true() -> None: def test_telemetry_server_exporter_force_flush_ignores_timeout() -> None: """Test that force_flush ignores the timeout parameter.""" - exporter = TelemetryServerSpanExporter(telemetry_server_url='http://localhost:4000') + exporter = TraceServerExporter(telemetry_server_url='http://localhost:4000') result = exporter.force_flush(timeout_millis=1) assert result is True -@patch('genkit.core.trace.default_exporter.httpx.Client') +@patch('genkit._core._trace._default_exporter.httpx.Client') def test_telemetry_server_exporter_export_sends_http_post(mock_client_class: MagicMock) -> None: """Test that export sends HTTP POST requests for each span.""" # Setup mock client @@ -232,7 +151,7 @@ def test_telemetry_server_exporter_export_sends_http_post(mock_client_class: Mag mock_client_class.return_value.__enter__ = MagicMock(return_value=mock_client) mock_client_class.return_value.__exit__ = MagicMock(return_value=None) - exporter = TelemetryServerSpanExporter(telemetry_server_url='http://localhost:4000') + exporter = TraceServerExporter(telemetry_server_url='http://localhost:4000') # Create a mock span mock_span = create_mock_span() @@ -245,7 +164,7 @@ def test_telemetry_server_exporter_export_sends_http_post(mock_client_class: Mag mock_client.post.assert_called_once() -@patch('genkit.core.trace.default_exporter.httpx.Client') +@patch('genkit._core._trace._default_exporter.httpx.Client') def test_telemetry_server_exporter_export_multiple_spans(mock_client_class: MagicMock) -> None: """Test that export sends HTTP POST for each span in the sequence.""" # Setup mock client @@ -253,7 +172,7 @@ def test_telemetry_server_exporter_export_multiple_spans(mock_client_class: Magi mock_client_class.return_value.__enter__ = MagicMock(return_value=mock_client) mock_client_class.return_value.__exit__ = MagicMock(return_value=None) - exporter = TelemetryServerSpanExporter(telemetry_server_url='http://localhost:4000') + exporter = TraceServerExporter(telemetry_server_url='http://localhost:4000') # Create multiple mock spans mock_spans = [create_mock_span() for _ in range(3)] diff --git a/py/packages/genkit/tests/genkit/core/trace/realtime_processor_test.py b/py/packages/genkit/tests/genkit/core/trace/realtime_processor_test.py index a35e809635..ca179ae936 100644 --- a/py/packages/genkit/tests/genkit/core/trace/realtime_processor_test.py +++ b/py/packages/genkit/tests/genkit/core/trace/realtime_processor_test.py @@ -27,7 +27,7 @@ from opentelemetry.sdk.trace import ReadableSpan, Span from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult -from genkit.core.trace.realtime_processor import RealtimeSpanProcessor +from genkit._core._trace._realtime_processor import RealtimeSpanProcessor class MockSpanExporter(SpanExporter): @@ -71,9 +71,9 @@ def test_realtime_processor_exports_on_start() -> None: # Call on_start processor.on_start(mock_span) - # Verify span was exported + # Verify span was exported (on_start passes list) assert len(exporter.exported_spans) == 1 - assert exporter.exported_spans[0] == [mock_span] + assert list(exporter.exported_spans[0]) == [mock_span] def test_realtime_processor_exports_on_end() -> None: @@ -90,9 +90,9 @@ def test_realtime_processor_exports_on_end() -> None: # Call on_end processor.on_end(mock_span) - # Verify span was exported + # Verify span was exported (SimpleSpanProcessor passes tuple) assert len(exporter.exported_spans) == 1 - assert exporter.exported_spans[0] == [mock_span] + assert list(exporter.exported_spans[0]) == [mock_span] def test_realtime_processor_exports_twice_for_full_lifecycle() -> None: @@ -112,36 +112,19 @@ def test_realtime_processor_exports_twice_for_full_lifecycle() -> None: processor.on_start(mock_span_start) processor.on_end(mock_span_end) - # Verify span was exported twice + # Verify span was exported twice (on_start uses list, on_end uses tuple) assert len(exporter.exported_spans) == 2 - assert exporter.exported_spans[0] == [mock_span_start] - assert exporter.exported_spans[1] == [mock_span_end] + assert list(exporter.exported_spans[0]) == [mock_span_start] + assert list(exporter.exported_spans[1]) == [mock_span_end] -def test_realtime_processor_force_flush_delegates_to_exporter() -> None: - """Test that force_flush is delegated to the underlying exporter.""" +def test_realtime_processor_force_flush() -> None: + """Test that force_flush works (inherited from SimpleSpanProcessor).""" exporter = MockSpanExporter() processor = RealtimeSpanProcessor(exporter) - # Call force_flush with custom timeout result = processor.force_flush(timeout_millis=5000) - # Verify delegation - assert result is True - assert exporter.force_flush_called is True - assert exporter.force_flush_timeout == 5000 - - -def test_realtime_processor_force_flush_returns_true_if_exporter_lacks_method() -> None: - """Test that force_flush returns True if exporter doesn't have the method.""" - # Create a minimal exporter without force_flush - mock_exporter = MagicMock(spec=['export', 'shutdown']) - - processor = RealtimeSpanProcessor(mock_exporter) - - # Call force_flush - should return True even without exporter support - result = processor.force_flush() - assert result is True @@ -170,7 +153,7 @@ def test_realtime_processor_on_start_with_parent_context() -> None: # Verify span was still exported assert len(exporter.exported_spans) == 1 - assert exporter.exported_spans[0] == [mock_span] + assert list(exporter.exported_spans[0]) == [mock_span] def test_realtime_processor_multiple_spans() -> None: diff --git a/py/packages/genkit/tests/genkit/lang/deprecations_test.py b/py/packages/genkit/tests/genkit/lang/deprecations_test.py deleted file mode 100644 index 6b88bc3989..0000000000 --- a/py/packages/genkit/tests/genkit/lang/deprecations_test.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for deprecation helpers.""" - -import sys -import unittest -import warnings - -if sys.version_info < (3, 11): - from strenum import StrEnum -else: - from enum import StrEnum - -from genkit.lang.deprecations import ( - DeprecationInfo, - DeprecationStatus, - deprecated_enum_metafactory, -) - -TEST_DEPRECATED_MODELS = { - 'GEMINI_1_0_PRO': DeprecationInfo(recommendation='GEMINI_2_0_PRO', status=DeprecationStatus.LEGACY), - 'GEMINI_1_5_PRO': DeprecationInfo(recommendation='GEMINI_2_0_PRO', status=DeprecationStatus.DEPRECATED), - 'GEMINI_1_5_FLASH': DeprecationInfo(recommendation='GEMINI_2_0_FLASH', status=DeprecationStatus.DEPRECATED), - 'GEMINI_1_5_FLASH_8B': DeprecationInfo(recommendation=None, status=DeprecationStatus.DEPRECATED), -} - - -class GeminiVersionTest(StrEnum, metaclass=deprecated_enum_metafactory(TEST_DEPRECATED_MODELS)): - """Test Gemini models enum.""" - - GEMINI_1_0_PRO = 'gemini-1.0-pro' - GEMINI_1_5_PRO = 'gemini-1.5-pro' - GEMINI_1_5_FLASH = 'gemini-1.5-flash' - GEMINI_1_5_FLASH_8B = 'gemini-1.5-flash-8b' - GEMINI_2_0_FLASH = 'gemini-2.0-flash' - GEMINI_2_0_PRO = 'gemini-2.0-pro' # Added for recommendation - - -class TestDeprecatedEnum(unittest.TestCase): - """Test deprecated enum members.""" - - def test_legacy_member_warning(self) -> None: - """Verify warning for legacy member with recommendation.""" - expected_regex = ( - r'GeminiVersionTest\.GEMINI_1_0_PRO is legacy; ' - r'use GeminiVersionTest\.GEMINI_2_0_PRO instead' - ) - warnings.simplefilter('always', DeprecationWarning) - try: - with self.assertWarnsRegex(DeprecationWarning, expected_regex): - member = GeminiVersionTest.GEMINI_1_0_PRO - self.assertEqual(member, GeminiVersionTest.GEMINI_1_0_PRO) - self.assertEqual(member.value, 'gemini-1.0-pro') - finally: - warnings.simplefilter('default', DeprecationWarning) - - def test_deprecated_member_with_recommendation_warning(self) -> None: - """Verify warning for deprecated member with recommendation.""" - expected_regex = ( - r'GeminiVersionTest\.GEMINI_1_5_PRO is deprecated; ' - r'use GeminiVersionTest\.GEMINI_2_0_PRO instead' - ) - warnings.simplefilter('always', DeprecationWarning) - try: - with self.assertWarnsRegex(DeprecationWarning, expected_regex): - member = GeminiVersionTest.GEMINI_1_5_PRO - self.assertEqual(member, GeminiVersionTest.GEMINI_1_5_PRO) - self.assertEqual(member.value, 'gemini-1.5-pro') - finally: - warnings.simplefilter('default', DeprecationWarning) - - def test_deprecated_member_without_recommendation_warning(self) -> None: - """Verify warning for deprecated member without recommendation.""" - expected_regex = r'GeminiVersionTest\.GEMINI_1_5_FLASH_8B is deprecated' - warnings.simplefilter('always', DeprecationWarning) - try: - with self.assertWarnsRegex(DeprecationWarning, expected_regex) as cm: - member = GeminiVersionTest.GEMINI_1_5_FLASH_8B - self.assertTrue( - hasattr(cm, 'warnings') and cm.warnings, - 'Warning was not captured by assertWarnsRegex', - ) - self.assertNotIn('instead', str(cm.warnings[0].message)) - self.assertEqual(member, GeminiVersionTest.GEMINI_1_5_FLASH_8B) - self.assertEqual(member.value, 'gemini-1.5-flash-8b') - finally: - warnings.simplefilter('default', DeprecationWarning) - - def test_supported_member_no_warning(self) -> None: - """Verify no warning for a supported member.""" - member_to_test = GeminiVersionTest.GEMINI_2_0_FLASH - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter('always') - member = getattr(GeminiVersionTest, member_to_test.name) - self.assertEqual( - len(w), - 0, - f'Expected no warnings for {member_to_test.name}, but got {len(w)}: {[warn.message for warn in w]}', - ) - self.assertEqual(member, member_to_test) - self.assertEqual(member.value, member_to_test.value) - - def test_recommended_member_no_warning(self) -> None: - """Verify no warning for a member used as a recommendation.""" - member_to_test = GeminiVersionTest.GEMINI_2_0_PRO - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter('always') - member = getattr(GeminiVersionTest, member_to_test.name) - self.assertEqual( - len(w), - 0, - f'Expected no warnings for {member_to_test.name}, but got {len(w)}: {[warn.message for warn in w]}', - ) - self.assertEqual(member, member_to_test) - self.assertEqual(member.value, member_to_test.value) - - def test_access_via_value_no_warning(self) -> None: - """Verify no warning when accessing members via value lookup.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter('always') - member_dep = GeminiVersionTest('gemini-1.5-pro') - member_leg = GeminiVersionTest('gemini-1.0-pro') - member_sup = GeminiVersionTest('gemini-2.0-flash') - self.assertEqual( - len(w), - 0, - f'Expected no warnings for value lookup, but got {len(w)}: {[warn.message for warn in w]}', - ) - self.assertEqual(member_dep, GeminiVersionTest.GEMINI_1_5_PRO) - self.assertEqual(member_leg, GeminiVersionTest.GEMINI_1_0_PRO) - self.assertEqual(member_sup, GeminiVersionTest.GEMINI_2_0_FLASH) - - def test_iteration_no_warning(self) -> None: - """Verify no warning when iterating over enum members.""" - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter('always') - members = list(GeminiVersionTest) - self.assertEqual( - len(w), - 0, - f'Expected no warnings during iteration, but got {len(w)}: {[warn.message for warn in w]}', - ) - self.assertEqual(len(members), len(GeminiVersionTest.__members__)) - self.assertIn(GeminiVersionTest.GEMINI_1_5_FLASH_8B, members) - self.assertIn(GeminiVersionTest.GEMINI_2_0_FLASH, members) - - -if __name__ == '__main__': - unittest.main() diff --git a/py/packages/genkit/tests/genkit/testing_test.py b/py/packages/genkit/tests/genkit/testing_test.py index e4b5d18e9f..777ef9086f 100644 --- a/py/packages/genkit/tests/genkit/testing_test.py +++ b/py/packages/genkit/tests/genkit/testing_test.py @@ -57,17 +57,8 @@ import pytest -from genkit.ai import Genkit -from genkit.core.typing import ( - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - Message, - Part, - Role, - TextPart, -) -from genkit.testing import ( +from genkit import ActionRunContext, Genkit, Message, ModelConfig, ModelRequest, ModelResponse, ModelResponseChunk +from genkit._ai._testing import ( EchoModel, GablorkenInput, ProgrammableModel, @@ -79,17 +70,24 @@ skip, test_models as run_model_tests, ) +from genkit._core._typing import ( + Part, + Role, + TextPart, +) -class MockActionRunContext: +class MockActionRunContext(ActionRunContext): """Mock context for testing model functions directly.""" def __init__(self) -> None: """Initialize with empty chunks list.""" - self.chunks: list[GenerateResponseChunk] = [] + super().__init__() + self.chunks: list[ModelResponseChunk] = [] - def send_chunk(self, chunk: GenerateResponseChunk) -> None: + def send_chunk(self, chunk: object) -> None: """Append a chunk to the chunks list.""" + assert isinstance(chunk, ModelResponseChunk) self.chunks.append(chunk) @@ -102,12 +100,13 @@ def ai() -> Genkit: class TestEchoModel: """Tests for EchoModel functionality.""" - def test_echo_model_basic(self) -> None: + @pytest.mark.asyncio + async def test_echo_model_basic(self) -> None: """Test basic echo functionality.""" echo = EchoModel() ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -117,7 +116,7 @@ def test_echo_model_basic(self) -> None: ) # pyright: ignore[reportArgumentType] - MockActionRunContext is compatible - response = echo.model_fn(request, ctx) # type: ignore[arg-type] + response = await echo.model_fn(request, ctx) assert response.message is not None text = response.message.content[0].root.text @@ -126,34 +125,36 @@ def test_echo_model_basic(self) -> None: assert 'user:' in text assert 'Hello world' in text - def test_echo_model_with_config(self) -> None: + @pytest.mark.asyncio + async def test_echo_model_with_config(self) -> None: """Test that echo includes config in response.""" echo = EchoModel() ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, content=[Part(root=TextPart(text='test'))], ), ], - config={'temperature': 0.5}, + config=ModelConfig(temperature=0.5), ) - response = echo.model_fn(request, ctx) # type: ignore[arg-type] + response = await echo.model_fn(request, ctx) assert response.message is not None text = response.message.content[0].root.text assert isinstance(text, str) assert 'temperature' in text - def test_echo_model_stream_countdown(self) -> None: + @pytest.mark.asyncio + async def test_echo_model_stream_countdown(self) -> None: """Test stream countdown functionality.""" echo = EchoModel(stream_countdown=True) ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -162,7 +163,8 @@ def test_echo_model_stream_countdown(self) -> None: ], ) - echo.model_fn(request, ctx) # type: ignore[arg-type] + # Model uses ctx.send_chunk() internally for streaming + await echo.model_fn(request, ctx) # Should have streamed 3, 2, 1 assert len(ctx.chunks) == 3 @@ -170,12 +172,13 @@ def test_echo_model_stream_countdown(self) -> None: assert ctx.chunks[1].content[0].root.text == '2' assert ctx.chunks[2].content[0].root.text == '1' - def test_echo_model_stores_request(self) -> None: + @pytest.mark.asyncio + async def test_echo_model_stores_request(self) -> None: """Test that echo stores the last request.""" echo = EchoModel() ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -184,7 +187,7 @@ def test_echo_model_stores_request(self) -> None: ], ) - echo.model_fn(request, ctx) # type: ignore[arg-type] + await echo.model_fn(request, ctx) assert echo.last_request is not None assert echo.last_request.messages[0].content[0].root.text == 'test' @@ -203,11 +206,12 @@ async def test_define_echo_model(self, ai: Genkit) -> None: class TestProgrammableModel: """Tests for ProgrammableModel functionality.""" - def test_programmable_model_basic(self) -> None: + @pytest.mark.asyncio + async def test_programmable_model_basic(self) -> None: """Test basic programmable model functionality.""" pm = ProgrammableModel() pm.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Response 1'))], @@ -216,7 +220,7 @@ def test_programmable_model_basic(self) -> None: ] ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -225,23 +229,24 @@ def test_programmable_model_basic(self) -> None: ], ) - response = pm.model_fn(request, ctx) # type: ignore[arg-type] + response = await pm.model_fn(request, ctx) assert response.message is not None assert response.message.content[0].root.text == 'Response 1' assert pm.request_count == 1 - def test_programmable_model_multiple_responses(self) -> None: + @pytest.mark.asyncio + async def test_programmable_model_multiple_responses(self) -> None: """Test multiple sequential responses.""" pm = ProgrammableModel() pm.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Response 1'))], ), ), - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Response 2'))], @@ -250,7 +255,7 @@ def test_programmable_model_multiple_responses(self) -> None: ] ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -259,8 +264,8 @@ def test_programmable_model_multiple_responses(self) -> None: ], ) - response1 = pm.model_fn(request, ctx) # type: ignore[arg-type] - response2 = pm.model_fn(request, ctx) # type: ignore[arg-type] + response1 = await pm.model_fn(request, ctx) + response2 = await pm.model_fn(request, ctx) assert response1.message is not None assert response2.message is not None @@ -268,11 +273,12 @@ def test_programmable_model_multiple_responses(self) -> None: assert response2.message.content[0].root.text == 'Response 2' assert pm.request_count == 2 - def test_programmable_model_chunks(self) -> None: + @pytest.mark.asyncio + async def test_programmable_model_chunks(self) -> None: """Test streaming programmed chunks.""" pm = ProgrammableModel() pm.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Final'))], @@ -281,13 +287,13 @@ def test_programmable_model_chunks(self) -> None: ] pm.chunks = [ [ - GenerateResponseChunk(content=[Part(root=TextPart(text='Chunk 1'))]), - GenerateResponseChunk(content=[Part(root=TextPart(text='Chunk 2'))]), + ModelResponseChunk(content=[Part(root=TextPart(text='Chunk 1'))]), + ModelResponseChunk(content=[Part(root=TextPart(text='Chunk 2'))]), ], ] ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -296,17 +302,19 @@ def test_programmable_model_chunks(self) -> None: ], ) - pm.model_fn(request, ctx) # type: ignore[arg-type] + # Model uses ctx.send_chunk() internally for streaming + await pm.model_fn(request, ctx) assert len(ctx.chunks) == 2 assert ctx.chunks[0].content[0].root.text == 'Chunk 1' assert ctx.chunks[1].content[0].root.text == 'Chunk 2' - def test_programmable_model_reset(self) -> None: + @pytest.mark.asyncio + async def test_programmable_model_reset(self) -> None: """Test reset clears state.""" pm = ProgrammableModel() pm.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Response'))], @@ -315,7 +323,7 @@ def test_programmable_model_reset(self) -> None: ] ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -324,7 +332,7 @@ def test_programmable_model_reset(self) -> None: ], ) - pm.model_fn(request, ctx) # type: ignore[arg-type] + await pm.model_fn(request, ctx) assert pm.request_count == 1 assert pm.last_request is not None @@ -335,11 +343,12 @@ def test_programmable_model_reset(self) -> None: assert pm.responses == [] assert pm.chunks is None - def test_programmable_model_stores_deep_copy(self) -> None: + @pytest.mark.asyncio + async def test_programmable_model_stores_deep_copy(self) -> None: """Test that last_request is a deep copy.""" pm = ProgrammableModel() pm.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Response'))], @@ -348,7 +357,7 @@ def test_programmable_model_stores_deep_copy(self) -> None: ] ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -357,7 +366,7 @@ def test_programmable_model_stores_deep_copy(self) -> None: ], ) - pm.model_fn(request, ctx) # type: ignore[arg-type] + await pm.model_fn(request, ctx) # Modify original request original_part = request.messages[0].content[0].root @@ -375,7 +384,7 @@ async def test_define_programmable_model(self, ai: Genkit) -> None: """Test define_programmable_model helper function.""" pm, _action = define_programmable_model(ai, name='testPM') pm.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Programmed response'))], @@ -392,7 +401,8 @@ async def test_define_programmable_model(self, ai: Genkit) -> None: class TestStaticResponseModel: """Tests for StaticResponseModel functionality.""" - def test_static_model_basic(self) -> None: + @pytest.mark.asyncio + async def test_static_model_basic(self) -> None: """Test basic static response model functionality.""" static = StaticResponseModel( message={ @@ -402,7 +412,7 @@ def test_static_model_basic(self) -> None: ) ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -411,12 +421,13 @@ def test_static_model_basic(self) -> None: ], ) - response = static.model_fn(request, ctx) # type: ignore[arg-type] + response = await static.model_fn(request, ctx) assert response.message is not None assert response.message.content[0].root.text == 'Static response' - def test_static_model_request_count(self) -> None: + @pytest.mark.asyncio + async def test_static_model_request_count(self) -> None: """Test request counting.""" static = StaticResponseModel( message={ @@ -426,7 +437,7 @@ def test_static_model_request_count(self) -> None: ) ctx = MockActionRunContext() - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -435,9 +446,9 @@ def test_static_model_request_count(self) -> None: ], ) - static.model_fn(request, ctx) # type: ignore[arg-type] - static.model_fn(request, ctx) # type: ignore[arg-type] - static.model_fn(request, ctx) # type: ignore[arg-type] + await static.model_fn(request, ctx) + await static.model_fn(request, ctx) + await static.model_fn(request, ctx) assert static.request_count == 3 @@ -500,48 +511,48 @@ async def test_test_models_with_echo_model(self, ai: Genkit) -> None: pm, _ = define_programmable_model(ai, name='testModel') pm.responses = [ # For basic hi test - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Hi'))], ), ), # For multimodal test (will skip since no media support) - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='plus'))], ), ), # For history test - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Nice to meet you'))], ), ), - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Your name is Glorb'))], ), ), # For system prompt test - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Bye'))], ), ), # For structured output test - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='{"name": "Jack", "occupation": "Lumberjack"}'))], ), ), # For tool calling test (will skip since no tools support) - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='9.407'))], @@ -569,7 +580,7 @@ async def test_test_models_report_format(self, ai: Genkit) -> None: """Test that report format matches JS implementation.""" pm, _ = define_programmable_model(ai, name='formatTestModel') pm.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Hi'))], @@ -597,7 +608,7 @@ async def test_test_models_multiple_models(self, ai: Genkit) -> None: """Test test_models with multiple models.""" pm1, _ = define_programmable_model(ai, name='model1') pm1.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Hi'))], @@ -607,7 +618,7 @@ async def test_test_models_multiple_models(self, ai: Genkit) -> None: pm2, _ = define_programmable_model(ai, name='model2') pm2.responses = [ - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Hello'))], @@ -630,7 +641,7 @@ async def test_test_models_handles_failures(self, ai: Genkit) -> None: pm, _ = define_programmable_model(ai, name='failingModel') pm.responses = [ # Return something that doesn't match expected pattern - GenerateResponse( + ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text='Goodbye'))], # Should be "Hi" diff --git a/py/packages/genkit/tests/genkit/veneer/reflection_server_test.py b/py/packages/genkit/tests/genkit/veneer/reflection_server_test.py index f5b8b1a12a..bda245752a 100644 --- a/py/packages/genkit/tests/genkit/veneer/reflection_server_test.py +++ b/py/packages/genkit/tests/genkit/veneer/reflection_server_test.py @@ -19,9 +19,9 @@ import httpx -from genkit.ai._aio import Genkit -from genkit.ai._server import ServerSpec -from genkit.core.environment import EnvVar, GenkitEnvironment +from genkit import Genkit +from genkit._core._environment import GENKIT_ENV, GenkitEnvironment +from genkit._core._reflection import ServerSpec def _find_free_port() -> int: @@ -43,7 +43,7 @@ def test_server_starts_on_construction() -> None: No run_main(), no lifespan hooks β€” construction is sufficient. """ port = _find_free_port() - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.DEV}): ai = Genkit(reflection_server_spec=ServerSpec(scheme='http', host='127.0.0.1', port=port)) resp = _wait_and_get(ai, '/api/__health') assert resp.status_code == 200 @@ -57,7 +57,7 @@ def test_flow_registered_after_construction_is_visible() -> None: See test_registry_reads_concurrent_with_writes for that. """ port = _find_free_port() - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.DEV}): ai = Genkit(reflection_server_spec=ServerSpec(scheme='http', host='127.0.0.1', port=port)) @ai.flow() @@ -81,7 +81,7 @@ def test_registry_reads_concurrent_with_writes() -> None: port = _find_free_port() errors: list[Exception] = [] - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.DEV}): ai = Genkit(reflection_server_spec=ServerSpec(scheme='http', host='127.0.0.1', port=port)) assert ai._reflection_ready.wait(timeout=5) # pyright: ignore[reportPrivateUsage] @@ -118,7 +118,7 @@ async def _f(x: str) -> str: def test_two_instances_serve_concurrently() -> None: """Two Genkit() instances in the same process don't interfere with each other.""" port1, port2 = _find_free_port(), _find_free_port() - with mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}): + with mock.patch.dict(os.environ, {GENKIT_ENV: GenkitEnvironment.DEV}): ai1 = Genkit(reflection_server_spec=ServerSpec(scheme='http', host='127.0.0.1', port=port1)) ai2 = Genkit(reflection_server_spec=ServerSpec(scheme='http', host='127.0.0.1', port=port2)) diff --git a/py/packages/genkit/tests/genkit/veneer/server_test.py b/py/packages/genkit/tests/genkit/veneer/server_test.py index d83f7f99ce..022d96f478 100644 --- a/py/packages/genkit/tests/genkit/veneer/server_test.py +++ b/py/packages/genkit/tests/genkit/veneer/server_test.py @@ -9,8 +9,8 @@ import pathlib import tempfile -from genkit.ai._runtime import RuntimeManager -from genkit.ai._server import ServerSpec +from genkit._ai._runtime import RuntimeManager +from genkit._core._reflection import ServerSpec def test_server_spec() -> None: diff --git a/py/packages/genkit/tests/genkit/veneer/veneer_resource_test.py b/py/packages/genkit/tests/genkit/veneer/veneer_resource_test.py index 8e2b6dd9b2..6014093e0b 100644 --- a/py/packages/genkit/tests/genkit/veneer/veneer_resource_test.py +++ b/py/packages/genkit/tests/genkit/veneer/veneer_resource_test.py @@ -24,10 +24,10 @@ import pytest -from genkit.ai import ActionRunContext, Genkit -from genkit.blocks.resource import ResourceInput -from genkit.core.action.types import ActionKind -from genkit.core.typing import Part, TextPart +from genkit import ActionRunContext, Genkit +from genkit._ai._resource import ResourceInput +from genkit._core._action import ActionKind +from genkit._core._typing import Part, TextPart @pytest.mark.asyncio @@ -38,7 +38,7 @@ async def test_define_resource_veneer() -> None: async def my_resource_fn(input: ResourceInput, ctx: ActionRunContext) -> dict[str, list[Part]]: return {'content': [Part(root=TextPart(text=f'Content for {input.uri}'))]} - act = ai.define_resource({'uri': 'http://example.com/foo'}, my_resource_fn) + act = ai.define_resource(fn=my_resource_fn, uri='http://example.com/foo') assert act.name == 'http://example.com/foo' assert act.metadata is not None @@ -51,5 +51,5 @@ async def my_resource_fn(input: ResourceInput, ctx: ActionRunContext) -> dict[st assert looked_up == act # Verify execution - output = await act.arun({'uri': 'http://example.com/foo'}) + output = await act.run({'uri': 'http://example.com/foo'}) assert 'Content for http://example.com/foo' in output.response['content'][0]['text'] diff --git a/py/packages/genkit/tests/genkit/veneer/veneer_test.py b/py/packages/genkit/tests/genkit/veneer/veneer_test.py index ef06e03502..ed8a7206e2 100644 --- a/py/packages/genkit/tests/genkit/veneer/veneer_test.py +++ b/py/packages/genkit/tests/genkit/veneer/veneer_test.py @@ -6,35 +6,42 @@ """Tests for the action module.""" import json -from typing import Any, cast +from collections.abc import Awaitable, Callable +from typing import Any import pytest from pydantic import BaseModel, Field -from genkit.ai import Genkit, Output, ToolRunContext, tool_response -from genkit.blocks.document import Document -from genkit.blocks.formats.types import FormatDef, Formatter, FormatterConfig -from genkit.blocks.model import MessageWrapper, ModelMiddlewareNext, text_from_message -from genkit.core.action import ActionRunContext -from genkit.core.action.types import ActionKind -from genkit.core.typing import ( +from genkit import ( + Document, + Genkit, + Message, + ModelResponse, + ModelResponseChunk, + ToolRunContext, + tool_response, +) +from genkit._ai._formats._types import FormatDef, Formatter, FormatterConfig +from genkit._ai._model import text_from_message +from genkit._ai._testing import ( + EchoModel, + ProgrammableModel, + define_echo_model, + define_programmable_model, +) +from genkit._core._action import ActionKind, ActionRunContext +from genkit._core._model import ModelRequest +from genkit._core._typing import ( BaseDataPoint, Details, - DocumentData, DocumentPart, EvalFnResponse, EvalRequest, EvalResponse, FinishReason, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - Message, Metadata, ModelInfo, - OutputConfig, Part, - RetrieverResponse, Role, Score, Supports, @@ -46,12 +53,6 @@ ToolResponse, ToolResponsePart, ) -from genkit.testing import ( - EchoModel, - ProgrammableModel, - define_echo_model, - define_programmable_model, -) # type SetupFixture = tuple[Genkit, EchoModel, ProgrammableModel] SetupFixture = tuple[Genkit, EchoModel, ProgrammableModel] @@ -73,15 +74,15 @@ async def test_generate_uses_default_model(setup_test: SetupFixture) -> None: """Test that the generate function uses the default model.""" ai, *_ = setup_test - want_txt = '[ECHO] user: "hi" {"temperature":11}' + want_txt = '[ECHO] user: "hi" {"temperature":11.0}' response = await ai.generate(prompt='hi', config={'temperature': 11}) assert response.text == want_txt - _, response = ai.generate_stream(prompt='hi', config={'temperature': 11}) + stream_result = ai.generate_stream(prompt='hi', config={'temperature': 11}) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt @pytest.mark.asyncio @@ -123,11 +124,11 @@ async def test_generate_with_explicit_model(setup_test: SetupFixture) -> None: response = await ai.generate(model='echoModel', prompt='hi', config={'temperature': 11}) - assert response.text == '[ECHO] user: "hi" {"temperature":11}' + assert response.text == '[ECHO] user: "hi" {"temperature":11.0}' - _, response = ai.generate_stream(model='echoModel', prompt='hi', config={'temperature': 11}) + stream_result = ai.generate_stream(model='echoModel', prompt='hi', config={'temperature': 11}) - assert (await response).text == '[ECHO] user: "hi" {"temperature":11}' + assert (await stream_result.response).text == '[ECHO] user: "hi" {"temperature":11.0}' @pytest.mark.asyncio @@ -137,7 +138,7 @@ async def test_generate_with_str_prompt(setup_test: SetupFixture) -> None: response = await ai.generate(prompt='hi', config={'temperature': 11}) - assert response.text == '[ECHO] user: "hi" {"temperature":11}' + assert response.text == '[ECHO] user: "hi" {"temperature":11.0}' @pytest.mark.asyncio @@ -145,15 +146,15 @@ async def test_generate_with_part_prompt(setup_test: SetupFixture) -> None: """Test that the generate function with a part prompt works.""" ai, *_ = setup_test - want_txt = '[ECHO] user: "hi" {"temperature":11}' + want_txt = '[ECHO] user: "hi" {"temperature":11.0}' - response = await ai.generate(prompt=Part(root=TextPart(text='hi')), config={'temperature': 11}) + response = await ai.generate(prompt=[Part(root=TextPart(text='hi'))], config={'temperature': 11}) assert response.text == want_txt - _, response = ai.generate_stream(prompt=Part(root=TextPart(text='hi')), config={'temperature': 11}) + stream_result = ai.generate_stream(prompt=[Part(root=TextPart(text='hi'))], config={'temperature': 11}) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt @pytest.mark.asyncio @@ -161,7 +162,7 @@ async def test_generate_with_part_list_prompt(setup_test: SetupFixture) -> None: """Test that the generate function with a list of parts prompt works.""" ai, *_ = setup_test - want_txt = '[ECHO] user: "hello","world" {"temperature":11}' + want_txt = '[ECHO] user: "hello","world" {"temperature":11.0}' response = await ai.generate( prompt=[Part(root=TextPart(text='hello')), Part(root=TextPart(text='world'))], @@ -170,12 +171,12 @@ async def test_generate_with_part_list_prompt(setup_test: SetupFixture) -> None: assert response.text == want_txt - _, response = ai.generate_stream( + stream_result = ai.generate_stream( prompt=[Part(root=TextPart(text='hello')), Part(root=TextPart(text='world'))], config={'temperature': 11}, ) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt @pytest.mark.asyncio @@ -183,15 +184,15 @@ async def test_generate_with_str_system(setup_test: SetupFixture) -> None: """Test that the generate function with a string system works.""" ai, *_ = setup_test - want_txt = '[ECHO] system: "talk like pirate" user: "hi" {"temperature":11}' + want_txt = '[ECHO] system: "talk like pirate" user: "hi" {"temperature":11.0}' response = await ai.generate(system='talk like pirate', prompt='hi', config={'temperature': 11}) assert response.text == want_txt - _, response = ai.generate_stream(system='talk like pirate', prompt='hi', config={'temperature': 11}) + stream_result = ai.generate_stream(system='talk like pirate', prompt='hi', config={'temperature': 11}) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt @pytest.mark.asyncio @@ -199,23 +200,23 @@ async def test_generate_with_part_system(setup_test: SetupFixture) -> None: """Test that the generate function with a part system works.""" ai, *_ = setup_test - want_txt = '[ECHO] system: "talk like pirate" user: "hi" {"temperature":11}' + want_txt = '[ECHO] system: "talk like pirate" user: "hi" {"temperature":11.0}' response = await ai.generate( - system=Part(root=TextPart(text='talk like pirate')), + system=[Part(root=TextPart(text='talk like pirate'))], prompt='hi', config={'temperature': 11}, ) assert response.text == want_txt - _, response = ai.generate_stream( - system=Part(root=TextPart(text='talk like pirate')), + stream_result = ai.generate_stream( + system=[Part(root=TextPart(text='talk like pirate'))], prompt='hi', config={'temperature': 11}, ) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt @pytest.mark.asyncio @@ -223,7 +224,7 @@ async def test_generate_with_part_list_system(setup_test: SetupFixture) -> None: """Test that the generate function with a list of parts system works.""" ai, *_ = setup_test - want_txt = '[ECHO] system: "talk","like pirate" user: "hi" {"temperature":11}' + want_txt = '[ECHO] system: "talk","like pirate" user: "hi" {"temperature":11.0}' response = await ai.generate( system=[Part(root=TextPart(text='talk')), Part(root=TextPart(text='like pirate'))], @@ -233,13 +234,13 @@ async def test_generate_with_part_list_system(setup_test: SetupFixture) -> None: assert response.text == want_txt - _, response = ai.generate_stream( + stream_result = ai.generate_stream( system=[Part(root=TextPart(text='talk')), Part(root=TextPart(text='like pirate'))], prompt='hi', config={'temperature': 11}, ) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt @pytest.mark.asyncio @@ -257,9 +258,9 @@ async def test_generate_with_messages(setup_test: SetupFixture) -> None: config={'temperature': 11}, ) - assert response.text == '[ECHO] user: "hi" {"temperature":11}' + assert response.text == '[ECHO] user: "hi" {"temperature":11.0}' - _, response = ai.generate_stream( + stream_result = ai.generate_stream( messages=[ Message( role=Role.USER, @@ -269,7 +270,7 @@ async def test_generate_with_messages(setup_test: SetupFixture) -> None: config={'temperature': 11}, ) - assert (await response).text == '[ECHO] user: "hi" {"temperature":11}' + assert (await stream_result.response).text == '[ECHO] user: "hi" {"temperature":11.0}' @pytest.mark.asyncio @@ -298,7 +299,7 @@ async def test_generate_with_system_prompt_messages( assert response.text == want_txt - _, response = ai.generate_stream( + stream_result = ai.generate_stream( system='talk like pirate', prompt='hi again', messages=[ @@ -313,7 +314,7 @@ async def test_generate_with_system_prompt_messages( ], ) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt @pytest.mark.asyncio @@ -325,7 +326,7 @@ class ToolInput(BaseModel): value: int | None = Field(None, description='value field') @ai.tool(name='testTool') - def test_tool(input: ToolInput) -> int: + async def test_tool(input: ToolInput) -> int: """The tool.""" return input.value or 0 @@ -362,14 +363,14 @@ def test_tool(input: ToolInput) -> int: assert echo.last_request is not None assert echo.last_request.tools == want_request - _, response = ai.generate_stream( + stream_result = ai.generate_stream( model='echoModel', prompt='hi', tool_choice=ToolChoice.REQUIRED, tools=['testTool'], ) - assert (await response).text == want_txt + assert (await stream_result.response).text == want_txt assert echo.last_request is not None assert echo.last_request.tools == want_request @@ -385,16 +386,16 @@ class ToolInput(BaseModel): value: int | None = Field(None, description='value field') @ai.tool(name='test_tool') - def test_tool(input: ToolInput) -> int: + async def test_tool(input: ToolInput) -> int: """The tool.""" return (input.value or 0) + 7 @ai.tool(name='test_interrupt') - def test_interrupt(input: ToolInput, ctx: ToolRunContext) -> None: + async def test_interrupt(input: ToolInput, ctx: ToolRunContext) -> None: """The interrupt.""" ctx.interrupt({'banana': 'yes please'}) - tool_request_msg = MessageWrapper( + tool_request_msg = Message( Message( role=Role.MODEL, content=[ @@ -407,13 +408,13 @@ def test_interrupt(input: ToolInput, ctx: ToolRunContext) -> None: ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=tool_request_msg, ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='tool called'))]), ) @@ -463,7 +464,7 @@ def test_interrupt(input: ToolInput, ctx: ToolRunContext) -> None: ] assert response.text == 'call these tools' - assert response.message == MessageWrapper( + assert response.message == Message( Message( role=Role.MODEL, content=[ @@ -498,16 +499,16 @@ class ToolInput(BaseModel): value: int | None = Field(None, description='value field') @ai.tool(name='test_tool') - def test_tool(input: ToolInput) -> int: + async def test_tool(input: ToolInput) -> int: """The tool.""" return (input.value or 0) + 7 @ai.tool(name='test_interrupt') - def test_interrupt(input: ToolInput, ctx: ToolRunContext) -> None: + async def test_interrupt(input: ToolInput, ctx: ToolRunContext) -> None: """The interrupt.""" ctx.interrupt({'banana': 'yes please'}) - tool_request_msg = MessageWrapper( + tool_request_msg = Message( Message( role=Role.MODEL, content=[ @@ -520,13 +521,13 @@ def test_interrupt(input: ToolInput, ctx: ToolRunContext) -> None: ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=tool_request_msg, ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='tool called'))]), ) @@ -647,11 +648,11 @@ class ToolInput(BaseModel): value: int | None = Field(None, description='value field') @ai.tool(name='testTool') - def test_tool(input: ToolInput) -> str: + async def test_tool(input: ToolInput) -> str: """The tool.""" return 'abc' - tool_request_msg = MessageWrapper( + tool_request_msg = Message( Message( role=Role.MODEL, content=[ @@ -660,13 +661,13 @@ def test_tool(input: ToolInput) -> str: ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=tool_request_msg, ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='tool called'))]), ) @@ -719,11 +720,11 @@ class ToolInput(BaseModel): value: int | None = Field(None, description='value field') @ai.tool(name='testTool') - def test_tool(input: ToolInput) -> str: + async def test_tool(input: ToolInput) -> str: """The tool.""" return 'abc' - tool_request_msg = MessageWrapper( + tool_request_msg = Message( Message( role=Role.MODEL, content=[ @@ -732,28 +733,28 @@ def test_tool(input: ToolInput) -> str: ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=tool_request_msg, ) ) pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='tool called'))]), ) ) pm.chunks = [ [ - GenerateResponseChunk( + ModelResponseChunk( role=Role(tool_request_msg.role), content=tool_request_msg.content, ) ], - [GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='tool called'))])], + [ModelResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='tool called'))])], ] - stream, aresponse = ai.generate_stream( + stream_result = ai.generate_stream( model='programmableModel', prompt='hi', tool_choice=ToolChoice.REQUIRED, @@ -761,7 +762,7 @@ def test_tool(input: ToolInput) -> str: ) chunks = [] - async for chunk in stream: + async for chunk in stream_result.stream: summary = '' if chunk.role: summary += f'{chunk.role} ' @@ -771,7 +772,7 @@ def test_tool(input: ToolInput) -> str: summary += f' {p.root.text}' chunks.append(summary) - response = await aresponse + response = await stream_result.response assert response.text == 'tool called' assert response.request is not None @@ -797,21 +798,21 @@ async def test_generate_stream_no_need_to_await_response( ai, _, pm, *_ = setup_test pm.responses.append( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='something else'))]), ) ) pm.chunks = [ [ - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='h'))]), - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='i'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='h'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='i'))]), ], ] - stream, _ = ai.generate_stream(model='programmableModel', prompt='do it') + stream_result = ai.generate_stream(model='programmableModel', prompt='do it') chunks = '' - async for chunk in stream: + async for chunk in stream_result.stream: chunks += chunk.text assert chunks == 'hi' @@ -825,54 +826,59 @@ class TestSchema(BaseModel): foo: int | None = Field(None, description='foo field') bar: str | None = Field(None, description='bar field') - want = GenerateRequest( + _schema = { + 'properties': { + 'foo': { + 'anyOf': [{'type': 'integer'}, {'type': 'null'}], + 'default': None, + 'description': 'foo field', + 'title': 'Foo', + }, + 'bar': { + 'anyOf': [{'type': 'string'}, {'type': 'null'}], + 'default': None, + 'description': 'bar field', + 'title': 'Bar', + }, + }, + 'title': 'TestSchema', + 'type': 'object', + } + want = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='hi'))]), ], - config={}, + config={}, # type: ignore[arg-type] tools=[], - output=OutputConfig( - format='json', - schema={ - 'properties': { - 'foo': { - 'anyOf': [{'type': 'integer'}, {'type': 'null'}], - 'default': None, - 'description': 'foo field', - 'title': 'Foo', - }, - 'bar': { - 'anyOf': [{'type': 'string'}, {'type': 'null'}], - 'default': None, - 'description': 'bar field', - 'title': 'Bar', - }, - }, - 'title': 'TestSchema', - 'type': 'object', - }, - constrained=True, - content_type='application/json', - ), + output_format='json', + output_schema=_schema, + output_constrained=True, + output_content_type='application/json', ) response = await ai.generate( model='echoModel', prompt='hi', - output=Output(schema=TestSchema, format='json', content_type='application/json', constrained=True), - output_instructions=False, + output_schema=TestSchema, + output_format='json', + output_content_type='application/json', + output_constrained=True, + output_instructions='', ) assert response.request == want - _, response = ai.generate_stream( + stream_result = ai.generate_stream( model='echoModel', prompt='hi', - output=Output(schema=TestSchema, format='json', content_type='application/json', constrained=True), - output_instructions=False, + output_schema=TestSchema, + output_format='json', + output_content_type='application/json', + output_constrained=True, + output_instructions='', ) - assert (await response).request == want + assert (await stream_result.response).request == want @pytest.mark.asyncio @@ -886,53 +892,52 @@ class TestSchema(BaseModel): foo: int | None = Field(None, description='foo field') bar: str | None = Field(None, description='bar field') - want = GenerateRequest( + _schema = { + 'properties': { + 'foo': { + 'anyOf': [{'type': 'integer'}, {'type': 'null'}], + 'default': None, + 'description': 'foo field', + 'title': 'Foo', + }, + 'bar': { + 'anyOf': [{'type': 'string'}, {'type': 'null'}], + 'default': None, + 'description': 'bar field', + 'title': 'Bar', + }, + }, + 'title': 'TestSchema', + 'type': 'object', + } + want = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='hi'))]), ], - config={}, + config={}, # type: ignore[arg-type] tools=[], - output=OutputConfig( - format='json', - schema={ - 'properties': { - 'foo': { - 'anyOf': [{'type': 'integer'}, {'type': 'null'}], - 'default': None, - 'description': 'foo field', - 'title': 'Foo', - }, - 'bar': { - 'anyOf': [{'type': 'string'}, {'type': 'null'}], - 'default': None, - 'description': 'bar field', - 'title': 'Bar', - }, - }, - 'title': 'TestSchema', - 'type': 'object', - }, - # these get populated by the format - constrained=True, - content_type='application/json', - ), + output_format='json', + output_schema=_schema, + # these get populated by the format + output_constrained=True, + output_content_type='application/json', ) response = await ai.generate( model='echoModel', prompt='hi', - output=Output(schema=TestSchema), + output_schema=TestSchema, ) assert response.request == want - _, response = ai.generate_stream( + stream_result = ai.generate_stream( model='echoModel', prompt='hi', - output=Output(schema=TestSchema), + output_schema=TestSchema, ) - assert (await response).request == want + assert (await stream_result.response).request == want @pytest.mark.asyncio @@ -946,52 +951,52 @@ class TestSchema(BaseModel): foo: int | None = Field(None, description='foo field') bar: str | None = Field(None, description='bar field') - want = GenerateRequest( + want = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='hi'))]), ], - config={}, + config={}, # type: ignore[arg-type] tools=[], - output=OutputConfig( - format='json', - schema={ - 'properties': { - 'foo': { - 'anyOf': [{'type': 'integer'}, {'type': 'null'}], - 'default': None, - 'description': 'foo field', - 'title': 'Foo', - }, - 'bar': { - 'anyOf': [{'type': 'string'}, {'type': 'null'}], - 'default': None, - 'description': 'bar field', - 'title': 'Bar', - }, + output_format='json', + output_schema={ + 'properties': { + 'foo': { + 'anyOf': [{'type': 'integer'}, {'type': 'null'}], + 'default': None, + 'description': 'foo field', + 'title': 'Foo', + }, + 'bar': { + 'anyOf': [{'type': 'string'}, {'type': 'null'}], + 'default': None, + 'description': 'bar field', + 'title': 'Bar', }, - 'title': 'TestSchema', - 'type': 'object', }, - constrained=False, - content_type='application/json', - ), + 'title': 'TestSchema', + 'type': 'object', + }, + output_constrained=False, + output_content_type='application/json', ) response = await ai.generate( model='echoModel', prompt='hi', - output=Output(schema=TestSchema, constrained=False), + output_schema=TestSchema, + output_constrained=False, ) assert response.request == want - _, response = ai.generate_stream( + stream_result = ai.generate_stream( model='echoModel', prompt='hi', - output=Output(schema=TestSchema, constrained=False), + output_schema=TestSchema, + output_constrained=False, ) - assert (await response).request == want + assert (await stream_result.response).request == want @pytest.mark.asyncio @@ -1001,10 +1006,12 @@ async def test_generate_with_middleware( """When middleware is provided, applies it.""" ai, *_ = setup_test - async def pre_middle(req: GenerateRequest, ctx: ActionRunContext, next: ModelMiddlewareNext) -> GenerateResponse: + async def pre_middle( + req: ModelRequest, ctx: ActionRunContext, next: Callable[..., Awaitable[ModelResponse]] + ) -> ModelResponse: txt = ''.join(text_from_message(m) for m in req.messages) return await next( - GenerateRequest( + ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text=f'PRE {txt}'))]), ], @@ -1012,11 +1019,13 @@ async def pre_middle(req: GenerateRequest, ctx: ActionRunContext, next: ModelMid ctx, ) - async def post_middle(req: GenerateRequest, ctx: ActionRunContext, next: ModelMiddlewareNext) -> GenerateResponse: - resp: GenerateResponse = await next(req, ctx) + async def post_middle( + req: ModelRequest, ctx: ActionRunContext, next: Callable[..., Awaitable[ModelResponse]] + ) -> ModelResponse: + resp: ModelResponse = await next(req, ctx) assert resp.message is not None txt = text_from_message(resp.message) - return GenerateResponse( + return ModelResponse( finish_reason=resp.finish_reason, message=Message(role=Role.USER, content=[Part(root=TextPart(text=f'{txt} POST'))]), ) @@ -1027,9 +1036,9 @@ async def post_middle(req: GenerateRequest, ctx: ActionRunContext, next: ModelMi assert response.text == want - _, response = ai.generate_stream(model='echoModel', prompt='hi', use=[pre_middle, post_middle]) + stream_result = ai.generate_stream(model='echoModel', prompt='hi', use=[pre_middle, post_middle]) - assert (await response).text == want + assert (await stream_result.response).text == want @pytest.mark.asyncio @@ -1040,11 +1049,11 @@ async def test_generate_passes_through_current_action_context( ai, *_ = setup_test async def inject_context( - req: GenerateRequest, ctx: ActionRunContext, next: ModelMiddlewareNext - ) -> GenerateResponse: + req: ModelRequest, ctx: ActionRunContext, next: Callable[..., Awaitable[ModelResponse]] + ) -> ModelResponse: txt = ''.join(text_from_message(m) for m in req.messages) return await next( - GenerateRequest( + ModelRequest( messages=[ Message( role=Role.USER, @@ -1055,11 +1064,11 @@ async def inject_context( ctx, ) - async def action_fn() -> GenerateResponse: + async def action_fn() -> ModelResponse: return await ai.generate(model='echoModel', prompt='hi', use=[inject_context]) action = ai.registry.register_action(name='test_action', kind=ActionKind.CUSTOM, fn=action_fn) - action_response = await action.arun(context={'foo': 'bar'}) + action_response = await action.run(context={'foo': 'bar'}) assert action_response.response.text == '''[ECHO] user: "hi {'foo': 'bar'}"''' @@ -1072,11 +1081,11 @@ async def test_generate_uses_explicitly_passed_in_context( ai, *_ = setup_test async def inject_context( - req: GenerateRequest, ctx: ActionRunContext, next: ModelMiddlewareNext - ) -> GenerateResponse: + req: ModelRequest, ctx: ActionRunContext, next: Callable[..., Awaitable[ModelResponse]] + ) -> ModelResponse: txt = ''.join(text_from_message(m) for m in req.messages) return await next( - GenerateRequest( + ModelRequest( messages=[ Message( role=Role.USER, @@ -1087,7 +1096,7 @@ async def inject_context( ctx, ) - async def action_fn() -> GenerateResponse: + async def action_fn() -> ModelResponse: return await ai.generate( model='echoModel', prompt='hi', @@ -1096,7 +1105,7 @@ async def action_fn() -> GenerateResponse: ) action = ai.registry.register_action(name='test_action', kind=ActionKind.CUSTOM, fn=action_fn) - action_response = await action.arun(context={'foo': 'bar'}) + action_response = await action.run(context={'foo': 'bar'}) assert action_response.response.text == '''[ECHO] user: "hi {'bar': 'baz'}"''' @@ -1105,14 +1114,30 @@ async def action_fn() -> GenerateResponse: async def test_generate_json_format_unconstrained_with_instructions( setup_test: SetupFixture, ) -> None: - """When Output is provided, format will default to json.""" + """When output_instructions is provided, instructions are injected.""" ai, *_ = setup_test class TestSchema(BaseModel): foo: int | None = Field(None, description='foo field') bar: str | None = Field(None, description='bar field') - want = GenerateRequest( + # Explicit instructions text to inject (matches formatter output for this schema) + instructions_text = ( + 'Output should be in JSON format and conform to the ' + 'following schema:\n\n```\n{\n "properties": {\n ' + '"foo": {\n "anyOf": [\n {\n ' + '"type": "integer"\n },\n {\n ' + '"type": "null"\n }\n ],\n ' + '"default": null,\n "description": "foo field",\n ' + '"title": "Foo"\n },\n "bar": {\n ' + '"anyOf": [\n {\n "type": "string"\n },\n ' + '{\n "type": "null"\n }\n ],\n ' + '"default": null,\n "description": "bar field",\n ' + '"title": "Bar"\n }\n },\n "title": "TestSchema",\n ' + '"type": "object"\n}\n```\n' + ) + + want = ModelRequest( messages=[ Message( role=Role.USER, @@ -1120,70 +1145,57 @@ class TestSchema(BaseModel): Part(root=TextPart(text='hi')), Part( root=TextPart( - text=( - 'Output should be in JSON format and conform to the ' - 'following schema:\n\n```\n{\n "properties": {\n ' - '"foo": {\n "anyOf": [\n {\n ' - '"type": "integer"\n },\n {\n ' - '"type": "null"\n }\n ],\n ' - '"default": null,\n "description": "foo field",\n ' - '"title": "Foo"\n },\n "bar": {\n ' - '"anyOf": [\n {\n "type": "string"\n },\n ' - '{\n "type": "null"\n }\n ],\n ' - '"default": null,\n "description": "bar field",\n ' - '"title": "Bar"\n }\n },\n "title": "TestSchema",\n ' - '"type": "object"\n}\n```\n' - ), + text=instructions_text, metadata=Metadata(root={'purpose': 'output'}), ) ), ], ) ], - config={}, + config={}, # type: ignore[arg-type] tools=[], - output=OutputConfig( - format='json', - schema={ - 'properties': { - 'foo': { - 'anyOf': [{'type': 'integer'}, {'type': 'null'}], - 'default': None, - 'description': 'foo field', - 'title': 'Foo', - }, - 'bar': { - 'anyOf': [{'type': 'string'}, {'type': 'null'}], - 'default': None, - 'description': 'bar field', - 'title': 'Bar', - }, + output_format='json', + output_schema={ + 'properties': { + 'foo': { + 'anyOf': [{'type': 'integer'}, {'type': 'null'}], + 'default': None, + 'description': 'foo field', + 'title': 'Foo', + }, + 'bar': { + 'anyOf': [{'type': 'string'}, {'type': 'null'}], + 'default': None, + 'description': 'bar field', + 'title': 'Bar', }, - 'title': 'TestSchema', - 'type': 'object', }, - constrained=False, - content_type='application/json', - ), + 'title': 'TestSchema', + 'type': 'object', + }, + output_constrained=False, + output_content_type='application/json', ) response = await ai.generate( model='echoModel', prompt='hi', - output=Output(schema=TestSchema, constrained=False), - output_instructions=True, + output_schema=TestSchema, + output_constrained=False, + output_instructions=instructions_text, ) assert response.request == want - _, response = ai.generate_stream( + stream_result = ai.generate_stream( model='echoModel', prompt='hi', - output=Output(schema=TestSchema, constrained=False), - output_instructions=True, + output_schema=TestSchema, + output_constrained=False, + output_instructions=instructions_text, ) - assert (await response).request == want + assert (await stream_result.response).request == want @pytest.mark.asyncio @@ -1213,24 +1225,24 @@ async def test_generate_simulates_doc_grounding( content=[Part(root=TextPart(text='hi'))], ), ], - docs=[DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))])], + docs=[Document(content=[DocumentPart(root=TextPart(text='doc content 1'))])], ) assert response.request is not None assert response.request.messages is not None assert response.request.messages[0] == want_msg - _, response = ai.generate_stream( + stream_result = ai.generate_stream( messages=[ Message( role=Role.USER, content=[Part(root=TextPart(text='hi'))], ), ], - docs=[DocumentData(content=[DocumentPart(root=TextPart(text='doc content 1'))])], + docs=[Document(content=[DocumentPart(root=TextPart(text='doc content 1'))])], ) - resp = await response + resp = await stream_result.response assert resp.request is not None assert resp.request.messages is not None assert resp.request.messages[0] == want_msg @@ -1258,7 +1270,7 @@ def message_parser(msg: Message) -> str: parts = [p.root.text or '' for p in msg.content if hasattr(p.root, 'text') and p.root.text] return f'banana {"".join(parts)}' # type: ignore[arg-type] - def chunk_parser(chunk: GenerateResponseChunk) -> str: + def chunk_parser(chunk: ModelResponseChunk) -> str: """Parse the chunk.""" parts = [p.root.text or '' for p in chunk.content if hasattr(p.root, 'text') and p.root.text] return f'banana chunk {"".join(parts)}' # type: ignore[arg-type] @@ -1288,7 +1300,7 @@ class TestSchema(BaseModel): pm.responses = [ ( - GenerateResponse( + ModelResponse( finish_reason=FinishReason.STOP, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='model says'))]), ) @@ -1296,29 +1308,30 @@ class TestSchema(BaseModel): ] pm.chunks = [ [ - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='1'))]), - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='2'))]), - GenerateResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='3'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='1'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='2'))]), + ModelResponseChunk(role=Role.MODEL, content=[Part(root=TextPart(text='3'))]), ] ] chunks = [] - stream, aresponse = ai.generate_stream( + stream_result = ai.generate_stream( model='programmableModel', prompt='hi', - output=Output(schema=TestSchema, format='banana'), + output_schema=TestSchema, + output_format='banana', ) - async for chunk in stream: + async for chunk in stream_result.stream: chunks.append(chunk.output) - response = await aresponse + response = await stream_result.response assert response.output == 'banana model says' assert chunks == ['banana chunk 1', 'banana chunk 2', 'banana chunk 3'] - assert response.request == GenerateRequest( + assert response.request == ModelRequest( messages=[ Message( role=Role.USER, @@ -1339,32 +1352,29 @@ class TestSchema(BaseModel): ], ), ], - config={}, + config={}, # type: ignore[arg-type] tools=[], - output=OutputConfig( - format='json', - schema={ - 'properties': { - 'foo': { - 'anyOf': [{'type': 'integer'}, {'type': 'null'}], - 'default': None, - 'description': 'foo field', - 'title': 'Foo', - }, - 'bar': { - 'anyOf': [{'type': 'string'}, {'type': 'null'}], - 'default': None, - 'description': 'bar field', - 'title': 'Bar', - }, + output_format='json', + output_schema={ + 'properties': { + 'foo': { + 'anyOf': [{'type': 'integer'}, {'type': 'null'}], + 'default': None, + 'description': 'foo field', + 'title': 'Foo', + }, + 'bar': { + 'anyOf': [{'type': 'string'}, {'type': 'null'}], + 'default': None, + 'description': 'bar field', + 'title': 'Bar', }, - 'title': 'TestSchema', - 'type': 'object', }, - # these get populated by the format - constrained=True, - content_type='application/banana', - ), + 'title': 'TestSchema', + 'type': 'object', + }, + output_constrained=True, + output_content_type='application/banana', ) @@ -1372,8 +1382,8 @@ def test_define_model_default_metadata(setup_test: SetupFixture) -> None: """Test that the define model function works.""" ai, _, _, *_ = setup_test - def foo_model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='banana!'))])) + async def foo_model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='banana!'))])) action = ai.define_model( name='foo', @@ -1393,8 +1403,8 @@ class Config(BaseModel): field_a: str = Field(description='a field') field_b: str = Field(description='b field') - def foo_model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='banana!'))])) + async def foo_model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='banana!'))])) action = ai.define_model( name='foo', @@ -1430,8 +1440,8 @@ def test_define_model_with_info(setup_test: SetupFixture) -> None: """Test that the define model function with info works.""" ai, _, _, *_ = setup_test - def foo_model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: - return GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='banana!'))])) + async def foo_model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + return ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='banana!'))])) action = ai.define_model( name='foo', @@ -1447,65 +1457,6 @@ def foo_model_fn(request: GenerateRequest, ctx: ActionRunContext) -> GenerateRes } -def test_define_retriever_default_metadata(setup_test: SetupFixture) -> None: - """Test that the define retriever function works.""" - ai, _, _, *_ = setup_test - - def my_retriever(doc: Document, config: dict[str, Any] | None = None) -> RetrieverResponse: - return RetrieverResponse(documents=[Document.from_text('Hello'), Document.from_text('World')]) - - action = ai.define_retriever( - name='fooRetriever', - fn=my_retriever, - ) - - assert action.metadata['retriever'] == { - 'label': 'fooRetriever', - } - - -def test_define_retriever_with_schema(setup_test: SetupFixture) -> None: - """Test that the define retriever function with schema works.""" - ai, _, _, *_ = setup_test - - class Config(BaseModel): - field_a: str = Field(description='a field') - field_b: str = Field(description='b field') - - def my_retriever(doc: Document, config: dict[str, Any] | None = None) -> RetrieverResponse: - return RetrieverResponse(documents=[Document.from_text('Hello'), Document.from_text('World')]) - - action = ai.define_retriever( - name='fooRetriever', - fn=my_retriever, - config_schema=Config, - ) - - assert action.metadata['retriever'] == { - 'customOptions': { - 'properties': { - 'field_a': { - 'description': 'a field', - 'title': 'Field A', - 'type': 'string', - }, - 'field_b': { - 'description': 'b field', - 'title': 'Field B', - 'type': 'string', - }, - }, - 'required': [ - 'field_a', - 'field_b', - ], - 'title': 'Config', - 'type': 'object', - }, - 'label': 'fooRetriever', - } - - def test_define_evaluator_simple(setup_test: SetupFixture) -> None: """Test that the define evaluator function works.""" ai, _, _, *_ = setup_test @@ -1611,27 +1562,27 @@ async def my_eval_fn(req: EvalRequest, options: object | None) -> list[EvalFnRes @pytest.mark.asyncio async def test_define_sync_flow(setup_test: SetupFixture) -> None: - """Test defining a synchronous flow.""" + """Test defining an async flow (renamed from sync test - sync flows no longer supported).""" ai, _, _, *_ = setup_test - @cast(Any, ai.flow()) - def my_flow(input: str, ctx: ActionRunContext | None = None) -> str: - if ctx: - ctx.send_chunk(1) - ctx.send_chunk(2) - ctx.send_chunk(3) + @ai.flow() + async def my_flow(input: str, ctx: ActionRunContext) -> str: + # Use ctx.send_chunk() for streaming + ctx.send_chunk(1) + ctx.send_chunk(2) + ctx.send_chunk(3) return input - assert my_flow('banana') == 'banana' + assert (await my_flow('banana')) == 'banana' - stream, response = my_flow.stream('banana2') + result = my_flow.stream('banana2') chunks = [] - async for chunk in stream: + async for chunk in result.stream: chunks.append(chunk) assert chunks == [1, 2, 3] - assert (await response).response == 'banana2' + assert await result.response == 'banana2' @pytest.mark.asyncio @@ -1640,23 +1591,23 @@ async def test_define_async_flow(setup_test: SetupFixture) -> None: ai, _, _, *_ = setup_test @ai.flow() - async def my_flow(input: str, ctx: ActionRunContext | None = None) -> str: - if ctx: - ctx.send_chunk(1) - ctx.send_chunk(2) - ctx.send_chunk(3) + async def my_flow(input: str, ctx: ActionRunContext) -> str: + # Use ctx.send_chunk() for streaming + ctx.send_chunk(1) + ctx.send_chunk(2) + ctx.send_chunk(3) return input assert (await my_flow('banana')) == 'banana' - stream, response = my_flow.stream('banana2') + result = my_flow.stream('banana2') chunks = [] - async for chunk in stream: + async for chunk in result.stream: chunks.append(chunk) assert chunks == [1, 2, 3] - assert (await response).response == 'banana2' + assert await result.response == 'banana2' @pytest.mark.asyncio diff --git a/py/plugins/README.md b/py/plugins/README.md index 70d8683d30..1269a69e7a 100644 --- a/py/plugins/README.md +++ b/py/plugins/README.md @@ -4,14 +4,14 @@ This directory contains all official Genkit plugins for Python. ## Plugin Architecture -All plugins inherit from `genkit.core.plugin.Plugin` and implement three +All plugins inherit from `genkit._core.plugin.Plugin` and implement three async methods. The registry calls them lazily β€” `init()` runs only on first use, not at registration time. ``` β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ Plugin (Abstract Base Class) β”‚ - β”‚ genkit.core.plugin.Plugin β”‚ + β”‚ genkit._core.plugin.Plugin β”‚ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”‚ β”‚ β”‚ name: str Plugin namespace (e.g., 'googleai') β”‚ @@ -98,14 +98,8 @@ first use, not at registration time. β”‚ β”‚ anthropic β”‚ β”‚ firebase β”‚ β”‚ β”‚ β”‚ β€’ Claude 3.5/4 β”‚ β”‚ β€’ Firebase Telemetry β”‚ β”‚ β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ -β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ -β”‚ β”‚ amazon-bedrock 🌐 β”‚ INTEGRATIONS β”‚ -β”‚ β”‚ β€’ Claude, Llama, Nova β”‚ ──────────── β”‚ -β”‚ β”‚ β€’ Titan, Mistral β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ -β”‚ β”‚ β€’ X-Ray telemetry β”‚ β”‚ flask β”‚ β”‚ -β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ β€’ HTTP endpoints β”‚ β”‚ -β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ -β”‚ β”‚ microsoft-foundry β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ +β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ +β”‚ β”‚ microsoft-foundry β”‚ β”‚ flask β”‚ β”‚ β”‚ β”‚ β€’ GPT-4o, Claude, Llama β”‚ β”‚ mcp β”‚ β”‚ β”‚ β”‚ β€’ 11,000+ models β”‚ β”‚ β€’ Model Context Protocolβ”‚ β”‚ β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ @@ -162,10 +156,7 @@ first use, not at registration time. β”‚ β†’ google-genai (Gemini 2.0) β”‚ β”‚ β”‚ β”‚ "I need Claude models" β”‚ -β”‚ β†’ anthropic (direct) OR amazon-bedrock OR microsoft-foundry β”‚ -β”‚ β”‚ -β”‚ "I'm on AWS and want managed models" β”‚ -β”‚ β†’ amazon-bedrock (Claude, Llama, Nova, Titan) β”‚ +β”‚ β†’ anthropic (direct) OR microsoft-foundry β”‚ β”‚ β”‚ β”‚ "I'm on Azure and want managed models" β”‚ β”‚ β†’ microsoft-foundry (GPT-4o, Claude, Llama, 11,000+ models) β”‚ @@ -245,7 +236,6 @@ first use, not at registration time. β”‚ β”‚ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”‚ β”‚ -β”‚ "I'm on AWS and want X-Ray" β†’ amazon-bedrock plugin β”‚ β”‚ "I'm on GCP and want Cloud Trace" β†’ google-cloud plugin β”‚ β”‚ "I'm on Azure and want App Insights" β†’ microsoft-foundry plugin β”‚ β”‚ "I'm using Firebase" β†’ firebase plugin (auto telemetry) β”‚ @@ -311,7 +301,6 @@ first use, not at registration time. |--------|--------|----------| | **google-genai** | Gemini, Imagen, Veo, Lyria | Multimodal AI, Google ecosystem | | **anthropic** | Claude 3.5, Claude 4 | Direct Claude access | -| **amazon-bedrock** 🌐 | Claude, Llama, Nova, Titan | AWS managed models (community) | | **microsoft-foundry** 🌐 | GPT-4o, Claude, Llama, 11,000+ | Azure AI, enterprise (community) | | **vertex-ai** | Model Garden (Claude, Llama) | GCP third-party models | | **ollama** | Llama, Mistral, Phi, etc. | Local/private deployment | @@ -337,7 +326,6 @@ first use, not at registration time. | Plugin | Backend | Features | |--------|---------|----------| | **google-cloud** | Cloud Trace, Logging | GCP native, log correlation | -| **amazon-bedrock** 🌐 | X-Ray | AWS native, SigV4 auth, built into model plugin (community) | | **microsoft-foundry** 🌐 | Application Insights | Azure Monitor, trace correlation, built into model plugin (community) | | **cloudflare-workers-ai** 🌐 | Any OTLP endpoint | Generic OTLP, Bearer auth, combined with models (community) | | **observability** 🌐 | Sentry, Honeycomb, Datadog, Grafana, Axiom | 3rd party presets (community) | @@ -376,9 +364,6 @@ All environment variables used by Genkit plugins. Configure these before running |----------|--------|----------|-------------|---------------| | `GEMINI_API_KEY` | google-genai | Yes | Google AI Studio API key | [Get API Key](https://aistudio.google.com/apikey) | | `ANTHROPIC_API_KEY` | anthropic | Yes | Anthropic API key | [Anthropic Console](https://console.anthropic.com/) | -| `AWS_REGION` | amazon-bedrock | Yes | AWS region (e.g., `us-east-1`) | [AWS Regions](https://docs.aws.amazon.com/general/latest/gr/bedrock.html) | -| `AWS_ACCESS_KEY_ID` | amazon-bedrock | Yes* | AWS access key | [AWS Credentials](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html) | -| `AWS_SECRET_ACCESS_KEY` | amazon-bedrock | Yes* | AWS secret key | [AWS Credentials](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html) | | `AZURE_AI_FOUNDRY_ENDPOINT` | microsoft-foundry | Yes | Azure AI Foundry endpoint URL | [Azure AI Foundry](https://ai.azure.com/) | | `AZURE_AI_FOUNDRY_API_KEY` | microsoft-foundry | Yes* | Azure AI Foundry API key | [Azure AI Foundry](https://ai.azure.com/) | | `OPENAI_API_KEY` | compat-oai | Yes | OpenAI API key | [OpenAI API Keys](https://platform.openai.com/api-keys) | @@ -511,7 +496,6 @@ Each plugin is a separate package. Install only what you need: # Model providers pip install genkit-google-genai-plugin pip install genkit-anthropic-plugin -pip install genkit-amazon-bedrock-plugin # Also includes X-Ray telemetry pip install genkit-microsoft-foundry-plugin # Telemetry @@ -531,7 +515,7 @@ pip install genkit-cohere-plugin ## Quick Start ```python -from genkit.ai import Genkit +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI # Initialize with your chosen plugin @@ -580,7 +564,7 @@ plugins are independent leaf nodes; only a few have inter-plugin dependencies. β”‚ β”‚ β”‚ INDEPENDENT PLUGINS (no inter-plugin dependencies): β”‚ β”‚ ───────────────────────────────────────────────── β”‚ -β”‚ google-genai, anthropic, amazon-bedrock, microsoft-foundry, β”‚ +β”‚ google-genai, anthropic, microsoft-foundry, β”‚ β”‚ ollama, xai, mistral, huggingface, cloudflare-workers-ai, β”‚ β”‚ cohere, google-cloud, firebase, observability, mcp, fastapi, β”‚ β”‚ evaluators, dev-local-vectorstore, checks β”‚ diff --git a/py/plugins/amazon-bedrock/LICENSE b/py/plugins/amazon-bedrock/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/plugins/amazon-bedrock/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/plugins/amazon-bedrock/README.md b/py/plugins/amazon-bedrock/README.md deleted file mode 100644 index a77ad4857d..0000000000 --- a/py/plugins/amazon-bedrock/README.md +++ /dev/null @@ -1,439 +0,0 @@ -# Genkit Amazon Bedrock Plugin - -> **Community Plugin** – This plugin is maintained by the community and is supported on a best-effort basis. It is not an official AWS product. -> -> **Preview** β€” This plugin is in preview and may have API changes in future releases. - -`genkit-plugin-amazon-bedrock` is a plugin for using Amazon Bedrock models with [Genkit](https://github.com/firebase/genkit). It provides access to foundation models AND AWS X-Ray telemetry in a single package. - -Amazon Bedrock is a fully managed service that provides access to foundation models from leading AI providers including Amazon, Anthropic, Meta, Mistral, Cohere, DeepSeek, and more through a unified API. - -## Documentation Links - -- **Amazon Bedrock Console**: https://console.aws.amazon.com/bedrock/ -- **Supported Models**: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html -- **Model Parameters**: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html -- **Converse API**: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html -- **Boto3 Bedrock Runtime**: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime.html - -## Installation - -```bash -pip install genkit-plugin-amazon-bedrock -``` - -## Setup - -You'll need AWS credentials configured. The plugin supports multiple authentication methods. - -### Option 1: Bedrock API Key (Simplest) - -Amazon Bedrock now supports API keys similar to OpenAI/Anthropic: - -```bash -export AWS_REGION="us-east-1" -export AWS_BEARER_TOKEN_BEDROCK="your-api-key" -``` - -Generate an API key in [Amazon Bedrock Console](https://console.aws.amazon.com/bedrock/) > API keys. - -**Important**: API keys require [inference profiles](#cross-region-inference-profiles) instead of direct model IDs. Use the `inference_profile()` helper: - -```python -from genkit.plugins.amazon_bedrock import inference_profile - -model = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0') -``` - -See: [Getting Started with Bedrock API Keys](https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started-api-keys.html) - -### Option 2: IAM Credentials (Recommended for Production) - -```bash -export AWS_REGION="us-east-1" -export AWS_ACCESS_KEY_ID="your-access-key-id" -export AWS_SECRET_ACCESS_KEY="your-secret-access-key" -``` - -### Option 3: AWS Profile - -```bash -export AWS_PROFILE="your-profile-name" -export AWS_REGION="us-east-1" -``` - -### Option 4: IAM Role (AWS Infrastructure) - -When running on EC2, Lambda, ECS, or EKS, credentials are automatically provided by the IAM role. - -```bash -export AWS_REGION="us-east-1" -# No credentials needed - IAM role provides them -``` - -### IAM Permissions (for Options 2-4) - -Your AWS credentials need the following permissions: - -```json -{ - "Version": "2012-10-17", - "Statement": [ - { - "Effect": "Allow", - "Action": [ - "bedrock:InvokeModel", - "bedrock:InvokeModelWithResponseStream" - ], - "Resource": "arn:aws:bedrock:*::foundation-model/*" - } - ] -} -``` - -## AWS X-Ray Telemetry - -This plugin includes built-in AWS X-Ray telemetry support. Enable it to send distributed traces to AWS X-Ray console: - -```python -from genkit.plugins.amazon_bedrock import add_aws_telemetry - -# Enable X-Ray telemetry (call once at startup) -add_aws_telemetry(region="us-east-1") -``` - -Traces will be visible in the [AWS X-Ray Console](https://console.aws.amazon.com/xray/home). - -### IAM Permissions for X-Ray - -Add these permissions for telemetry: - -```json -{ - "Effect": "Allow", - "Action": ["xray:PutTraceSegments", "xray:PutTelemetryRecords"], - "Resource": "*" -} -``` - -Or attach the managed policy: `AWSXrayWriteOnlyPolicy` - -## Basic Usage - -```python -from genkit import Genkit -from genkit.plugins.amazon_bedrock import AmazonBedrock, bedrock_model - -ai = Genkit( - plugins=[ - AmazonBedrock(region="us-east-1") - ], - model=bedrock_model("anthropic.claude-sonnet-4-5-20250929-v1:0"), -) - -response = await ai.generate(prompt="Tell me a joke.") -print(response.text) -``` - -### With Explicit Credentials - -```python -ai = Genkit( - plugins=[ - AmazonBedrock( - region="us-east-1", - access_key_id="your-access-key", - secret_access_key="your-secret-key", - ) - ], - model=bedrock_model("anthropic.claude-sonnet-4-5-20250929-v1:0"), -) -``` - -### With AWS Profile - -```python -ai = Genkit( - plugins=[ - AmazonBedrock( - region="us-east-1", - profile_name="my-aws-profile", - ) - ], - model=bedrock_model("anthropic.claude-sonnet-4-5-20250929-v1:0"), -) -``` - -## Supported Model Providers - -| Provider | Model Examples | Model ID Prefix | -|----------|---------------|-----------------| -| Amazon | Nova Pro, Nova Lite, Nova Micro | `amazon.nova-*` | -| Anthropic | Claude Sonnet 4.5, Claude Opus 4.5 | `anthropic.claude-*` | -| AI21 Labs | Jamba 1.5 Large, Jamba 1.5 Mini | `ai21.jamba-*` | -| Cohere | Command R, Command R+ | `cohere.command-*` | -| DeepSeek | DeepSeek-R1, DeepSeek-V3 | `deepseek.*` | -| Google | Gemma 3 4B, Gemma 3 12B | `google.gemma-*` | -| Meta | Llama 3.3 70B, Llama 4 Maverick | `meta.llama*` | -| MiniMax | MiniMax M2 | `minimax.*` | -| Mistral | Mistral Large 3, Pixtral Large | `mistral.*` | -| Moonshot | Kimi K2 Thinking | `moonshot.*` | -| NVIDIA | Nemotron Nano 9B, 12B | `nvidia.*` | -| OpenAI | GPT-OSS 120B, GPT-OSS 20B | `openai.*` | -| Qwen | Qwen3 32B, Qwen3 235B | `qwen.*` | -| Writer | Palmyra X4, Palmyra X5 | `writer.*` | - -## Model Examples - -### Anthropic Claude - -```python -from genkit.plugins.amazon_bedrock import ( - AmazonBedrock, - bedrock_model, - claude_sonnet_4_5, - claude_opus_4_5, -) - -ai = Genkit( - plugins=[AmazonBedrock(region="us-east-1")], - model=claude_sonnet_4_5, -) - -response = await ai.generate(prompt="Explain quantum computing") -``` - -### Amazon Nova - -```python -from genkit.plugins.amazon_bedrock import nova_pro, nova_lite - -response = await ai.generate( - model=nova_pro, - prompt="Describe the image", - # Nova supports images and video -) -``` - -### Meta Llama - -```python -from genkit.plugins.amazon_bedrock import llama_3_3_70b, llama_4_maverick - -response = await ai.generate( - model=llama_4_maverick, - prompt="Write a poem about AI", -) -``` - -### DeepSeek R1 (Reasoning) - -```python -from genkit.plugins.amazon_bedrock import deepseek_r1 - -response = await ai.generate( - model=deepseek_r1, - prompt="Solve this step by step: What is 15% of 240?", -) -# Response includes reasoning_content for reasoning models -``` - -## Configuration - -The plugin supports model-specific configuration parameters: - -```python -from genkit.plugins.amazon_bedrock import BedrockConfig - -response = await ai.generate( - prompt="Tell me a story", - config=BedrockConfig( - temperature=0.7, - max_tokens=1000, - top_p=0.9, - stop_sequences=["THE END"], - ), -) -``` - -### Common Parameters - -| Parameter | Type | Description | -|-----------|------|-------------| -| `temperature` | float (0.0-1.0) | Sampling temperature. Higher = more random. | -| `max_tokens` | int | Maximum tokens to generate. | -| `top_p` | float (0.0-1.0) | Nucleus sampling probability. | -| `top_k` | int | Top-k sampling (model-specific). | -| `stop_sequences` | list[str] | Stop sequences to end generation. | - -### Model-Specific Configs - -Each model family has its own configuration class with provider-specific parameters: - -```python -from genkit.plugins.amazon_bedrock import AnthropicConfig, CohereConfig, MetaLlamaConfig - -# Anthropic Claude with top_k -response = await ai.generate( - model=claude_sonnet_4_5, - prompt="...", - config=AnthropicConfig( - temperature=0.7, - top_k=40, - ), -) - -# Cohere with documents for RAG -response = await ai.generate( - model=bedrock_model("cohere.command-r-plus-v1:0"), - prompt="...", - config=CohereConfig( - temperature=0.5, - k=50, - p=0.9, - ), -) - -# Meta Llama -response = await ai.generate( - model=llama_3_3_70b, - prompt="...", - config=MetaLlamaConfig( - temperature=0.6, - max_gen_len=1024, - ), -) -``` - -## Multimodal Support - -Models like Claude, Nova, and Llama 4 support images: - -```python -from genkit.types import Media, MediaPart, Part, TextPart - -response = await ai.generate( - model=claude_sonnet_4_5, - prompt=[ - Part(root=TextPart(text="What's in this image?")), - Part(root=MediaPart(media=Media(url="https://example.com/image.jpg"))), - ], -) -``` - -## Embeddings - -```python -from genkit.blocks.document import Document - -# Amazon Titan Embeddings -response = await ai.embed( - embedder="amazon-bedrock/amazon.titan-embed-text-v2:0", - input=[Document.from_text("Hello, world!")], -) - -# Cohere Embeddings -response = await ai.embed( - embedder="amazon-bedrock/cohere.embed-english-v3", - input=[Document.from_text("Hello, world!")], -) -``` - -## Cross-Region Inference Profiles - -**When using API keys (`AWS_BEARER_TOKEN_BEDROCK`)**, you must use inference profiles instead of direct model IDs. The plugin provides helpers for this: - -### Inference Profile Helper - -```python -from genkit.plugins.amazon_bedrock import inference_profile - -# Auto-detects region from AWS_REGION environment variable -model = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0') -# β†’ 'amazon-bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0' (if AWS_REGION=us-east-1) - -# Or specify region explicitly -model = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0', 'eu-west-1') -# β†’ 'amazon-bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0' -``` - -### Regional Prefixes - -| Region | Prefix | Example Regions | -|--------|--------|-----------------| -| United States | `us.` | us-east-1, us-west-2 | -| Europe | `eu.` | eu-west-1, eu-central-1 | -| Asia Pacific | `apac.` | ap-northeast-1, ap-southeast-1 | - -### When to Use Inference Profiles - -| Auth Method | Model ID Format | -|-------------|-----------------| -| API Key (`AWS_BEARER_TOKEN_BEDROCK`) | Inference profile required: `us.anthropic.claude-...` | -| IAM Credentials | Direct model ID works: `anthropic.claude-...` | - -See: [Cross-Region Inference](https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html) - -## Streaming - -Streaming is supported via the Genkit streaming API: - -```python -async for chunk in ai.generate_stream( - model=claude_sonnet_4_5, - prompt="Write a long story", -): - print(chunk.text, end="", flush=True) -``` - -## Tool Use (Function Calling) - -```python -from genkit.ai import tool - -@tool() -def get_weather(city: str) -> str: - """Get the current weather for a city.""" - return f"The weather in {city} is sunny." - -response = await ai.generate( - model=claude_sonnet_4_5, - prompt="What's the weather in Tokyo?", - tools=[get_weather], -) -``` - -## References - -### AWS Documentation - -- [Amazon Bedrock User Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) -- [Supported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) -- [Model Parameters](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html) -- [Converse API](https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html) -- [Tool Use](https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use.html) - -### Model Provider Documentation - -- [Anthropic Claude](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html) -- [Amazon Nova](https://docs.aws.amazon.com/nova/latest/userguide/what-is-nova.html) -- [Meta Llama](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html) -- [Mistral AI](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html) -- [Cohere](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html) -- [DeepSeek](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-deepseek.html) - -## Disclaimer - -This is a community plugin and is not officially supported or endorsed by Amazon Web Services. - -"Amazon", "AWS", "Amazon Bedrock", and related marks are trademarks of Amazon.com, Inc. -or its affiliates. This plugin is developed independently and is not affiliated with, -endorsed by, or sponsored by Amazon. - -The use of AWS APIs is subject to AWS's terms of service. Users are responsible for -ensuring their usage complies with AWS's API terms and any applicable rate limits or -usage policies. - -## License - -Apache 2.0 diff --git a/py/plugins/amazon-bedrock/pyproject.toml b/py/plugins/amazon-bedrock/pyproject.toml deleted file mode 100644 index 1362db476e..0000000000 --- a/py/plugins/amazon-bedrock/pyproject.toml +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Framework :: AsyncIO", - "Framework :: Pydantic", - "Framework :: Pydantic :: 2", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Typing :: Typed", - "License :: OSI Approved :: Apache Software License", -] -dependencies = [ - "genkit", - "aioboto3>=13.0.0", - "boto3>=1.35.0", - "botocore>=1.35.0", - "opentelemetry-sdk-extension-aws>=2.1.0", - "opentelemetry-propagator-aws-xray>=1.0.0", - "opentelemetry-exporter-otlp-proto-http>=1.20.0", - "strenum>=0.4.15; python_version < '3.11'", -] -description = "Genkit Amazon Bedrock Plugin - Models and AWS Observability (Community)" -keywords = [ - "genkit", - "ai", - "llm", - "machine-learning", - "artificial-intelligence", - "generative-ai", - "aws", - "bedrock", - "amazon", - "xray", - "telemetry", - "observability", -] -license = "Apache-2.0" -name = "genkit-plugin-amazon-bedrock" -readme = "README.md" -requires-python = ">=3.10" -version = "0.5.1" - -[project.urls] -"Bug Tracker" = "https://github.com/firebase/genkit/issues" -Changelog = "https://github.com/firebase/genkit/blob/main/amazon-bedrock/CHANGELOG.md" -"Documentation" = "https://firebase.google.com/docs/genkit" -"Homepage" = "https://github.com/firebase/genkit" -"Repository" = "https://github.com/firebase/genkit/tree/main/py" - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -only-include = ["src/genkit/plugins/amazon_bedrock"] -sources = ["src"] diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/__init__.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/__init__.py deleted file mode 100644 index 7fa860559c..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/__init__.py +++ /dev/null @@ -1,308 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS Bedrock plugin for Genkit. - -This plugin provides access to AWS Bedrock models through the Genkit framework. -AWS Bedrock is a fully managed service that provides access to foundation models -from multiple providers through a unified API. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ AWS Bedrock β”‚ Amazon's AI model marketplace. One place to β”‚ - β”‚ β”‚ access Claude, Llama, Titan, and more. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Converse API β”‚ A unified way to talk to ANY Bedrock model. β”‚ - β”‚ β”‚ Same code works for Claude, Llama, Nova, etc. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Model ID β”‚ The name of a specific model. Like β”‚ - β”‚ β”‚ "anthropic.claude-3-sonnet-20240229-v1:0". β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Inference Profile β”‚ A cross-region alias for a model. Required β”‚ - β”‚ β”‚ when using API keys instead of IAM roles. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Region β”‚ Which AWS data center to use. Pick one near β”‚ - β”‚ β”‚ you (us-east-1, eu-west-1, ap-northeast-1). β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ IAM Role β”‚ AWS's way of granting permissions. Like a β”‚ - β”‚ β”‚ badge that lets your code access models. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Nova β”‚ Amazon's own AI models (Pro, Lite, Micro). β”‚ - β”‚ β”‚ Good balance of cost and performance. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Inline Bytes β”‚ Bedrock needs actual image data, not URLs. β”‚ - β”‚ β”‚ We fetch images for you automatically. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW AWS BEDROCK PROCESSES YOUR REQUEST β”‚ - β”‚ β”‚ - β”‚ Your Code β”‚ - β”‚ ai.generate(prompt="Describe this image", media=[image_url]) β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) Request goes to AmazonBedrock plugin β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AmazonBedrock β”‚ β€’ Adds AWS credentials β”‚ - β”‚ β”‚ Plugin β”‚ β€’ Converts model ID β†’ inference profile β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Fetch image bytes (Bedrock needs inline data) β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ BedrockModel β”‚ β€’ Downloads image from URL β”‚ - β”‚ β”‚ (httpx async) β”‚ β€’ Converts to Converse API format β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Bedrock Converse API call β”‚ - β”‚ β–Ό β”‚ - β”‚ ════════════════════════════════════════════════════ β”‚ - β”‚ β”‚ Internet (HTTPS to bedrock.{region}.amazonaws.com) β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AWS Bedrock β”‚ Routes to the right provider β”‚ - β”‚ β”‚ β”‚ (Claude, Llama, Nova, etc.) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (4) Streaming response β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Your App β”‚ response.text = "I see a cute kitten..." β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Architecture Overview:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ AWS Bedrock Plugin β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Plugin Entry Point (__init__.py) β”‚ - β”‚ β”œβ”€β”€ AmazonBedrock - Plugin class β”‚ - β”‚ β”œβ”€β”€ bedrock_model() / inference_profile() - Helper functions β”‚ - β”‚ └── Pre-defined models (claude_sonnet_4_5, deepseek_r1, nova_pro, ...) β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ plugin.py - Plugin Implementation β”‚ - β”‚ β”œβ”€β”€ AmazonBedrock class (registers models/embedders) β”‚ - β”‚ └── Configuration and boto3 client initialization β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ typing.py - Type-Safe Configuration Classes β”‚ - β”‚ β”œβ”€β”€ BedrockConfig (common base) β”‚ - β”‚ β”œβ”€β”€ Provider-specific configs (AnthropicConfig, MetaLlamaConfig, ...) β”‚ - β”‚ └── Enums (CohereSafetyMode, CohereToolChoice, StabilityMode, ...) β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ models/model.py - Model Implementation β”‚ - β”‚ β”œβ”€β”€ BedrockModel (Converse API integration) β”‚ - β”‚ β”œβ”€β”€ Automatic inference profile conversion for API key auth β”‚ - β”‚ β”œβ”€β”€ Async media URL fetching (httpx) - Bedrock requires inline bytes β”‚ - β”‚ └── JSON mode via prompt engineering (no native support) β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ models/model_info.py - Model Registry β”‚ - β”‚ β”œβ”€β”€ SUPPORTED_BEDROCK_MODELS (12+ providers) β”‚ - β”‚ └── SUPPORTED_EMBEDDING_MODELS β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Supported Model Providers: - - Amazon (Nova Pro, Nova Lite, Nova Micro, Titan) - - Anthropic (Claude Opus, Sonnet, Haiku) - - AI21 Labs (Jamba 1.5) - - Cohere (Command R, Command R+, Embed) - - DeepSeek (R1, V3) - - Google (Gemma 3) - - Meta (Llama 3.x, Llama 4) - - MiniMax (M2) - - Mistral AI (Large 3, Pixtral, Ministral) - - Moonshot AI (Kimi K2) - - NVIDIA (Nemotron) - - OpenAI (GPT-OSS) - - Qwen (Qwen3) - - Writer (Palmyra) - -Documentation Links: - - AWS Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html - - Supported Models: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html - - Converse API: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html - -Example:: - - from genkit import Genkit - from genkit.plugins.amazon_bedrock import AmazonBedrock, claude_sonnet_4_5 - - ai = Genkit( - plugins=[AmazonBedrock(region='us-east-1')], - model=claude_sonnet_4_5, - ) - - response = await ai.generate(prompt='Tell me a joke.') - print(response.text) - -Trademark Notice: - This is a community plugin and is not officially supported by Amazon Web Services. - "Amazon", "AWS", "Amazon Bedrock", and related marks are trademarks of - Amazon.com, Inc. or its affiliates. Model names are trademarks of their - respective owners. -""" - -from .plugin import ( - AMAZON_BEDROCK_PLUGIN_NAME, - AmazonBedrock, - bedrock_model, - bedrock_name, - # Pre-defined model references - Anthropic Claude - claude_3_5_haiku, - claude_3_haiku, - claude_haiku_4_5, - claude_opus_4_1, - claude_opus_4_5, - claude_opus_4_6, - claude_sonnet_4, - claude_sonnet_4_5, - # Pre-defined model references - Cohere - command_r, - command_r_plus, - # Pre-defined model references - DeepSeek - deepseek_r1, - deepseek_v3, - get_config_schema_for_model, - # Inference profile helpers (for API key authentication) - get_inference_profile_prefix, - inference_profile, - # Pre-defined model references - AI21 Jamba - jamba_large, - jamba_mini, - # Pre-defined model references - Meta Llama - llama_3_1_70b, - llama_3_1_405b, - llama_3_3_70b, - llama_4_maverick, - llama_4_scout, - # Pre-defined model references - Mistral - mistral_large, - mistral_large_3, - # Pre-defined model references - Amazon Nova - nova_lite, - nova_micro, - nova_premier, - nova_pro, - pixtral_large, -) -from .telemetry import add_aws_telemetry -from .typing import ( - # Model-Specific Configs - AI21JambaConfig, - AmazonNovaConfig, - AnthropicConfig, - # Base/Common Configs - BedrockConfig, - CohereConfig, - # Enums - CohereSafetyMode, - CohereToolChoice, - DeepSeekConfig, - # Mixin - GenkitCommonConfigMixin, - GoogleGemmaConfig, - MetaLlamaConfig, - MiniMaxConfig, - MistralConfig, - MoonshotConfig, - NvidiaConfig, - OpenAIConfig, - QwenConfig, - StabilityConfig, - # Embedding Config - TextEmbeddingConfig, - TitanConfig, - WriterConfig, -) - -__all__ = [ - 'AMAZON_BEDROCK_PLUGIN_NAME', - 'AI21JambaConfig', - # Plugin - 'AmazonBedrock', - # Model-Specific Configs (16 providers) - 'AmazonNovaConfig', - 'AnthropicConfig', - # Base/Common Configs - 'BedrockConfig', - 'CohereConfig', - # Enums - 'CohereSafetyMode', - 'CohereToolChoice', - 'DeepSeekConfig', - # Mixin - 'GenkitCommonConfigMixin', - 'GoogleGemmaConfig', - 'MetaLlamaConfig', - 'MiniMaxConfig', - 'MistralConfig', - 'MoonshotConfig', - 'NvidiaConfig', - 'OpenAIConfig', - 'QwenConfig', - 'StabilityConfig', - # Embedding Config - 'TextEmbeddingConfig', - 'TitanConfig', - 'WriterConfig', - # Helper functions - 'bedrock_model', - 'bedrock_name', - 'claude_3_5_haiku', - 'claude_3_haiku', - 'claude_haiku_4_5', - 'claude_opus_4_1', - 'claude_opus_4_5', - 'claude_opus_4_6', - 'claude_sonnet_4', - # Pre-defined model references - Anthropic Claude - 'claude_sonnet_4_5', - 'command_r', - # Pre-defined model references - Cohere - 'command_r_plus', - # Pre-defined model references - DeepSeek - 'deepseek_r1', - 'deepseek_v3', - 'get_config_schema_for_model', - # Inference profile helpers (for API key authentication) - 'get_inference_profile_prefix', - 'inference_profile', - # Pre-defined model references - AI21 Jamba - 'jamba_large', - 'jamba_mini', - 'llama_3_1_70b', - 'llama_3_1_405b', - # Pre-defined model references - Meta Llama - 'llama_3_3_70b', - 'llama_4_maverick', - 'llama_4_scout', - 'mistral_large', - # Pre-defined model references - Mistral - 'mistral_large_3', - 'nova_lite', - 'nova_micro', - 'nova_premier', - # Pre-defined model references - Amazon Nova - 'nova_pro', - 'pixtral_large', - # Telemetry - 'add_aws_telemetry', -] diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/__init__.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/__init__.py deleted file mode 100644 index 45fb99058f..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS Bedrock models subpackage.""" - -from .model import BedrockModel -from .model_info import ( - SUPPORTED_BEDROCK_MODELS, - SUPPORTED_EMBEDDING_MODELS, - get_model_info, -) - -__all__ = [ - 'SUPPORTED_BEDROCK_MODELS', - 'SUPPORTED_EMBEDDING_MODELS', - 'BedrockModel', - 'get_model_info', -] diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/converters.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/converters.py deleted file mode 100644 index 6cb8e98d21..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/converters.py +++ /dev/null @@ -1,658 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS Bedrock format conversion utilities. - -Pure-function helpers for converting between Genkit types and AWS Bedrock -Converse API formats. Extracted from the model module for independent -unit testing. - -See: - - Converse API: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html - - Boto3 Reference: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html -""" - -import base64 -import json -import re -from typing import Any - -from genkit.plugins.amazon_bedrock.typing import BedrockConfig -from genkit.types import ( - FinishReason, - GenerateRequest, - GenerationCommonConfig, - GenerationUsage, - Media, - Message, - Part, - Role, - TextPart, - ToolDefinition, - ToolRequest, - ToolRequestPart, - ToolResponsePart, -) - -__all__ = [ - 'FINISH_REASON_MAP', - 'INFERENCE_PROFILE_PREFIXES', - 'INFERENCE_PROFILE_SUPPORTED_PROVIDERS', - 'build_json_instruction', - 'build_media_block', - 'build_usage', - 'convert_media_data_uri', - 'from_bedrock_content', - 'get_effective_model_id', - 'is_image_media', - 'map_finish_reason', - 'normalize_config', - 'parse_tool_call_args', - 'separate_system_messages', - 'to_bedrock_content', - 'maybe_strip_fences', - 'strip_markdown_fences', - 'StreamingFenceStripper', - 'to_bedrock_role', - 'to_bedrock_tool', -] - -# Mapping from Bedrock stop reasons to Genkit finish reasons. -# -# See: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseOutput.html -FINISH_REASON_MAP: dict[str, FinishReason] = { - 'end_turn': FinishReason.STOP, - 'stop_sequence': FinishReason.STOP, - 'max_tokens': FinishReason.LENGTH, - 'tool_use': FinishReason.STOP, - 'content_filtered': FinishReason.BLOCKED, - 'guardrail_intervened': FinishReason.BLOCKED, -} - -# Inference profile prefixes that indicate an ID already has a regional prefix. -INFERENCE_PROFILE_PREFIXES = ('us.', 'eu.', 'apac.') - -# Model provider prefixes that support cross-region inference profiles. -# Only these providers can use inference profile IDs with regional prefixes. -INFERENCE_PROFILE_SUPPORTED_PROVIDERS = ( - 'anthropic.', - 'amazon.', - 'meta.', - 'mistral.', - 'cohere.', - 'deepseek.', -) - - -def map_finish_reason(stop_reason: str) -> FinishReason: - """Map a Bedrock stop reason string to a Genkit FinishReason. - - Args: - stop_reason: The stop reason from the Bedrock API response. - - Returns: - The corresponding Genkit FinishReason, or UNKNOWN if unmapped. - """ - return FINISH_REASON_MAP.get(stop_reason, FinishReason.UNKNOWN) - - -def to_bedrock_role(role: Role | str) -> str: - """Convert a Genkit role to a Bedrock role string. - - The Bedrock Converse API only supports 'user' and 'assistant' roles. - Tool responses are sent as 'user' messages. - - Args: - role: Genkit Role enum or string. - - Returns: - Bedrock role string ('user' or 'assistant'). - """ - if isinstance(role, str): - str_role_map = { - 'user': 'user', - 'model': 'assistant', - 'assistant': 'assistant', - 'tool': 'user', - } - return str_role_map.get(role.lower(), 'user') - - role_map = { - Role.USER: 'user', - Role.MODEL: 'assistant', - Role.TOOL: 'user', - } - return role_map.get(role, 'user') - - -def separate_system_messages( - messages: list[Message], -) -> tuple[list[str], list[Message]]: - """Separate system messages from conversation messages. - - The Bedrock Converse API requires system messages to be passed - separately from conversation messages. - - Args: - messages: List of Genkit messages. - - Returns: - Tuple of (system_texts, conversation_messages). - """ - system_texts: list[str] = [] - conversation_messages: list[Message] = [] - - for msg in messages: - if msg.role == Role.SYSTEM or (isinstance(msg.role, str) and msg.role.lower() == 'system'): - text_parts: list[str] = [] - for part in msg.content: - root = part.root if isinstance(part, Part) else part - if isinstance(root, TextPart): - text_parts.append(root.text) - if text_parts: - system_texts.append(''.join(text_parts)) - else: - conversation_messages.append(msg) - - return system_texts, conversation_messages - - -def to_bedrock_tool(tool: ToolDefinition) -> dict[str, Any]: - """Convert a Genkit tool definition to Bedrock format. - - Args: - tool: Genkit ToolDefinition. - - Returns: - Bedrock-compatible tool specification. - """ - input_schema = tool.input_schema or {'type': 'object', 'properties': {}} - - return { - 'toolSpec': { - 'name': tool.name, - 'description': tool.description or '', - 'inputSchema': { - 'json': input_schema, - }, - }, - } - - -def to_bedrock_content(part: Part) -> dict[str, Any] | None: - """Convert a single Genkit Part to a Bedrock content block. - - Handles TextPart, ToolRequestPart, and ToolResponsePart. - MediaPart requires async URL fetching so is handled separately - in the model class. - - Args: - part: A Genkit Part. - - Returns: - Bedrock content block dict, or None if the part type needs - special handling (e.g. MediaPart). - """ - root = part.root if isinstance(part, Part) else part - - if isinstance(root, TextPart): - return {'text': root.text} - - if isinstance(root, ToolRequestPart): - tool_req = root.tool_request - return { - 'toolUse': { - 'toolUseId': tool_req.ref or '', - 'name': tool_req.name, - 'input': tool_req.input if isinstance(tool_req.input, dict) else {}, - }, - } - - if isinstance(root, ToolResponsePart): - tool_resp = root.tool_response - output = tool_resp.output - result_content = [{'text': output}] if isinstance(output, str) else [{'json': output}] - return { - 'toolResult': { - 'toolUseId': tool_resp.ref or '', - 'content': result_content, - }, - } - - return None - - -def from_bedrock_content(content_blocks: list[dict[str, Any]]) -> list[Part]: - """Convert Bedrock response content blocks to Genkit parts. - - Handles text, toolUse, and reasoningContent blocks. - - Args: - content_blocks: List of Bedrock content blocks from the API response. - - Returns: - List of Genkit Part objects. - """ - parts: list[Part] = [] - - for block in content_blocks: - if 'text' in block: - parts.append(Part(root=TextPart(text=block['text']))) - - if 'toolUse' in block: - tool_use = block['toolUse'] - parts.append( - Part( - root=ToolRequestPart( - tool_request=ToolRequest( - ref=tool_use.get('toolUseId', ''), - name=tool_use.get('name', ''), - input=tool_use.get('input', {}), - ) - ) - ) - ) - - if 'reasoningContent' in block: - reasoning = block['reasoningContent'] - if 'reasoningText' in reasoning: - reasoning_text = reasoning['reasoningText'] - if isinstance(reasoning_text, dict) and 'text' in reasoning_text: - parts.append(Part(root=TextPart(text=f'[Reasoning]\n{reasoning_text["text"]}\n[/Reasoning]\n'))) - elif isinstance(reasoning_text, str): - parts.append(Part(root=TextPart(text=f'[Reasoning]\n{reasoning_text}\n[/Reasoning]\n'))) - - return parts - - -def strip_markdown_fences(text: str) -> str: - r"""Strip markdown code fences from a JSON response. - - Models sometimes wrap JSON output in markdown fences like - ``\`\`\`json ... \`\`\``` even when instructed to output raw - JSON. This helper removes the fences. - - Args: - text: The response text, possibly wrapped in fences. - - Returns: - The text with markdown fences removed, or the original - text if no fences are found. - """ - stripped = text.strip() - match = re.match(r'^```(?:json)?\s*\n?(.*?)\n?\s*```$', stripped, re.DOTALL) - if match: - return match.group(1).strip() - return text - - -def maybe_strip_fences(request: GenerateRequest, parts: list[Part]) -> list[Part]: - """Strip markdown fences from text parts when JSON output is expected. - - Args: - request: The original generate request. - parts: The response content parts. - - Returns: - Parts with fences stripped from text if JSON was requested. - """ - if not request.output or request.output.format != 'json': - return parts - - cleaned: list[Part] = [] - changed = False - for part in parts: - if isinstance(part.root, TextPart) and part.root.text: - cleaned_text = strip_markdown_fences(part.root.text) - if cleaned_text != part.root.text: - cleaned.append(Part(root=TextPart(text=cleaned_text))) - changed = True - else: - cleaned.append(part) - else: - cleaned.append(part) - return cleaned if changed else parts - - -class StreamingFenceStripper: - r"""Strips markdown fences from streamed text chunks. - - Buffers the first few characters to detect an opening - ``\`\`\`json`` fence. If found, the prefix is discarded and - subsequent chunks have the closing ``\`\`\``` stripped eagerly. - - Usage:: - - stripper = StreamingFenceStripper(json_mode=True) - for raw_text in stream: - text = stripper.process(raw_text) - if text: - send_chunk(text) - leftover = stripper.flush() - if leftover: - send_chunk(leftover) - """ - - _BUF_LIMIT = 12 # max length of "```json\n" with whitespace - - def __init__(self, *, json_mode: bool) -> None: - """Initialize the stripper. - - Args: - json_mode: When True, detect and strip fences; otherwise pass through. - """ - self._json_mode = json_mode - self._buf = '' - self._fence_detected = False - self._checked = False - self._pending_nl = '' - - def _flush_buf(self) -> str: - """Flush the buffer, stripping the opening fence if found.""" - self._checked = True - match = re.match(r'^```(?:json)?\s*\n?', self._buf) - if match: - self._fence_detected = True - remainder = self._buf[match.end() :] - self._buf = '' - return remainder - text = self._buf - self._buf = '' - return text - - def process(self, text: str) -> str: - """Process a text chunk, returning cleaned text to emit. - - May return an empty string if the text is being buffered. - """ - if not self._json_mode: - return text - - if not self._checked: - self._buf += text - if len(self._buf) >= self._BUF_LIMIT or '\n' in self._buf: - text = self._flush_buf() - else: - return '' - - if self._fence_detected: - raw = text - text = re.sub(r'\n?```\s*$', '', text) - fence_stripped = text != raw - if fence_stripped: - # A closing fence was removed β€” drop the deferred newline - # that preceded it. - self._pending_nl = '' - return text - # Defer a trailing newline β€” it may precede a closing fence - # in the next chunk. - prefix = self._pending_nl - if text.endswith('\n'): - self._pending_nl = '\n' - text = prefix + text[:-1] - else: - self._pending_nl = '' - text = prefix + text - return text - - return text - - def flush(self) -> str: - """Flush any remaining buffered text. Call after the stream ends.""" - if self._buf and not self._checked: - return self._flush_buf() - return self._pending_nl - - -def parse_tool_call_args(args_str: str) -> dict[str, Any] | str: - """Parse tool call arguments from a JSON string. - - Gracefully handles invalid JSON by returning the raw string. - - Args: - args_str: JSON string of tool call arguments. - - Returns: - Parsed dict if valid JSON, otherwise the raw string. - """ - if not args_str: - return {} - try: - return json.loads(args_str) - except (json.JSONDecodeError, TypeError): - return args_str - - -def build_usage(usage_data: dict[str, Any]) -> GenerationUsage: - """Build GenerationUsage from Bedrock usage data. - - Args: - usage_data: Usage dict from the Bedrock API response. - - Returns: - GenerationUsage with token counts. - """ - return GenerationUsage( - input_tokens=usage_data.get('inputTokens', 0), - output_tokens=usage_data.get('outputTokens', 0), - total_tokens=usage_data.get('totalTokens', 0), - ) - - -def normalize_config(config: object) -> BedrockConfig: - """Normalize config to BedrockConfig. - - Handles dicts with camelCase keys, GenerationCommonConfig, and - BedrockConfig passthrough. - - Args: - config: Request configuration (dict, BedrockConfig, or GenerationCommonConfig). - - Returns: - Normalized BedrockConfig instance. - """ - if config is None: - return BedrockConfig() - - if isinstance(config, BedrockConfig): - return config - - if isinstance(config, GenerationCommonConfig): - max_tokens = int(config.max_output_tokens) if config.max_output_tokens is not None else None - return BedrockConfig( - temperature=config.temperature, - max_tokens=max_tokens, - top_p=config.top_p, - stop_sequences=config.stop_sequences, - ) - - if isinstance(config, dict): - mapped: dict[str, Any] = {} - key_map: dict[str, str] = { - 'maxOutputTokens': 'max_tokens', - 'maxTokens': 'max_tokens', - 'topP': 'top_p', - 'topK': 'top_k', - 'stopSequences': 'stop_sequences', - } - for key, value in config.items(): - str_key = str(key) - mapped_key = key_map.get(str_key, str_key) - mapped[mapped_key] = value - return BedrockConfig(**mapped) - - return BedrockConfig() - - -def build_json_instruction(request: GenerateRequest) -> str | None: - """Build a JSON output instruction for the system prompt. - - The Bedrock Converse API doesn't have native JSON mode. Instead, - we inject instructions into the system prompt to guide the model. - - Args: - request: The generation request. - - Returns: - JSON instruction string if JSON output is requested, None otherwise. - """ - if not request.output: - return None - - if request.output.format != 'json': - return None - - instruction_parts = [ - 'IMPORTANT: You MUST respond with valid JSON only.', - 'Do not include any text before or after the JSON.', - 'Do not wrap the JSON in markdown code blocks.', - ] - - if request.output.schema: - schema_str = json.dumps(request.output.schema, indent=2) - instruction_parts.append(f'Your response MUST conform to this JSON schema:\n{schema_str}') - - return '\n'.join(instruction_parts) - - -def convert_media_data_uri(media: Media) -> tuple[bytes, str, bool]: - """Convert a data URI media to raw bytes and format. - - Only handles data: URIs. For HTTP URLs, returns empty values - and is_data_uri=False, signaling that async fetching is needed. - - Args: - media: Genkit Media object. - - Returns: - Tuple of (media_bytes, format_str, is_data_uri). - """ - url = media.url - content_type = media.content_type or '' - - if not url.startswith('data:'): - return b'', '', False - - format_str = _format_from_content_type(content_type, url) - - parts = url.split(',', 1) - if len(parts) == 2: - media_bytes = base64.b64decode(parts[1]) - return media_bytes, format_str, True - - return b'', format_str, False - - -def is_image_media(content_type: str, url: str) -> bool: - """Determine if media is an image (vs video). - - Args: - content_type: MIME type. - url: Media URL. - - Returns: - True if the media is an image. - """ - if content_type.startswith('image/'): - return True - if content_type.startswith('video/'): - return False - return any(ext in url.lower() for ext in ['.jpg', '.jpeg', '.png', '.gif', '.webp']) - - -def build_media_block(media_bytes: bytes, format_str: str, is_image: bool) -> dict[str, Any]: - """Build a Bedrock media content block from raw bytes. - - Args: - media_bytes: The raw media bytes. - format_str: Format string (e.g., 'jpeg', 'png', 'mp4'). - is_image: True for image, False for video. - - Returns: - Bedrock-compatible media content block. - """ - media_type = 'image' if is_image else 'video' - return { - media_type: { - 'format': format_str, - 'source': {'bytes': media_bytes}, - }, - } - - -def get_effective_model_id( - model_id: str, - *, - bearer_token: str | None = None, - aws_region: str | None = None, -) -> str: - """Get the effective model ID, adding inference profile prefix if needed. - - When using API key authentication (bearer token), AWS Bedrock requires - inference profile IDs with regional prefixes (us., eu., apac.) instead - of direct model IDs for supported providers. - - Args: - model_id: The base model ID. - bearer_token: AWS bearer token for API key auth (from env). - aws_region: AWS region string (from env). - - Returns: - The model ID to use for the API call. - """ - if model_id.startswith(INFERENCE_PROFILE_PREFIXES): - return model_id - - if not bearer_token: - return model_id - - if not model_id.startswith(INFERENCE_PROFILE_SUPPORTED_PROVIDERS): - return model_id - - if not aws_region: - return model_id - - region_lower = aws_region.lower() - if region_lower.startswith(('us-', 'us_')): - prefix = 'us' - elif region_lower.startswith(('eu-', 'eu_')): - prefix = 'eu' - elif region_lower.startswith(('ap-', 'ap_', 'cn-', 'cn_', 'me-', 'me_', 'af-', 'af_', 'sa-', 'sa_')): - prefix = 'apac' - else: - prefix = 'us' - - return f'{prefix}.{model_id}' - - -def _format_from_content_type(content_type: str, url: str) -> str: - """Extract media format string from content type or URL. - - Args: - content_type: MIME type (e.g., 'image/png'). - url: Media URL as fallback for format detection. - - Returns: - Format string (e.g., 'jpeg', 'png', 'mp4'). - """ - if content_type: - return content_type.split('/')[-1] - - for ext in ['jpeg', 'jpg', 'png', 'gif', 'webp', 'mp4', 'webm', 'mov']: - if f'.{ext}' in url.lower(): - return ext if ext != 'jpg' else 'jpeg' - - return 'jpeg' diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/model.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/model.py deleted file mode 100644 index 52449f751a..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/model.py +++ /dev/null @@ -1,932 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS Bedrock model implementation for Genkit. - -This module implements the model interface for AWS Bedrock using the Converse API. - -See: -- Converse API: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html -- Boto3 Reference: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html - -Key Features ------------- -- Chat completions using Converse API -- Tool/function calling support -- Streaming responses via ConverseStream -- Multimodal inputs (images, video for supported models) -- Reasoning content extraction (for DeepSeek-R1, etc.) - -Implementation Notes & Edge Cases ---------------------------------- - -**Media URL Fetching (Bedrock-Specific Requirement)** - -Unlike other AI providers (Anthropic, OpenAI, Google GenAI, xAI) that accept media URLs -directly in their APIs and fetch the content server-side, AWS Bedrock's Converse API -**only accepts inline bytes**. The API does not support URL references. - -This means we must fetch media content client-side before sending to Bedrock:: - - # Other providers (e.g., Anthropic): - {'type': 'url', 'url': 'https://example.com/image.jpg'} # API fetches it - - # AWS Bedrock requires: - {'image': {'format': 'jpeg', 'source': {'bytes': b'...actual bytes...'}}} - -We use ``httpx.AsyncClient`` for true async HTTP requests. This approach: - -- Uses httpx which is already a genkit core dependency -- True async I/O (no thread pool needed) -- Doesn't block the event loop during network I/O -- Includes a 30-second timeout for fetch operations -- Supports both images and videos -- Better error handling with ``HTTPStatusError`` for HTTP errors - -**User-Agent Header Requirement** - -Some servers (notably Wikipedia/Wikimedia) block requests without a proper ``User-Agent`` -header, returning HTTP 403 Forbidden. We include a standard User-Agent header to ensure -compatibility:: - - headers = { - 'User-Agent': 'Genkit/1.0 (https://github.com/firebase/genkit; Python httpx)', - 'Accept': 'image/*,video/*,*/*', - } - -**Base64 Data URL Handling** - -Data URLs (``data:image/png;base64,...``) are handled inline without network requests. -The base64 payload is decoded directly to bytes. - -**JSON Output Mode (Prompt Engineering)** - -The Bedrock Converse API doesn't have a native JSON mode like OpenAI's ``response_format``. -When JSON output is requested via ``request.output.format == 'json'``, we inject -instructions into the system prompt to guide the model:: - - IMPORTANT: You MUST respond with valid JSON only. - Do not include any text before or after the JSON. - Do not wrap the JSON in markdown code blocks. - Your response MUST conform to this JSON schema: - {...schema...} - -This prompt engineering approach works across all Bedrock models but is not as -reliable as native JSON mode. Models may occasionally include extra text. - -**Automatic Inference Profile Conversion** - -When using API key authentication (``AWS_BEARER_TOKEN_BEDROCK``), AWS Bedrock -requires inference profile IDs with regional prefixes instead of direct model IDs. -The plugin automatically detects API key auth and adds the appropriate regional -prefix (``us.``, ``eu.``, ``apac.``) based on ``AWS_REGION``:: - - # User specifies direct model ID - model = 'amazon-bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0' - - # With API key auth and AWS_REGION=us-east-1, automatically converts to: - modelId = 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - -This means users don't need to manually use ``inference_profile()`` - the plugin -handles the conversion transparently. If the model ID already has a regional -prefix, no conversion is performed. - -**Logging & Error Handling** - -All API calls and media fetches are logged for debugging: - -- ``logger.debug()`` for successful operations (request start, media fetch) -- ``logger.exception()`` for failures (API errors, fetch failures) - -Exceptions from boto3 are logged with full context before being re-raised, -ensuring errors are visible in logs even when caught by upstream code. - -**Streaming Implementation** - -Streaming uses the ``converse_stream`` API. Tool use deltas are accumulated -across multiple events and assembled into complete tool requests at the end -of the stream. -""" - -import base64 -import json -import os -from typing import Any - -import httpx - -from genkit.ai import ActionRunContext -from genkit.core.http_client import get_cached_client -from genkit.core.logging import get_logger -from genkit.plugins.amazon_bedrock.models.converters import ( - StreamingFenceStripper, - maybe_strip_fences, -) -from genkit.plugins.amazon_bedrock.typing import BedrockConfig -from genkit.types import ( - FinishReason, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationCommonConfig, - GenerationUsage, - Media, - MediaPart, - Message, - Part, - Role, - TextPart, - ToolDefinition, - ToolRequest, - ToolRequestPart, - ToolResponsePart, -) - -# Logger for this module -logger = get_logger(__name__) - -# Regional prefixes for inference profiles (us, eu, apac) -_INFERENCE_PROFILE_PREFIXES = ('us.', 'eu.', 'apac.') - -# Model provider prefixes that support cross-region inference profiles. -# These providers can use regional prefixes (us., eu., apac.) with API key auth. -# Other providers (ai21., stability.) require direct model IDs only. -_INFERENCE_PROFILE_SUPPORTED_PROVIDERS = ( - 'anthropic.', # Claude models - 'amazon.', # Nova, Titan models - 'meta.', # Llama models - 'mistral.', # Mistral models - 'cohere.', # Command R models - 'deepseek.', # DeepSeek models -) - -# Mapping from Bedrock stop reasons to Genkit finish reasons -FINISH_REASON_MAP: dict[str, FinishReason] = { - 'end_turn': FinishReason.STOP, - 'stop_sequence': FinishReason.STOP, - 'max_tokens': FinishReason.LENGTH, - 'tool_use': FinishReason.STOP, - 'content_filtered': FinishReason.BLOCKED, - 'guardrail_intervened': FinishReason.BLOCKED, -} - - -class BedrockModel: - """AWS Bedrock model for chat completions using the Converse API. - - This class handles the conversion between Genkit's message format - and the AWS Bedrock Converse API format. - - Attributes: - model_id: The Bedrock model ID (e.g., 'anthropic.claude-sonnet-4-5-20250929-v1:0'). - client: boto3 bedrock-runtime client instance. - """ - - def __init__( - self, - model_id: str, - client: Any, # noqa: ANN401 - ) -> None: - """Initialize the Bedrock model. - - Args: - model_id: Bedrock model ID (e.g., 'anthropic.claude-sonnet-4-5-20250929-v1:0'). - client: boto3 bedrock-runtime client instance. - """ - self.model_id = model_id - self.client = client - - def _get_effective_model_id(self) -> str: - """Get the effective model ID, adding inference profile prefix if needed. - - When using API key authentication (AWS_BEARER_TOKEN_BEDROCK), AWS Bedrock - requires inference profile IDs with regional prefixes (us., eu., apac.) - instead of direct model IDs for supported providers. - - **Important**: Not all model providers support cross-region inference profiles. - Only certain providers (Anthropic, Amazon, Meta, Mistral, Cohere, DeepSeek) - support the regional prefix. Other providers (AI21, Stability) require - direct model IDs and will fail if a regional prefix is added. - - This method automatically detects if API keys are being used and adds - the appropriate regional prefix based on AWS_REGION if: - 1. The model ID doesn't already have a prefix - 2. The model provider supports inference profiles - - Returns: - The model ID to use for the API call. - """ - # Check if already has an inference profile prefix - if self.model_id.startswith(_INFERENCE_PROFILE_PREFIXES): - return self.model_id - - # Check if using API key authentication - if 'AWS_BEARER_TOKEN_BEDROCK' not in os.environ: - return self.model_id - - # Check if this model provider supports inference profiles - if not self.model_id.startswith(_INFERENCE_PROFILE_SUPPORTED_PROVIDERS): - logger.debug( - 'Model provider does not support inference profiles, using direct model ID', - model_id=self.model_id, - ) - return self.model_id - - # API key auth requires inference profiles - add regional prefix - region = os.environ.get('AWS_REGION') or os.environ.get('AWS_DEFAULT_REGION') - if not region: - logger.warning('API key auth detected but AWS_REGION not set. Using direct model ID which may fail.') - return self.model_id - - # Determine regional prefix - region_lower = region.lower() - if region_lower.startswith('us-') or region_lower.startswith('us_'): - prefix = 'us' - elif region_lower.startswith('eu-') or region_lower.startswith('eu_'): - prefix = 'eu' - elif region_lower.startswith(('ap-', 'ap_', 'cn-', 'cn_', 'me-', 'me_', 'af-', 'af_', 'sa-', 'sa_')): - prefix = 'apac' - else: - prefix = 'us' # Default to US for unknown regions - - effective_id = f'{prefix}.{self.model_id}' - logger.debug( - 'Auto-converting model ID to inference profile', - original=self.model_id, - effective=effective_id, - ) - return effective_id - - async def generate( - self, - request: GenerateRequest, - ctx: ActionRunContext | None = None, - ) -> GenerateResponse: - """Generate a response from AWS Bedrock. - - Args: - request: The generation request containing messages and config. - ctx: Action run context for streaming support. - - Returns: - GenerateResponse with the model's output. - """ - config = self._normalize_config(request.config) - params = await self._build_request_body(request, config) - streaming = ctx is not None and ctx.is_streaming - - logger.debug( - 'Bedrock generate request', - model_id=self.model_id, - streaming=streaming, - ) - - try: - if streaming and ctx is not None: - return await self._generate_streaming(params, ctx, request) - - # Non-streaming request using Converse API - response = await self.client.converse(**params) - except Exception as e: - logger.exception( - 'Bedrock API call failed', - model_id=self.model_id, - error=str(e), - ) - raise - - # Extract the output message - output = response.get('output', {}) - message_data = output.get('message', {}) - - # Convert response to Genkit format - content = self._from_bedrock_content(message_data.get('content', [])) - content = maybe_strip_fences(request, content) - response_message = Message(role=Role.MODEL, content=content) - - # Build usage statistics - usage_data = response.get('usage', {}) - usage = GenerationUsage( - input_tokens=usage_data.get('inputTokens', 0), - output_tokens=usage_data.get('outputTokens', 0), - total_tokens=usage_data.get('totalTokens', 0), - ) - - stop_reason = response.get('stopReason', '') - finish_reason = FINISH_REASON_MAP.get(stop_reason, FinishReason.UNKNOWN) - - return GenerateResponse( - message=response_message, - usage=usage, - finish_reason=finish_reason, - request=request, - ) - - async def _generate_streaming( - self, - params: dict[str, Any], - ctx: ActionRunContext, - request: GenerateRequest, - ) -> GenerateResponse: - """Handle streaming generation using ConverseStream. - - Args: - params: Request parameters for the API. - ctx: Action run context for sending chunks. - request: Original generation request. - - Returns: - Final GenerateResponse after streaming completes. - """ - try: - # Use converse_stream for streaming - response = await self.client.converse_stream(**params) - except Exception as e: - logger.exception( - 'Bedrock streaming API call failed', - model_id=self.model_id, - error=str(e), - ) - raise - - accumulated_content: list[Part] = [] - accumulated_tool_uses: dict[str, dict[str, Any]] = {} - final_usage: GenerationUsage | None = None - stop_reason: str = '' - - json_mode = bool(request.output and request.output.format == 'json') - fence_stripper = StreamingFenceStripper(json_mode=json_mode) - - def _send_text(text: str) -> None: - """Send a text chunk and accumulate it.""" - if not text: - return - part = Part(root=TextPart(text=text)) - accumulated_content.append(part) - ctx.send_chunk(GenerateResponseChunk(role=Role.MODEL, content=[part], index=0)) - - # Process the event stream - stream = response.get('stream', []) - async for event in stream: - # Handle content block delta (text chunks) - if 'contentBlockDelta' in event: - delta = event['contentBlockDelta'].get('delta', {}) - if 'text' in delta: - _send_text(fence_stripper.process(delta['text'])) - # Handle tool use delta (input fragments only). - # contentBlockStart always fires first with name/toolUseId, - # so by the time we get here the entry should already exist. - if 'toolUse' in delta: - tool_use = delta['toolUse'] - block_idx = str(event['contentBlockDelta'].get('contentBlockIndex', 0)) - if block_idx not in accumulated_tool_uses: - # Defensive fallback β€” should not happen in practice - # because contentBlockStart always precedes deltas. - accumulated_tool_uses[block_idx] = { - 'toolUseId': tool_use.get('toolUseId', ''), - 'name': tool_use.get('name', ''), - 'input': '', - } - if 'input' in tool_use: - tool_input = tool_use['input'] - if isinstance(tool_input, str): - accumulated_tool_uses[block_idx]['input'] += tool_input - - # Handle content block start (for tool use) - if 'contentBlockStart' in event: - start = event['contentBlockStart'].get('start', {}) - if 'toolUse' in start: - tool_use = start['toolUse'] - block_index = event['contentBlockStart'].get('contentBlockIndex', 0) - accumulated_tool_uses[str(block_index)] = { - 'toolUseId': tool_use.get('toolUseId', ''), - 'name': tool_use.get('name', ''), - 'input': '', - } - - # Handle metadata (usage) - if 'metadata' in event: - metadata = event['metadata'] - usage_data = metadata.get('usage', {}) - final_usage = GenerationUsage( - input_tokens=usage_data.get('inputTokens', 0), - output_tokens=usage_data.get('outputTokens', 0), - total_tokens=usage_data.get('totalTokens', 0), - ) - - # Handle content block stop β€” emit tool request chunks - if 'contentBlockStop' in event: - block_idx = str(event['contentBlockStop'].get('contentBlockIndex', 0)) - if block_idx in accumulated_tool_uses: - tool_data = accumulated_tool_uses[block_idx] - try: - tool_input = json.loads(tool_data['input']) if tool_data['input'] else {} - except json.JSONDecodeError: - tool_input = tool_data['input'] - ctx.send_chunk( - GenerateResponseChunk( - role=Role.MODEL, - content=[ - Part( - root=ToolRequestPart( - tool_request=ToolRequest( - ref=tool_data['toolUseId'], - name=tool_data['name'], - input=tool_input, - ) - ) - ) - ], - index=0, - ) - ) - - # Handle message stop - if 'messageStop' in event: - stop_reason = event['messageStop'].get('stopReason', '') - - # Flush any remaining fence buffer (short streams). - _send_text(fence_stripper.flush()) - - # Add accumulated tool uses to content - for tool_data in accumulated_tool_uses.values(): - try: - tool_input = json.loads(tool_data['input']) if tool_data['input'] else {} - except json.JSONDecodeError: - tool_input = tool_data['input'] - - accumulated_content.append( - Part( - root=ToolRequestPart( - tool_request=ToolRequest( - ref=tool_data['toolUseId'], - name=tool_data['name'], - input=tool_input, - ) - ) - ) - ) - - finish_reason = FINISH_REASON_MAP.get(stop_reason, FinishReason.UNKNOWN) - - return GenerateResponse( - message=Message(role=Role.MODEL, content=accumulated_content), - usage=final_usage, - finish_reason=finish_reason, - request=request, - ) - - def _normalize_config(self, config: object) -> BedrockConfig: - """Normalize config to BedrockConfig. - - Args: - config: Request configuration (dict, BedrockConfig, or GenerationCommonConfig). - - Returns: - Normalized BedrockConfig instance. - """ - if config is None: - return BedrockConfig() - - if isinstance(config, BedrockConfig): - return config - - if isinstance(config, GenerationCommonConfig): - max_tokens = int(config.max_output_tokens) if config.max_output_tokens is not None else None - return BedrockConfig( - temperature=config.temperature, - max_tokens=max_tokens, - top_p=config.top_p, - stop_sequences=config.stop_sequences, - ) - - if isinstance(config, dict): - # Handle camelCase to snake_case mapping - mapped: dict[str, Any] = {} - key_map: dict[str, str] = { - 'maxOutputTokens': 'max_tokens', - 'maxTokens': 'max_tokens', - 'topP': 'top_p', - 'topK': 'top_k', - 'stopSequences': 'stop_sequences', - } - for key, value in config.items(): - str_key = str(key) - mapped_key = key_map.get(str_key, str_key) - mapped[mapped_key] = value - return BedrockConfig(**mapped) - - return BedrockConfig() - - async def _build_request_body( - self, - request: GenerateRequest, - config: BedrockConfig, - ) -> dict[str, Any]: - """Build the AWS Bedrock Converse API request body. - - Args: - request: The generation request. - config: Normalized configuration. - - Returns: - Dictionary suitable for client.converse(). - """ - # Separate system messages from conversation messages - system_messages, conversation_messages = self._separate_system_messages(request.messages) - - body: dict[str, Any] = { - 'modelId': self._get_effective_model_id(), - 'messages': await self._to_bedrock_messages(conversation_messages), - } - - # Handle JSON output format by injecting instructions into system prompt - # The Converse API doesn't have native JSON mode, so we use prompt engineering - json_instruction = self._build_json_instruction(request) - if json_instruction: - system_messages.append(json_instruction) - - # Add system prompt if present - if system_messages: - body['system'] = [{'text': msg} for msg in system_messages] - - # Build inference config - inference_config: dict[str, Any] = {} - - if config.max_tokens is not None: - inference_config['maxTokens'] = config.max_tokens - elif config.max_output_tokens is not None: - inference_config['maxTokens'] = config.max_output_tokens - - if config.temperature is not None: - inference_config['temperature'] = config.temperature - if config.top_p is not None: - inference_config['topP'] = config.top_p - if config.stop_sequences is not None: - inference_config['stopSequences'] = config.stop_sequences - - if inference_config: - body['inferenceConfig'] = inference_config - - # Handle tools - if request.tools: - body['toolConfig'] = { - 'tools': [self._to_bedrock_tool(t) for t in request.tools], - } - - return body - - def _build_json_instruction(self, request: GenerateRequest) -> str | None: - """Build a JSON output instruction based on request.output configuration. - - The Bedrock Converse API doesn't have native JSON mode like OpenAI's response_format. - Instead, we inject instructions into the system prompt to ensure JSON output. - - Args: - request: The generation request. - - Returns: - JSON instruction string if JSON output is requested, None otherwise. - """ - if not request.output: - return None - - output_format = request.output.format - schema = request.output.schema - - if output_format != 'json': - return None - - # Build instruction for JSON output - instruction_parts = [ - 'IMPORTANT: You MUST respond with valid JSON only.', - 'Do not include any text before or after the JSON.', - 'Do not wrap the JSON in markdown code blocks.', - ] - - if schema: - # Include the schema in the instruction - schema_str = json.dumps(schema, indent=2) - instruction_parts.append(f'Your response MUST conform to this JSON schema:\n{schema_str}') - - return '\n'.join(instruction_parts) - - def _separate_system_messages( - self, - messages: list[Message], - ) -> tuple[list[str], list[Message]]: - """Separate system messages from conversation messages. - - The Converse API requires system messages to be passed separately. - - Args: - messages: List of Genkit messages. - - Returns: - Tuple of (system_texts, conversation_messages). - """ - system_texts: list[str] = [] - conversation_messages: list[Message] = [] - - for msg in messages: - if msg.role == Role.SYSTEM or (isinstance(msg.role, str) and msg.role.lower() == 'system'): - # Extract text from system message - text_parts: list[str] = [] - for part in msg.content: - root = part.root if isinstance(part, Part) else part - if isinstance(root, TextPart): - text_parts.append(root.text) - if text_parts: - system_texts.append(''.join(text_parts)) - else: - conversation_messages.append(msg) - - return system_texts, conversation_messages - - def _to_bedrock_tool(self, tool: ToolDefinition) -> dict[str, Any]: - """Convert a Genkit tool definition to Bedrock format. - - Args: - tool: Genkit ToolDefinition. - - Returns: - Bedrock-compatible tool specification. - """ - input_schema = tool.input_schema or {'type': 'object', 'properties': {}} - - return { - 'toolSpec': { - 'name': tool.name, - 'description': tool.description or '', - 'inputSchema': { - 'json': input_schema, - }, - }, - } - - async def _to_bedrock_messages( - self, - messages: list[Message], - ) -> list[dict[str, Any]]: - """Convert Genkit messages to Bedrock Converse API message format. - - Args: - messages: List of Genkit messages (excluding system messages). - - Returns: - List of Bedrock-compatible message dictionaries. - """ - bedrock_msgs: list[dict[str, Any]] = [] - - for msg in messages: - role = self._to_bedrock_role(msg.role) - content: list[dict[str, Any]] = [] - - for part in msg.content: - root = part.root if isinstance(part, Part) else part - - if isinstance(root, TextPart): - content.append({'text': root.text}) - - elif isinstance(root, MediaPart): - media = root.media - content.append(await self._convert_media_to_bedrock(media)) - - elif isinstance(root, ToolRequestPart): - # Tool use from assistant - tool_req = root.tool_request - content.append({ - 'toolUse': { - 'toolUseId': tool_req.ref or '', - 'name': tool_req.name, - 'input': tool_req.input if isinstance(tool_req.input, dict) else {}, - }, - }) - - elif isinstance(root, ToolResponsePart): - # Tool result from user - tool_resp = root.tool_response - output = tool_resp.output - result_content = [{'text': output}] if isinstance(output, str) else [{'json': output}] - - content.append({ - 'toolResult': { - 'toolUseId': tool_resp.ref or '', - 'content': result_content, - }, - }) - - if content: - bedrock_msgs.append({ - 'role': role, - 'content': content, - }) - - return bedrock_msgs - - async def _convert_media_to_bedrock(self, media: Media) -> dict[str, Any]: - """Convert Genkit Media to Bedrock image/video format. - - Args: - media: Genkit Media object. - - Returns: - Bedrock-compatible media content block. - - Raises: - ValueError: If the media URL cannot be fetched or is invalid. - """ - url = media.url - content_type = media.content_type or '' - - # Determine if this is an image or video - is_image = content_type.startswith('image/') or any( - ext in url.lower() for ext in ['.jpg', '.jpeg', '.png', '.gif', '.webp'] - ) - is_video = content_type.startswith('video/') or any(ext in url.lower() for ext in ['.mp4', '.webm', '.mov']) - - # Extract format from content type or URL - if content_type: - format_str = content_type.split('/')[-1] - else: - # Guess from URL extension - for ext in ['jpeg', 'jpg', 'png', 'gif', 'webp', 'mp4', 'webm', 'mov']: - if f'.{ext}' in url.lower(): - format_str = ext if ext != 'jpg' else 'jpeg' - break - else: - format_str = 'jpeg' # Default - - # Handle base64 data URLs - if url.startswith('data:'): - # Parse data URL: data:image/png;base64, - parts = url.split(',', 1) - if len(parts) == 2: - media_bytes = base64.b64decode(parts[1]) - if is_image or not is_video: - return { - 'image': { - 'format': format_str, - 'source': {'bytes': media_bytes}, - }, - } - else: - return { - 'video': { - 'format': format_str, - 'source': {'bytes': media_bytes}, - }, - } - - # For regular URLs, fetch the content asynchronously - # Bedrock doesn't support URL sources directly, we must provide bytes - media_bytes, format_str = await self._fetch_media_from_url(url, format_str) - - if is_image or not is_video: - return { - 'image': { - 'format': format_str, - 'source': {'bytes': media_bytes}, - }, - } - else: - return { - 'video': { - 'format': format_str, - 'source': {'bytes': media_bytes}, - }, - } - - # TODO(#4360): Replace with downloadRequestMedia middleware (G15 parity). - # User-Agent is required because many servers (e.g. Wikipedia) return - # 403 Forbidden for the default httpx user-agent string. - _DOWNLOAD_HEADERS: dict[str, str] = { - 'User-Agent': 'Genkit/1.0 (https://github.com/firebase/genkit; genkit@google.com)', - 'Accept': 'image/*,video/*,*/*', - } - - async def _fetch_media_from_url(self, url: str, default_format: str) -> tuple[bytes, str]: - """Fetch media content from a URL asynchronously. - - Args: - url: The URL to fetch media from. - default_format: Default format to use if not detected from response. - - Returns: - Tuple of (media_bytes, format_string). - - Raises: - ValueError: If the URL cannot be fetched. - """ - logger.debug('Fetching media from URL', url=url[:100]) - - try: - client = get_cached_client( - cache_key='amazon-bedrock/media-fetch', - headers=self._DOWNLOAD_HEADERS, - timeout=30.0, - follow_redirects=True, - ) - response = await client.get(url) - response.raise_for_status() - - media_bytes = response.content - format_str = default_format - - # Update format from response content-type if available - resp_content_type = response.headers.get('content-type', '') - if resp_content_type and '/' in resp_content_type: - format_str = resp_content_type.split('/')[-1].split(';')[0] - if format_str == 'jpg': - format_str = 'jpeg' - - logger.debug('Fetched media', size=len(media_bytes), format=format_str) - return media_bytes, format_str - except httpx.HTTPStatusError as e: - logger.exception('HTTP error fetching media URL', url=url[:100], status=e.response.status_code) - raise ValueError(f'HTTP {e.response.status_code} fetching media from URL: {url[:100]}...') from e - except Exception as e: - logger.exception('Failed to fetch media URL', url=url[:100], error=str(e)) - raise ValueError(f'Failed to fetch media from URL: {url[:100]}... Error: {e}') from e - - def _to_bedrock_role(self, role: Role | str) -> str: - """Convert Genkit role to Bedrock role. - - Args: - role: Genkit message role. - - Returns: - Bedrock role string ('user' or 'assistant'). - """ - if isinstance(role, str): - str_role_map = { - 'user': 'user', - 'model': 'assistant', - 'assistant': 'assistant', - 'tool': 'user', # Tool responses come from user role - } - return str_role_map.get(role.lower(), 'user') - - role_map = { - Role.USER: 'user', - Role.MODEL: 'assistant', - Role.TOOL: 'user', # Tool responses come from user role - } - return role_map.get(role, 'user') - - def _from_bedrock_content(self, content_blocks: list[dict[str, Any]]) -> list[Part]: - """Convert Bedrock response content to Genkit parts. - - Args: - content_blocks: List of Bedrock content blocks. - - Returns: - List of Genkit Part objects. - """ - parts: list[Part] = [] - - for block in content_blocks: - # Handle text content - if 'text' in block: - parts.append(Part(root=TextPart(text=block['text']))) - - # Handle tool use - if 'toolUse' in block: - tool_use = block['toolUse'] - parts.append( - Part( - root=ToolRequestPart( - tool_request=ToolRequest( - ref=tool_use.get('toolUseId', ''), - name=tool_use.get('name', ''), - input=tool_use.get('input', {}), - ) - ) - ) - ) - - # Handle reasoning content (DeepSeek-R1, etc.) - if 'reasoningContent' in block: - reasoning = block['reasoningContent'] - if 'reasoningText' in reasoning: - # Include reasoning as a text part with prefix - reasoning_text = reasoning['reasoningText'] - if isinstance(reasoning_text, dict) and 'text' in reasoning_text: - parts.append(Part(root=TextPart(text=f'[Reasoning]\n{reasoning_text["text"]}\n[/Reasoning]\n'))) - elif isinstance(reasoning_text, str): - parts.append(Part(root=TextPart(text=f'[Reasoning]\n{reasoning_text}\n[/Reasoning]\n'))) - - return parts diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/model_info.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/model_info.py deleted file mode 100644 index 7d8a50c6da..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/models/model_info.py +++ /dev/null @@ -1,605 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS Bedrock model info. - -This module defines the supported models for the AWS Bedrock plugin. -AWS Bedrock provides access to foundation models from multiple providers -through a unified API. - -See: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html - -Supported Model Categories -========================== -+------------------+--------------------------------------------------+ -| Provider | Models | -+------------------+--------------------------------------------------+ -| Amazon | Nova Pro, Nova Lite, Nova Micro, Titan | -| Anthropic | Claude Opus, Sonnet, Haiku (4.x, 3.x series) | -| AI21 Labs | Jamba 1.5 Large, Mini | -| Cohere | Command R, Command R+, Embed, Rerank | -| DeepSeek | R1, V3.1 | -| Google | Gemma 3 (4B, 12B, 27B) | -| Meta | Llama 3.x, Llama 4 Maverick/Scout | -| MiniMax | M2 | -| Mistral AI | Large 3, Pixtral, Magistral, Ministral | -| Moonshot AI | Kimi K2 Thinking | -| NVIDIA | Nemotron Nano | -| OpenAI | GPT-OSS 120B, 20B | -| Qwen | Qwen3 32B, 235B, Coder, VL | -| Writer | Palmyra X4, X5 | -+------------------+--------------------------------------------------+ - -Model ID Format: - Bedrock model IDs follow the pattern: `{provider}.{model-name}-{version}` - Examples: - - anthropic.claude-sonnet-4-5-20250929-v1:0 - - meta.llama3-3-70b-instruct-v1:0 - - amazon.nova-pro-v1:0 - -Cross-Region Inference: - For cross-region inference profiles, use the region-prefixed ID: - - us.anthropic.claude-sonnet-4-5-20250929-v1:0 - - us.deepseek.r1-v1:0 - -Trademark Notice: - Model names are trademarks of their respective owners. This plugin is - developed independently and is not affiliated with or endorsed by the - model providers. -""" - -from genkit.types import ModelInfo, Supports - -# Model capability definitions -MULTIMODAL_MODEL_SUPPORTS = Supports( - multiturn=True, - media=True, - tools=True, - system_role=True, - output=['text', 'json'], -) - -TEXT_ONLY_MODEL_SUPPORTS = Supports( - multiturn=True, - media=False, - tools=True, - system_role=True, - output=['text', 'json'], -) - -TEXT_ONLY_NO_TOOLS_SUPPORTS = Supports( - multiturn=True, - media=False, - tools=False, - system_role=True, - output=['text'], -) - -# Reasoning models typically output text only -REASONING_MODEL_SUPPORTS = Supports( - multiturn=True, - media=False, - tools=False, - system_role=True, - output=['text'], -) - -# Anthropic Claude model supports -CLAUDE_MODEL_SUPPORTS = Supports( - multiturn=True, - media=True, - tools=True, - system_role=True, - output=['text', 'json'], -) - -# Amazon Nova model supports (text, image, video input) -NOVA_MODEL_SUPPORTS = Supports( - multiturn=True, - media=True, - tools=True, - system_role=True, - output=['text', 'json'], -) - -NOVA_MICRO_SUPPORTS = Supports( - multiturn=True, - media=False, # Micro is text-only - tools=True, - system_role=True, - output=['text', 'json'], -) - -# Meta Llama model supports -LLAMA_TEXT_SUPPORTS = Supports( - multiturn=True, - media=False, - tools=True, - system_role=True, - output=['text', 'json'], -) - -LLAMA_MULTIMODAL_SUPPORTS = Supports( - multiturn=True, - media=True, - tools=True, - system_role=True, - output=['text', 'json'], -) - -# Mistral model supports -MISTRAL_MODEL_SUPPORTS = Supports( - multiturn=True, - media=False, - tools=True, - system_role=True, - output=['text', 'json'], -) - -MISTRAL_MULTIMODAL_SUPPORTS = Supports( - multiturn=True, - media=True, - tools=True, - system_role=True, - output=['text', 'json'], -) - -# Cohere model supports -COHERE_MODEL_SUPPORTS = Supports( - multiturn=True, - media=False, - tools=True, - system_role=True, - output=['text', 'json'], -) - -# DeepSeek model supports (reasoning) -DEEPSEEK_MODEL_SUPPORTS = Supports( - multiturn=True, - media=False, - tools=False, # DeepSeek-R1 doesn't support tools in Bedrock - system_role=True, - output=['text'], -) - -SUPPORTED_BEDROCK_MODELS: dict[str, ModelInfo] = { - # ========================================================================= - # Amazon Nova Models - # See: https://docs.aws.amazon.com/nova/latest/userguide/what-is-nova.html - # ========================================================================= - 'amazon.nova-pro-v1:0': ModelInfo( - label='Amazon Nova Pro', - versions=['amazon.nova-pro-v1:0'], - supports=NOVA_MODEL_SUPPORTS, - ), - 'amazon.nova-lite-v1:0': ModelInfo( - label='Amazon Nova Lite', - versions=['amazon.nova-lite-v1:0'], - supports=NOVA_MODEL_SUPPORTS, - ), - 'amazon.nova-micro-v1:0': ModelInfo( - label='Amazon Nova Micro', - versions=['amazon.nova-micro-v1:0'], - supports=NOVA_MICRO_SUPPORTS, - ), - 'amazon.nova-premier-v1:0': ModelInfo( - label='Amazon Nova Premier', - versions=['amazon.nova-premier-v1:0'], - supports=NOVA_MODEL_SUPPORTS, - ), - 'amazon.nova-2-lite-v1:0': ModelInfo( - label='Amazon Nova 2 Lite', - versions=['amazon.nova-2-lite-v1:0'], - supports=NOVA_MODEL_SUPPORTS, - ), - # ========================================================================= - # Anthropic Claude Models - # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html - # ========================================================================= - 'anthropic.claude-sonnet-4-5-20250929-v1:0': ModelInfo( - label='Claude Sonnet 4.5', - versions=['anthropic.claude-sonnet-4-5-20250929-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - 'anthropic.claude-sonnet-4-20250514-v1:0': ModelInfo( - label='Claude Sonnet 4', - versions=['anthropic.claude-sonnet-4-20250514-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - 'anthropic.claude-opus-4-5-20251101-v1:0': ModelInfo( - label='Claude Opus 4.5', - versions=['anthropic.claude-opus-4-5-20251101-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - 'anthropic.claude-opus-4-1-20250805-v1:0': ModelInfo( - label='Claude Opus 4.1', - versions=['anthropic.claude-opus-4-1-20250805-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - 'anthropic.claude-opus-4-6-20260205-v1:0': ModelInfo( - label='Claude Opus 4.6', - versions=['anthropic.claude-opus-4-6-20260205-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - 'anthropic.claude-haiku-4-5-20251001-v1:0': ModelInfo( - label='Claude Haiku 4.5', - versions=['anthropic.claude-haiku-4-5-20251001-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - 'anthropic.claude-3-5-haiku-20241022-v1:0': ModelInfo( - label='Claude 3.5 Haiku', - versions=['anthropic.claude-3-5-haiku-20241022-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - 'anthropic.claude-3-haiku-20240307-v1:0': ModelInfo( - label='Claude 3 Haiku', - versions=['anthropic.claude-3-haiku-20240307-v1:0'], - supports=CLAUDE_MODEL_SUPPORTS, - ), - # ========================================================================= - # AI21 Labs Jamba Models - # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-jamba.html - # ========================================================================= - 'ai21.jamba-1-5-large-v1:0': ModelInfo( - label='Jamba 1.5 Large', - versions=['ai21.jamba-1-5-large-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'ai21.jamba-1-5-mini-v1:0': ModelInfo( - label='Jamba 1.5 Mini', - versions=['ai21.jamba-1-5-mini-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - # ========================================================================= - # Cohere Models - # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html - # ========================================================================= - 'cohere.command-r-plus-v1:0': ModelInfo( - label='Command R+', - versions=['cohere.command-r-plus-v1:0'], - supports=COHERE_MODEL_SUPPORTS, - ), - 'cohere.command-r-v1:0': ModelInfo( - label='Command R', - versions=['cohere.command-r-v1:0'], - supports=COHERE_MODEL_SUPPORTS, - ), - # ========================================================================= - # DeepSeek Models - # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-deepseek.html - # ========================================================================= - 'deepseek.r1-v1:0': ModelInfo( - label='DeepSeek R1', - versions=['deepseek.r1-v1:0'], - supports=DEEPSEEK_MODEL_SUPPORTS, - ), - 'deepseek.v3-v1:0': ModelInfo( - label='DeepSeek V3.1', - versions=['deepseek.v3-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - # ========================================================================= - # Google Gemma Models - # ========================================================================= - 'google.gemma-3-4b-it': ModelInfo( - label='Gemma 3 4B IT', - versions=['google.gemma-3-4b-it'], - supports=MULTIMODAL_MODEL_SUPPORTS, - ), - 'google.gemma-3-12b-it': ModelInfo( - label='Gemma 3 12B IT', - versions=['google.gemma-3-12b-it'], - supports=MULTIMODAL_MODEL_SUPPORTS, - ), - 'google.gemma-3-27b-it': ModelInfo( - label='Gemma 3 27B IT', - versions=['google.gemma-3-27b-it'], - supports=MULTIMODAL_MODEL_SUPPORTS, - ), - # ========================================================================= - # Meta Llama Models - # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html - # ========================================================================= - 'meta.llama3-3-70b-instruct-v1:0': ModelInfo( - label='Llama 3.3 70B Instruct', - versions=['meta.llama3-3-70b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama3-1-405b-instruct-v1:0': ModelInfo( - label='Llama 3.1 405B Instruct', - versions=['meta.llama3-1-405b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama3-1-70b-instruct-v1:0': ModelInfo( - label='Llama 3.1 70B Instruct', - versions=['meta.llama3-1-70b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama3-1-8b-instruct-v1:0': ModelInfo( - label='Llama 3.1 8B Instruct', - versions=['meta.llama3-1-8b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama3-70b-instruct-v1:0': ModelInfo( - label='Llama 3 70B Instruct', - versions=['meta.llama3-70b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama3-8b-instruct-v1:0': ModelInfo( - label='Llama 3 8B Instruct', - versions=['meta.llama3-8b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama3-2-90b-instruct-v1:0': ModelInfo( - label='Llama 3.2 90B Instruct', - versions=['meta.llama3-2-90b-instruct-v1:0'], - supports=LLAMA_MULTIMODAL_SUPPORTS, - ), - 'meta.llama3-2-11b-instruct-v1:0': ModelInfo( - label='Llama 3.2 11B Instruct', - versions=['meta.llama3-2-11b-instruct-v1:0'], - supports=LLAMA_MULTIMODAL_SUPPORTS, - ), - 'meta.llama3-2-3b-instruct-v1:0': ModelInfo( - label='Llama 3.2 3B Instruct', - versions=['meta.llama3-2-3b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama3-2-1b-instruct-v1:0': ModelInfo( - label='Llama 3.2 1B Instruct', - versions=['meta.llama3-2-1b-instruct-v1:0'], - supports=LLAMA_TEXT_SUPPORTS, - ), - 'meta.llama4-maverick-17b-instruct-v1:0': ModelInfo( - label='Llama 4 Maverick 17B Instruct', - versions=['meta.llama4-maverick-17b-instruct-v1:0'], - supports=LLAMA_MULTIMODAL_SUPPORTS, - ), - 'meta.llama4-scout-17b-instruct-v1:0': ModelInfo( - label='Llama 4 Scout 17B Instruct', - versions=['meta.llama4-scout-17b-instruct-v1:0'], - supports=LLAMA_MULTIMODAL_SUPPORTS, - ), - # ========================================================================= - # MiniMax Models - # ========================================================================= - 'minimax.minimax-m2': ModelInfo( - label='MiniMax M2', - versions=['minimax.minimax-m2'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - # ========================================================================= - # Mistral AI Models - # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html - # ========================================================================= - 'mistral.mistral-large-3-675b-instruct': ModelInfo( - label='Mistral Large 3', - versions=['mistral.mistral-large-3-675b-instruct'], - supports=MISTRAL_MULTIMODAL_SUPPORTS, - ), - 'mistral.mistral-large-2407-v1:0': ModelInfo( - label='Mistral Large (24.07)', - versions=['mistral.mistral-large-2407-v1:0'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - 'mistral.mistral-large-2402-v1:0': ModelInfo( - label='Mistral Large (24.02)', - versions=['mistral.mistral-large-2402-v1:0'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - 'mistral.mistral-small-2402-v1:0': ModelInfo( - label='Mistral Small (24.02)', - versions=['mistral.mistral-small-2402-v1:0'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - 'mistral.mistral-7b-instruct-v0:2': ModelInfo( - label='Mistral 7B Instruct', - versions=['mistral.mistral-7b-instruct-v0:2'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - 'mistral.mixtral-8x7b-instruct-v0:1': ModelInfo( - label='Mixtral 8x7B Instruct', - versions=['mistral.mixtral-8x7b-instruct-v0:1'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - 'mistral.pixtral-large-2502-v1:0': ModelInfo( - label='Pixtral Large (25.02)', - versions=['mistral.pixtral-large-2502-v1:0'], - supports=MISTRAL_MULTIMODAL_SUPPORTS, - ), - 'mistral.magistral-small-2509': ModelInfo( - label='Magistral Small', - versions=['mistral.magistral-small-2509'], - supports=MISTRAL_MULTIMODAL_SUPPORTS, - ), - 'mistral.ministral-3-3b-instruct': ModelInfo( - label='Ministral 3B', - versions=['mistral.ministral-3-3b-instruct'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - 'mistral.ministral-3-8b-instruct': ModelInfo( - label='Ministral 8B', - versions=['mistral.ministral-3-8b-instruct'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - 'mistral.ministral-3-14b-instruct': ModelInfo( - label='Ministral 14B', - versions=['mistral.ministral-3-14b-instruct'], - supports=MISTRAL_MODEL_SUPPORTS, - ), - # ========================================================================= - # Moonshot AI Models - # ========================================================================= - 'moonshot.kimi-k2-thinking': ModelInfo( - label='Kimi K2 Thinking', - versions=['moonshot.kimi-k2-thinking'], - supports=REASONING_MODEL_SUPPORTS, - ), - # ========================================================================= - # NVIDIA Models - # ========================================================================= - 'nvidia.nemotron-nano-9b-v2': ModelInfo( - label='Nemotron Nano 9B', - versions=['nvidia.nemotron-nano-9b-v2'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'nvidia.nemotron-nano-12b-v2': ModelInfo( - label='Nemotron Nano 12B VL', - versions=['nvidia.nemotron-nano-12b-v2'], - supports=MULTIMODAL_MODEL_SUPPORTS, - ), - # ========================================================================= - # OpenAI Models (GPT-OSS on Bedrock) - # ========================================================================= - 'openai.gpt-oss-120b-1:0': ModelInfo( - label='GPT-OSS 120B', - versions=['openai.gpt-oss-120b-1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'openai.gpt-oss-20b-1:0': ModelInfo( - label='GPT-OSS 20B', - versions=['openai.gpt-oss-20b-1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'openai.gpt-oss-safeguard-120b': ModelInfo( - label='GPT-OSS Safeguard 120B', - versions=['openai.gpt-oss-safeguard-120b'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'openai.gpt-oss-safeguard-20b': ModelInfo( - label='GPT-OSS Safeguard 20B', - versions=['openai.gpt-oss-safeguard-20b'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - # ========================================================================= - # Qwen Models - # ========================================================================= - 'qwen.qwen3-32b-v1:0': ModelInfo( - label='Qwen3 32B', - versions=['qwen.qwen3-32b-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'qwen.qwen3-235b-a22b-2507-v1:0': ModelInfo( - label='Qwen3 235B A22B', - versions=['qwen.qwen3-235b-a22b-2507-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'qwen.qwen3-coder-30b-a3b-v1:0': ModelInfo( - label='Qwen3 Coder 30B', - versions=['qwen.qwen3-coder-30b-a3b-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'qwen.qwen3-coder-480b-a35b-v1:0': ModelInfo( - label='Qwen3 Coder 480B', - versions=['qwen.qwen3-coder-480b-a35b-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'qwen.qwen3-next-80b-a3b': ModelInfo( - label='Qwen3 Next 80B', - versions=['qwen.qwen3-next-80b-a3b'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'qwen.qwen3-vl-235b-a22b': ModelInfo( - label='Qwen3 VL 235B', - versions=['qwen.qwen3-vl-235b-a22b'], - supports=MULTIMODAL_MODEL_SUPPORTS, - ), - # ========================================================================= - # Writer Models - # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-writer-palmyra.html - # ========================================================================= - 'writer.palmyra-x4-v1:0': ModelInfo( - label='Palmyra X4', - versions=['writer.palmyra-x4-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - 'writer.palmyra-x5-v1:0': ModelInfo( - label='Palmyra X5', - versions=['writer.palmyra-x5-v1:0'], - supports=TEXT_ONLY_MODEL_SUPPORTS, - ), - # ========================================================================= - # Amazon Titan Text Models - # ========================================================================= - 'amazon.titan-tg1-large': ModelInfo( - label='Titan Text Large', - versions=['amazon.titan-tg1-large'], - supports=TEXT_ONLY_NO_TOOLS_SUPPORTS, - ), -} - -# Embedding models -# See: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html -SUPPORTED_EMBEDDING_MODELS: dict[str, dict] = { - 'amazon.titan-embed-text-v2:0': { - 'label': 'Titan Text Embeddings V2', - 'dimensions': 1024, - 'supports': {'input': ['text']}, - }, - 'amazon.titan-embed-text-v1': { - 'label': 'Titan Embeddings G1 - Text', - 'dimensions': 1536, - 'supports': {'input': ['text']}, - }, - 'amazon.titan-embed-image-v1': { - 'label': 'Titan Multimodal Embeddings G1', - 'dimensions': 1024, - 'supports': {'input': ['text', 'image']}, - }, - 'amazon.nova-2-multimodal-embeddings-v1:0': { - 'label': 'Nova Multimodal Embeddings', - 'dimensions': 1024, - 'supports': {'input': ['text', 'image', 'audio', 'video']}, - }, - 'cohere.embed-english-v3': { - 'label': 'Cohere Embed English', - 'dimensions': 1024, - 'supports': {'input': ['text']}, - }, - 'cohere.embed-multilingual-v3': { - 'label': 'Cohere Embed Multilingual', - 'dimensions': 1024, - 'supports': {'input': ['text']}, - }, - 'cohere.embed-v4:0': { - 'label': 'Cohere Embed v4', - 'dimensions': 1024, - 'supports': {'input': ['text', 'image']}, - }, -} - - -def get_model_info(name: str) -> ModelInfo: - """Get model info for a given model name. - - For the full model catalog, see: - https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html - - Args: - name: The name of the model (e.g., 'anthropic.claude-sonnet-4-5-20250929-v1:0'). - - Returns: - ModelInfo for the model. - """ - if name in SUPPORTED_BEDROCK_MODELS: - return SUPPORTED_BEDROCK_MODELS[name] - - # Default info for unknown models - assume text-only capable - # This allows users to use any model dynamically - return ModelInfo( - label=f'Bedrock - {name}', - supports=TEXT_ONLY_MODEL_SUPPORTS, - ) diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/plugin.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/plugin.py deleted file mode 100644 index b831818b4f..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/plugin.py +++ /dev/null @@ -1,726 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS Bedrock plugin for Genkit. - -This plugin provides access to AWS Bedrock models through the Genkit framework. -AWS Bedrock is a fully managed service that provides access to foundation models -from multiple providers through a unified API. - -Documentation Links: - - AWS Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html - - Supported Models: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html - - Converse API: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html - - Boto3 Reference: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime.html - -The plugin supports: - - Chat completion models via Converse API - - Text embedding models (Titan, Cohere, Nova) - - Tool/function calling - - Streaming responses - - Multimodal inputs (images, video for supported models) - -Authentication: - The plugin uses the standard AWS credential chain: - 1. Environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) - 2. AWS credentials file (~/.aws/credentials) - 3. IAM role (for EC2, Lambda, ECS, etc.) - 4. AWS profile - -Example:: - - from genkit import Genkit - from genkit.plugins.amazon_bedrock import AmazonBedrock, bedrock_model - - ai = Genkit( - plugins=[AmazonBedrock(region='us-east-1')], - model=bedrock_model('anthropic.claude-sonnet-4-5-20250929-v1:0'), - ) - - response = await ai.generate(prompt='Tell me a joke.') - print(response.text) - -Trademark Notice: - This is a community plugin and is not officially supported by Amazon Web Services. - "Amazon", "AWS", "Amazon Bedrock", and related marks are trademarks of - Amazon.com, Inc. or its affiliates. -""" - -import json -import os -from typing import Any - -import aioboto3 -from botocore.config import Config - -from genkit.ai import Plugin -from genkit.blocks.embedding import EmbedderOptions, EmbedderSupports, embedder_action_metadata -from genkit.blocks.model import model_action_metadata -from genkit.core.action import Action, ActionMetadata -from genkit.core.logging import get_logger -from genkit.core.registry import ActionKind -from genkit.plugins.amazon_bedrock.models.model import BedrockModel -from genkit.plugins.amazon_bedrock.models.model_info import ( - SUPPORTED_BEDROCK_MODELS, - SUPPORTED_EMBEDDING_MODELS, - get_model_info, -) -from genkit.plugins.amazon_bedrock.typing import ( - AI21JambaConfig, - AmazonNovaConfig, - AnthropicConfig, - BedrockConfig, - CohereConfig, - DeepSeekConfig, - GoogleGemmaConfig, - MetaLlamaConfig, - MiniMaxConfig, - MistralConfig, - MoonshotConfig, - NvidiaConfig, - OpenAIConfig, - QwenConfig, - StabilityConfig, - TitanConfig, - WriterConfig, -) -from genkit.types import Embedding, EmbedRequest, EmbedResponse - -# Regional prefixes for inference profiles (must match model.py) -_INFERENCE_PROFILE_PREFIXES = ('us.', 'eu.', 'apac.') - - -def _strip_inference_profile_prefix(model_id: str) -> str: - """Strip the regional inference profile prefix from a model ID. - - Converts inference profile IDs (e.g., ``us.amazon.titan-embed-text-v2:0``) - back to base model IDs (``amazon.titan-embed-text-v2:0``). - - The InvokeModel API (used for embeddings) does NOT support inference profile - IDs β€” only Converse/ConverseStream do. This helper ensures we always pass - base model IDs to InvokeModel. - - Args: - model_id: A model ID, possibly with a regional prefix. - - Returns: - The base model ID with any regional prefix stripped. - """ - for prefix in _INFERENCE_PROFILE_PREFIXES: - if model_id.startswith(prefix): - return model_id[len(prefix) :] - return model_id - - -_MODEL_CONFIG_PREFIX_MAP: dict[str, type] = { - # Amazon models - 'amazon.nova': AmazonNovaConfig, - 'amazon.titan': TitanConfig, - # Anthropic Claude models - 'anthropic.claude': AnthropicConfig, - # AI21 Labs Jamba models - 'ai21.jamba': AI21JambaConfig, - # Cohere models - 'cohere.command': CohereConfig, - 'cohere.embed': CohereConfig, - # DeepSeek models - 'deepseek': DeepSeekConfig, - # Google Gemma models - 'google.gemma': GoogleGemmaConfig, - # Meta Llama models - 'meta.llama': MetaLlamaConfig, - # MiniMax models - 'minimax': MiniMaxConfig, - # Mistral AI models - 'mistral': MistralConfig, - # Moonshot AI models - 'moonshot': MoonshotConfig, - # NVIDIA models - 'nvidia': NvidiaConfig, - # OpenAI models (GPT-OSS on Bedrock) - 'openai': OpenAIConfig, - # Qwen models - 'qwen': QwenConfig, - # Writer models - 'writer': WriterConfig, - # Stability AI models - 'stability': StabilityConfig, -} -"""Mapping from model ID prefixes to their configuration classes.""" - - -def get_config_schema_for_model(model_id: str) -> type: - """Get the appropriate config schema for a model based on its ID. - - This function maps model IDs to their model-specific configuration classes, - enabling the DevUI to show relevant parameters for each model family. - - Handles both direct model IDs and cross-region inference profile IDs: - - Direct: 'anthropic.claude-sonnet-4-5-20250929-v1:0' - - Inference profile: 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - - Args: - model_id: The model ID or inference profile ID. - - Returns: - The appropriate config class for the model. Returns BedrockConfig as default. - """ - model_lower = model_id.lower() - - # Strip regional prefix for inference profiles (e.g., 'us.', 'eu.', 'ap.') - # Inference profile format: {region}.{provider}.{model} - if '.' in model_lower: - parts = model_lower.split('.', 1) - # Check if first part is a region code (2-3 chars like 'us', 'eu', 'ap') - if len(parts[0]) <= 3 and parts[0].isalpha(): - model_lower = parts[1] - - for prefix, config_class in _MODEL_CONFIG_PREFIX_MAP.items(): - if model_lower.startswith(prefix): - return config_class - - # Default: standard Bedrock config - return BedrockConfig - - -# Plugin name -AMAZON_BEDROCK_PLUGIN_NAME = 'amazon-bedrock' - -# Logger for this module -logger = get_logger(__name__) - - -def bedrock_name(model_id: str) -> str: - """Get fully qualified AWS Bedrock model name. - - Args: - model_id: The Bedrock model ID (e.g., 'anthropic.claude-sonnet-4-5-20250929-v1:0'). - - Returns: - Fully qualified model name (e.g., 'amazon-bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0'). - """ - return f'{AMAZON_BEDROCK_PLUGIN_NAME}/{model_id}' - - -class AmazonBedrock(Plugin): - """AWS Bedrock plugin for Genkit. - - This plugin provides access to AWS Bedrock models including: - - Amazon: Nova Pro, Nova Lite, Nova Micro, Titan - - Anthropic: Claude Opus, Sonnet, Haiku - - AI21 Labs: Jamba 1.5 - - Cohere: Command R, Command R+, Embed - - DeepSeek: R1, V3 - - Google: Gemma 3 - - Meta: Llama 3.x, Llama 4 - - MiniMax: M2 - - Mistral AI: Large 3, Pixtral, Ministral - - Moonshot AI: Kimi K2 - - NVIDIA: Nemotron - - OpenAI: GPT-OSS - - Qwen: Qwen3 - - Writer: Palmyra - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html - - Attributes: - name: Plugin name ('amazon-bedrock'). - """ - - name = AMAZON_BEDROCK_PLUGIN_NAME - - def __init__( - self, - region: str | None = None, - access_key_id: str | None = None, - secret_access_key: str | None = None, - session_token: str | None = None, - profile_name: str | None = None, - connect_timeout: int = 60, - read_timeout: int = 3600, - **boto_config: Any, # noqa: ANN401 - ) -> None: - """Initialize the AWS Bedrock plugin. - - Args: - region: AWS region (e.g., 'us-east-1'). Falls back to AWS_REGION env var. - access_key_id: AWS access key ID. Falls back to AWS_ACCESS_KEY_ID env var. - secret_access_key: AWS secret access key. Falls back to AWS_SECRET_ACCESS_KEY env var. - session_token: AWS session token for temporary credentials. - profile_name: AWS profile name to use from credentials file. - connect_timeout: Connection timeout in seconds. Default: 60. - read_timeout: Read timeout in seconds. Default: 3600 (1 hour for Nova models). - **boto_config: Additional parameters passed to boto3 Config. - - Example: - # Using environment variables (recommended): - plugin = AmazonBedrock(region='us-east-1') - - # Using explicit credentials: - plugin = AmazonBedrock( - region='us-east-1', - access_key_id='your-access-key', - secret_access_key='your-secret-key', - ) - - # Using AWS profile: - plugin = AmazonBedrock( - region='us-east-1', - profile_name='my-profile', - ) - """ - # Resolve region from environment - resolved_region = region or os.environ.get('AWS_REGION') or os.environ.get('AWS_DEFAULT_REGION') - - if not resolved_region: - raise ValueError('AWS region is required. Set AWS_REGION environment variable or pass region parameter.') - - # Build botocore config with timeouts - # Nova models have 60-minute inference timeout, so we set read_timeout high - self._botocore_config = Config( - connect_timeout=connect_timeout, - read_timeout=read_timeout, - **boto_config, - ) - - # Build session kwargs - session_kwargs: dict[str, Any] = { - 'region_name': resolved_region, - } - - if profile_name: - session_kwargs['profile_name'] = profile_name - - # Create aioboto3 session (sync β€” session creation is lightweight) - self._session = aioboto3.Session(**session_kwargs) - - # Client kwargs for bedrock-runtime - self._client_kwargs: dict[str, Any] = { - 'config': self._botocore_config, - } - - # Add explicit credentials if provided - if access_key_id and secret_access_key: - self._client_kwargs['aws_access_key_id'] = access_key_id - self._client_kwargs['aws_secret_access_key'] = secret_access_key - if session_token: - self._client_kwargs['aws_session_token'] = session_token - - # Async client β€” created in init(), closed in close() - self._client: Any = None # noqa: ANN401 - self._client_ctx: Any = None # noqa: ANN401 - self._region = resolved_region - - logger.debug( - 'Initialized AWS Bedrock plugin', - region=resolved_region, - ) - - async def init(self) -> list[Action]: - """Initialize plugin and register supported models. - - Creates the async bedrock-runtime client and registers all - supported models and embedders. - - Returns: - List of Action objects for supported models and embedders. - """ - # Create the async bedrock-runtime client. - # aioboto3 requires async context managers for client creation. - # We enter the context here and exit in close(). - self._client_ctx = self._session.client('bedrock-runtime', **self._client_kwargs) - self._client = await self._client_ctx.__aenter__() - - actions: list[Action] = [] - - # Register all supported models from predefined list - for model_id in SUPPORTED_BEDROCK_MODELS: - actions.append(self._create_model_action(bedrock_name(model_id))) - - # Register all supported embedding models - for model_id in SUPPORTED_EMBEDDING_MODELS: - actions.append(self._create_embedder_action(bedrock_name(model_id))) - - return actions - - async def close(self) -> None: - """Close the async bedrock-runtime client. - - Exits the aioboto3 async context manager entered in init(), - releasing network connections and other resources. - """ - if self._client_ctx is not None: - await self._client_ctx.__aexit__(None, None, None) - self._client = None - self._client_ctx = None - - async def resolve(self, action_type: ActionKind, name: str) -> Action | None: - """Resolve an action by type and name. - - This enables lazy loading of models not pre-registered during init(). - - Args: - action_type: The kind of action to resolve (MODEL or EMBEDDER). - name: The namespaced name of the action. - - Returns: - Action object if resolvable, None otherwise. - """ - if action_type == ActionKind.MODEL: - return self._create_model_action(name) - elif action_type == ActionKind.EMBEDDER: - return self._create_embedder_action(name) - return None - - def _create_model_action(self, name: str) -> Action: - """Create an Action object for a chat completion model. - - Args: - name: The namespaced model name (e.g., 'amazon-bedrock/anthropic.claude-...'). - - Returns: - Action object for the model. - """ - # Extract model ID (remove plugin prefix) - prefix = f'{AMAZON_BEDROCK_PLUGIN_NAME}/' - model_id = name[len(prefix) :] if name.startswith(prefix) else name - - model = BedrockModel( - model_id=model_id, - client=self._client, - ) - model_info = get_model_info(model_id) - - # Get the appropriate config schema for this model family - config_schema = get_config_schema_for_model(model_id) - - return Action( - kind=ActionKind.MODEL, - name=name, - fn=model.generate, - metadata=model_action_metadata( - name=name, - info=model_info.model_dump(by_alias=True, exclude_none=True), - config_schema=config_schema, - ).metadata, - ) - - def _create_embedder_action(self, name: str) -> Action: - """Create an Action object for an embedding model. - - Args: - name: The namespaced embedder name. - - Returns: - Action object for the embedder. - """ - prefix = f'{AMAZON_BEDROCK_PLUGIN_NAME}/' - model_id = name[len(prefix) :] if name.startswith(prefix) else name - - # Get embedder info - embedder_info = SUPPORTED_EMBEDDING_MODELS.get( - model_id, - { - 'label': f'Bedrock - {model_id}', - 'dimensions': 1024, - 'supports': {'input': ['text']}, - }, - ) - - async def embed_fn(request: EmbedRequest) -> EmbedResponse: - """Generate embeddings using AWS Bedrock. - - The InvokeModel API does NOT support inference profile IDs - (regional prefixes like ``us.``, ``eu.``, ``apac.``). Only - the Converse/ConverseStream APIs support them. We must strip - any inference profile prefix before calling InvokeModel. - """ - # Strip inference profile prefix β€” InvokeModel requires base model IDs - base_model_id = _strip_inference_profile_prefix(model_id) - - # Extract text from document content - texts: list[str] = [] - for doc in request.input: - text_parts: list[str] = [] - for part in doc.content: - if hasattr(part.root, 'text') and part.root.text: - text_parts.append(str(part.root.text)) - doc_text = ''.join(text_parts) - texts.append(doc_text) - - embeddings: list[Embedding] = [] - - # Process each text (Bedrock embedding API typically handles one at a time) - for text in texts: - # Build request body based on model type - if base_model_id.startswith('amazon.titan-embed'): - body = {'inputText': text} - elif base_model_id.startswith('cohere.embed'): - body = { - 'texts': [text], - 'input_type': 'search_document', - } - elif base_model_id.startswith('amazon.nova'): - body = {'inputText': text} - else: - body = {'inputText': text} - - # Call InvokeModel for embeddings (uses base model ID, not inference profile) - response = await self._client.invoke_model( - modelId=base_model_id, - body=json.dumps(body), - contentType='application/json', - accept='application/json', - ) - - # Parse response - response_body = json.loads(await response['body'].read()) - - # Extract embedding based on model type - if base_model_id.startswith('amazon.titan-embed'): - embedding_vector = response_body.get('embedding', []) - elif base_model_id.startswith('cohere.embed'): - embedding_vector = response_body.get('embeddings', [[]])[0] - elif base_model_id.startswith('amazon.nova'): - embedding_vector = response_body.get('embedding', []) - else: - embedding_vector = response_body.get('embedding', []) - - embeddings.append(Embedding(embedding=embedding_vector)) - - return EmbedResponse(embeddings=embeddings) - - return Action( - kind=ActionKind.EMBEDDER, - name=name, - fn=embed_fn, - metadata=embedder_action_metadata( - name=name, - options=EmbedderOptions( - label=embedder_info['label'], - supports=EmbedderSupports(input=embedder_info['supports']['input']), - dimensions=embedder_info.get('dimensions'), - ), - ).metadata, - ) - - async def list_actions(self) -> list[ActionMetadata]: - """List all available models and embedders. - - Returns: - List of ActionMetadata for all supported models and embedders. - """ - actions: list[ActionMetadata] = [] - - # Add model metadata from predefined list - for model_id, model_info in SUPPORTED_BEDROCK_MODELS.items(): - config_schema = get_config_schema_for_model(model_id) - actions.append( - model_action_metadata( - name=bedrock_name(model_id), - info=model_info.model_dump(by_alias=True, exclude_none=True), - config_schema=config_schema, - ) - ) - - # Add embedder metadata from predefined list - for model_id, embed_info in SUPPORTED_EMBEDDING_MODELS.items(): - actions.append( - embedder_action_metadata( - name=bedrock_name(model_id), - options=EmbedderOptions( - label=embed_info['label'], - supports=EmbedderSupports(input=embed_info['supports']['input']), - dimensions=embed_info.get('dimensions'), - ), - ) - ) - - return actions - - -def bedrock_model(model_id: str) -> str: - """Get fully qualified AWS Bedrock model name. - - Convenience function for specifying models. - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html - - Args: - model_id: The Bedrock model ID (e.g., 'anthropic.claude-sonnet-4-5-20250929-v1:0'). - - Returns: - Fully qualified model name (e.g., 'amazon-bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0'). - - Example: - ai = Genkit( - plugins=[AmazonBedrock(region='us-east-1')], - model=bedrock_model('anthropic.claude-sonnet-4-5-20250929-v1:0'), - ) - - # Or for other models: - response = await ai.generate( - model=bedrock_model('meta.llama3-3-70b-instruct-v1:0'), - prompt='Explain quantum computing.', - ) - """ - return bedrock_name(model_id) - - -# ============================================================================= -# Inference Profile Region Helpers -# ============================================================================= -# AWS Bedrock cross-region inference profiles require a regional prefix. -# When using API keys (AWS_BEARER_TOKEN_BEDROCK), you MUST use inference profiles. -# When using IAM credentials, you can use either direct model IDs or inference profiles. -# -# See: https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html - - -def get_inference_profile_prefix(region: str | None = None) -> str: - """Get the inference profile prefix for a given AWS region. - - Args: - region: AWS region code (e.g., 'us-east-1', 'eu-west-1', 'ap-northeast-1'). - If None, uses AWS_REGION or AWS_DEFAULT_REGION environment variable. - - Returns: - The inference profile prefix ('us', 'eu', or 'apac'). - - Raises: - ValueError: If no region is specified and AWS_REGION is not set. - - Example:: - - >>> get_inference_profile_prefix('us-east-1') - 'us' - >>> get_inference_profile_prefix('eu-west-1') - 'eu' - >>> get_inference_profile_prefix('ap-northeast-1') - 'apac' - """ - if region is None: - region = os.environ.get('AWS_REGION') or os.environ.get('AWS_DEFAULT_REGION') - - if region is None: - raise ValueError( - 'AWS region is required for inference profiles. ' - 'Set AWS_REGION environment variable or pass region parameter.' - ) - - region_lower = region.lower() - - if region_lower.startswith('us-') or region_lower.startswith('us.'): - return 'us' - elif region_lower.startswith('eu-') or region_lower.startswith('eu.'): - return 'eu' - elif region_lower.startswith('ap-') or region_lower.startswith('ap.'): - return 'apac' - elif region_lower.startswith('ca-'): - # Canada is routed through US inference profiles - return 'us' - elif region_lower.startswith('sa-'): - # South America is routed through US inference profiles - return 'us' - elif region_lower.startswith('me-'): - # Middle East - check AWS docs, but typically EU - return 'eu' - elif region_lower.startswith('af-'): - # Africa - check AWS docs, but typically EU - return 'eu' - else: - # Default to US - return 'us' - - -def inference_profile(model_id: str, region: str | None = None) -> str: - """Convert a model ID to an inference profile ID for the given region. - - Use this when you need to use API keys (AWS_BEARER_TOKEN_BEDROCK) which - require inference profiles instead of direct model IDs. - - Args: - model_id: Base model ID (e.g., 'anthropic.claude-sonnet-4-5-20250929-v1:0'). - region: AWS region code. If None, uses AWS_REGION environment variable. - - Returns: - The Genkit model reference with inference profile ID. - - Example:: - - # Using environment variable AWS_REGION=eu-west-1 - >>> inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0') - 'amazon-bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0' - - # Explicit region - >>> inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0', 'ap-northeast-1') - 'amazon-bedrock/apac.anthropic.claude-sonnet-4-5-20250929-v1:0' - """ - prefix = get_inference_profile_prefix(region) - return bedrock_name(f'{prefix}.{model_id}') - - -# ============================================================================= -# Pre-defined Model References -# ============================================================================= -# These use DIRECT model IDs (without regional prefix) which work with: -# - IAM credentials (AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY) -# - IAM roles (EC2, Lambda, ECS, etc.) -# -# For API keys (AWS_BEARER_TOKEN_BEDROCK), use the inference_profile() helper: -# model = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0') -# -# Or use the region-specific helpers defined below. - -# Anthropic Claude models β€” grouped by family, descending version -claude_opus_4_6 = bedrock_name('anthropic.claude-opus-4-6-20260205-v1:0') -claude_opus_4_5 = bedrock_name('anthropic.claude-opus-4-5-20251101-v1:0') -claude_opus_4_1 = bedrock_name('anthropic.claude-opus-4-1-20250805-v1:0') -claude_sonnet_4_5 = bedrock_name('anthropic.claude-sonnet-4-5-20250929-v1:0') -claude_sonnet_4 = bedrock_name('anthropic.claude-sonnet-4-20250514-v1:0') -claude_haiku_4_5 = bedrock_name('anthropic.claude-haiku-4-5-20251001-v1:0') -claude_3_5_haiku = bedrock_name('anthropic.claude-3-5-haiku-20241022-v1:0') -claude_3_haiku = bedrock_name('anthropic.claude-3-haiku-20240307-v1:0') - -# Amazon Nova models -nova_pro = bedrock_name('amazon.nova-pro-v1:0') -nova_lite = bedrock_name('amazon.nova-lite-v1:0') -nova_micro = bedrock_name('amazon.nova-micro-v1:0') -nova_premier = bedrock_name('amazon.nova-premier-v1:0') - -# Meta Llama models -llama_3_3_70b = bedrock_name('meta.llama3-3-70b-instruct-v1:0') -llama_3_1_405b = bedrock_name('meta.llama3-1-405b-instruct-v1:0') -llama_3_1_70b = bedrock_name('meta.llama3-1-70b-instruct-v1:0') -llama_4_maverick = bedrock_name('meta.llama4-maverick-17b-instruct-v1:0') -llama_4_scout = bedrock_name('meta.llama4-scout-17b-instruct-v1:0') - -# Mistral models -mistral_large_3 = bedrock_name('mistral.mistral-large-3-675b-instruct') -mistral_large = bedrock_name('mistral.mistral-large-2407-v1:0') -pixtral_large = bedrock_name('mistral.pixtral-large-2502-v1:0') - -# DeepSeek models -deepseek_r1 = bedrock_name('deepseek.r1-v1:0') -deepseek_v3 = bedrock_name('deepseek.v3-v1:0') - -# Cohere models -command_r_plus = bedrock_name('cohere.command-r-plus-v1:0') -command_r = bedrock_name('cohere.command-r-v1:0') - -# AI21 Jamba models -jamba_large = bedrock_name('ai21.jamba-1-5-large-v1:0') -jamba_mini = bedrock_name('ai21.jamba-1-5-mini-v1:0') diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/py.typed b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/py.typed deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/telemetry/__init__.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/telemetry/__init__.py deleted file mode 100644 index 37ae3c8012..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/telemetry/__init__.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS telemetry integration for Genkit. - -This package provides telemetry export to AWS's observability suite, -enabling monitoring and debugging of Genkit applications through AWS X-Ray -for distributed tracing and CloudWatch for metrics and logs. - -Module Structure:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Module β”‚ Purpose β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ tracing.py β”‚ Main entry point, X-Ray OTLP exporter, SigV4 auth β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Quick Start: - ```python - from genkit.plugins.amazon_bedrock import add_aws_telemetry - - # Enable telemetry with defaults (uses AWS_REGION env var) - add_aws_telemetry() - - # Or with explicit region - add_aws_telemetry(region='us-west-2') - - # Or disable PII redaction (caution!) - add_aws_telemetry(log_input_and_output=True) - ``` - -Cross-Language Parity: - This implementation maintains feature parity with: - - Google Cloud plugin: py/plugins/google-cloud/ - -See Also: - - tracing.py module docstring for detailed architecture documentation - -AWS Documentation: - X-Ray: - - Overview: https://docs.aws.amazon.com/xray/ - - OTLP Endpoint: https://docs.aws.amazon.com/xray/latest/devguide/xray-api-sendingdata.html - - IAM Roles: AWSXrayWriteOnlyPolicy - - CloudWatch: - - Overview: https://docs.aws.amazon.com/cloudwatch/ - - OTLP Endpoint: https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/CloudWatch-OTLPEndpoint.html - - OpenTelemetry AWS: - - ADOT Python: https://aws-otel.github.io/docs/getting-started/python-sdk - - SDK Extension: https://pypi.org/project/opentelemetry-sdk-extension-aws/ -""" - -from .tracing import add_aws_telemetry - -__all__ = ['add_aws_telemetry'] diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/telemetry/tracing.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/telemetry/tracing.py deleted file mode 100644 index 7bfa90ca31..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/telemetry/tracing.py +++ /dev/null @@ -1,851 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Telemetry and tracing functionality for the Genkit AWS plugin. - -This module provides functionality for collecting and exporting telemetry data -from Genkit operations to AWS. It uses OpenTelemetry for tracing and exports -span data to AWS X-Ray for monitoring and debugging purposes. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Span β”‚ A "timer" that records how long something took. β”‚ - β”‚ β”‚ Like a stopwatch for one task in your code. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Trace β”‚ A collection of spans showing a request's journey. β”‚ - β”‚ β”‚ Like breadcrumbs through your code. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ X-Ray β”‚ AWS service that collects and visualizes traces. β”‚ - β”‚ β”‚ Like a detective board connecting all the clues. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ OTLP β”‚ OpenTelemetry Protocol - a standard way to send β”‚ - β”‚ β”‚ traces. Like a universal shipping label format. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ SigV4 β”‚ AWS's way of signing requests to prove identity. β”‚ - β”‚ β”‚ Like a secret handshake with AWS. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Propagator β”‚ Passes trace IDs between services. Like a relay β”‚ - β”‚ β”‚ baton so X-Ray can connect the dots. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ ID Generator β”‚ Creates trace IDs in X-Ray's special format. β”‚ - β”‚ β”‚ X-Ray needs timestamps baked into IDs. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Sampler β”‚ Decides which traces to keep. Like a bouncer β”‚ - β”‚ β”‚ deciding which requests get recorded. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Exporter β”‚ Ships your traces to X-Ray. Like a postal service β”‚ - β”‚ β”‚ for your telemetry data. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ PII Redaction β”‚ Removes sensitive data from traces. Like blacking β”‚ - β”‚ β”‚ out private info before sharing. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW YOUR CODE GETS TRACED β”‚ - β”‚ β”‚ - β”‚ Your Genkit App β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) You call a flow, model, or tool β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Flow β”‚ ──▢ β”‚ Model β”‚ ──▢ β”‚ Tool β”‚ Each creates a "span" β”‚ - β”‚ β”‚ (span) β”‚ β”‚ (span) β”‚ β”‚ (span) β”‚ (a timing record) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Spans collected into a trace β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Trace β”‚ All spans for one request β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Adjustments applied β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AwsAdjustingTraceExporter β”‚ β”‚ - β”‚ β”‚ β€’ Redact PII (input/output)β”‚ β”‚ - β”‚ β”‚ β€’ Fix zero-duration spans β”‚ β”‚ - β”‚ β”‚ β€’ Add error markers β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (4) Sent to AWS via OTLP/HTTP β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AwsXRayOtlpExporter β”‚ β”‚ - β”‚ β”‚ (with SigV4 auth) β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (5) HTTPS to X-Ray endpoint β”‚ - β”‚ β–Ό β”‚ - β”‚ ════════════════════════════════════════════════════ β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AWS X-Ray β”‚ View traces in AWS Console β”‚ - β”‚ β”‚ (your traces!) β”‚ Debug latency, errors, etc. β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Architecture Overview:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ TELEMETRY DATA FLOW β”‚ - β”‚ β”‚ - β”‚ Genkit Actions (flows, models, tools) β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ OpenTelemetry β”‚ Creates spans with genkit:* attributes β”‚ - β”‚ β”‚ Tracer β”‚ (type, name, input, output, state, path, etc.) β”‚ - β”‚ β”‚ + AwsXRayIdGen β”‚ Uses X-Ray-compatible trace ID format β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AwsAdjustingTraceExporter β”‚ β”‚ - β”‚ β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ β”‚ - β”‚ β”‚ β”‚ AdjustingTraceExporter._adjust() β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ - Redact genkit:input/output β†’ "" β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ - Mark error spans with /http/status_code: 599 β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ - Normalize labels for X-Ray compatibility β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ - TimeAdjustedSpan ensures end > start β”‚ β”‚ β”‚ - β”‚ β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AwsXRayOtlpExporter β”‚ β”‚ - β”‚ β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ β”‚ - β”‚ β”‚ β”‚ OTLP/HTTP Export with SigV4 Authentication β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ Endpoint: https://xray.{region}.amazonaws.com/v1/traces β”‚ β”‚ β”‚ - β”‚ β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ AWS X-Ray β”‚ β”‚ - β”‚ β”‚ Service β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Components: - 1. **AwsTelemetry**: Central manager class that encapsulates configuration - and manages the lifecycle of tracing and logging setup. - - 2. **AwsXRayOtlpExporter**: Custom OTLP/HTTP exporter with SigV4 signing. - Uses botocore credentials to authenticate requests to AWS X-Ray. - - 3. **AwsAdjustingTraceExporter**: Extends AdjustingTraceExporter to add - AWS-specific handling before spans are adjusted and exported. - - 4. **TimeAdjustedSpan**: Wrapper that ensures spans have non-zero duration - (X-Ray requires end_time > start_time). - - 5. **AwsXRayIdGenerator**: From opentelemetry-sdk-extension-aws, generates - X-Ray-compatible trace IDs with Unix timestamp in first 32 bits. - - 6. **AwsXRayPropagator**: From opentelemetry-propagator-aws-xray, handles - trace context propagation across AWS services. - -X-Ray Trace ID Requirements: - AWS X-Ray requires trace IDs to have Unix epoch time in the first 32 bits. - The AwsXRayIdGenerator ensures compatibility. Spans with timestamps older - than 30 days are rejected by X-Ray. - -Log-Trace Correlation: - The AwsTelemetry manager injects X-Ray trace context into structlog logs, - enabling correlation between logs and traces in CloudWatch. The format - uses `_X_AMZN_TRACE_ID` which matches AWS X-Ray conventions. - -Configuration Options:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Option β”‚ Type β”‚ Default β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ region β”‚ str β”‚ AWS_REGION β”‚ AWS region β”‚ - β”‚ log_input_and_output β”‚ bool β”‚ False β”‚ Disable redaction β”‚ - β”‚ force_dev_export β”‚ bool β”‚ True β”‚ Export in dev β”‚ - β”‚ disable_traces β”‚ bool β”‚ False β”‚ Skip traces β”‚ - β”‚ sampler β”‚ Sampler β”‚ AlwaysOn β”‚ Trace sampler β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Region Resolution Order: - 1. Explicit region parameter - 2. AWS_REGION environment variable - 3. AWS_DEFAULT_REGION environment variable - -Usage: - ```python - from genkit.plugins.amazon_bedrock import add_aws_telemetry - - # Enable telemetry with default settings (PII redaction enabled) - add_aws_telemetry() - - # Enable telemetry with explicit region - add_aws_telemetry(region='us-west-2') - - # Enable input/output logging (disable PII redaction) - add_aws_telemetry(log_input_and_output=True) - - # Force export in dev environment - add_aws_telemetry(force_dev_export=True) - ``` - -AWS Documentation References: - X-Ray: - - Overview: https://docs.aws.amazon.com/xray/ - - OTLP Endpoint: https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/CloudWatch-OTLPEndpoint.html - - IAM Role: AWSXrayWriteOnlyPolicy - - ADOT Python: - - Getting Started: https://aws-otel.github.io/docs/getting-started/python-sdk - - Migration Guide: https://docs.aws.amazon.com/xray/latest/devguide/migrate-xray-to-opentelemetry-python.html -""" - -import os -import uuid -from collections.abc import Callable, Mapping, MutableMapping, Sequence -from typing import Any, cast - -import requests -import structlog -from botocore.auth import SigV4Auth -from botocore.awsrequest import AWSRequest -from botocore.session import Session as BotocoreSession -from opentelemetry import propagate, trace -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter -from opentelemetry.propagators.aws import AwsXRayPropagator -from opentelemetry.sdk.extension.aws.trace import AwsXRayIdGenerator -from opentelemetry.sdk.resources import SERVICE_INSTANCE_ID, SERVICE_NAME, Resource -from opentelemetry.sdk.trace import ReadableSpan, TracerProvider -from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult -from opentelemetry.sdk.trace.sampling import Sampler -from requests.adapters import HTTPAdapter -from urllib3.util.retry import Retry - -from genkit.core.environment import is_dev_environment -from genkit.core.trace.adjusting_exporter import AdjustingTraceExporter, RedactedSpan -from genkit.core.tracing import add_custom_exporter - -logger = structlog.get_logger(__name__) - -# X-Ray OTLP endpoint pattern -XRAY_OTLP_ENDPOINT_PATTERN = 'https://xray.{region}.amazonaws.com/v1/traces' - - -def _resolve_region(region: str | None = None) -> str | None: - """Resolve the AWS region from various sources. - - Resolution order: - 1. Explicit region parameter - 2. AWS_REGION environment variable - 3. AWS_DEFAULT_REGION environment variable - - Args: - region: Explicitly provided AWS region. - - Returns: - The resolved region or None if not found. - """ - if region: - return region - - # Check environment variables in order of priority - for env_var in ('AWS_REGION', 'AWS_DEFAULT_REGION'): - env_value = os.environ.get(env_var) - if env_value: - return env_value - - return None - - -class AwsTelemetry: - """Central manager for AWS Telemetry configuration. - - Encapsulates configuration and manages the lifecycle of Tracing and Logging - setup, ensuring consistent state (like region) across all telemetry components. - - Example: - ```python - telemetry = AwsTelemetry(region='us-west-2') - telemetry.initialize() - ``` - """ - - def __init__( - self, - region: str | None = None, - sampler: Sampler | None = None, - log_input_and_output: bool = False, - force_dev_export: bool = True, - disable_traces: bool = False, - ) -> None: - """Initialize the AWS Telemetry manager. - - Args: - region: AWS region for X-Ray endpoint. - sampler: Trace sampler. - log_input_and_output: If False, redacts sensitive data. - force_dev_export: If True, exports even in dev environment. - disable_traces: If True, traces are not exported. - - Raises: - ValueError: If region cannot be resolved. - """ - self.sampler = sampler - self.log_input_and_output = log_input_and_output - self.force_dev_export = force_dev_export - self.disable_traces = disable_traces - - # Resolve region immediately - self.region = _resolve_region(region) - - if self.region is None: - raise ValueError( - 'AWS region is required. Set AWS_REGION environment variable ' - 'or pass region parameter to add_aws_telemetry().' - ) - - def initialize(self) -> None: - """Actuate the telemetry configuration. - - Sets up logging with trace correlation and configures tracing export. - """ - is_dev = is_dev_environment() - should_export = self.force_dev_export or not is_dev - - if not should_export: - logger.debug('Telemetry export disabled in dev environment') - return - - self._configure_logging() - self._configure_tracing() - - def _configure_logging(self) -> None: - """Configure structlog with X-Ray trace correlation. - - Injects X-Ray trace ID into log records for correlation in CloudWatch. - """ - try: - current_config = structlog.get_config() - processors = current_config.get('processors', []) - - # Check if our processor is already registered (by function name) - # Note: Use the exact function name being added, not the method name - if not any(getattr(p, '__name__', '') == 'inject_xray_trace_context' for p in processors): - - def inject_xray_trace_context( - _logger: Any, # noqa: ANN401 - method_name: str, - event_dict: MutableMapping[str, Any], - ) -> Mapping[str, Any]: - """Inject X-Ray trace context into log event.""" - return self._inject_trace_context(event_dict) - - new_processors = list(processors) - # Insert before the last processor (usually the renderer) - new_processors.insert(max(0, len(new_processors) - 1), inject_xray_trace_context) - cfg = structlog.get_config() - structlog.configure( - processors=new_processors, - wrapper_class=cfg.get('wrapper_class'), - context_class=cfg.get('context_class'), - logger_factory=cfg.get('logger_factory'), - cache_logger_on_first_use=cfg.get('cache_logger_on_first_use'), - ) - logger.debug('Configured structlog for AWS X-Ray trace correlation') - - except Exception as e: - logger.warning('Failed to configure structlog for trace correlation', error=str(e)) - - def _configure_tracing(self) -> None: - """Configure trace export to AWS X-Ray.""" - if self.disable_traces: - return - - # Region is guaranteed to be set by __init__ (raises ValueError if None) - assert self.region is not None - - # Configure X-Ray propagator for trace context - propagate.set_global_textmap(AwsXRayPropagator()) - - # Create resource with service info - resource = Resource.create({ - SERVICE_NAME: 'genkit', - SERVICE_INSTANCE_ID: str(uuid.uuid4()), - }) - - # Create TracerProvider with X-Ray ID generator - provider = TracerProvider( - resource=resource, - id_generator=AwsXRayIdGenerator(), - sampler=self.sampler, - ) - trace.set_tracer_provider(provider) - - # Create the base X-Ray OTLP exporter - base_exporter = AwsXRayOtlpExporter( - region=self.region, - error_handler=lambda e: _handle_tracing_error(e), - ) - - # Wrap with AwsAdjustingTraceExporter for PII redaction - trace_exporter = AwsAdjustingTraceExporter( - exporter=base_exporter, - log_input_and_output=self.log_input_and_output, - region=self.region, - error_handler=lambda e: _handle_tracing_error(e), - ) - - add_custom_exporter(trace_exporter, 'aws_xray_telemetry') - - logger.info( - 'AWS X-Ray telemetry configured', - region=self.region, - endpoint=XRAY_OTLP_ENDPOINT_PATTERN.format(region=self.region), - ) - - def _inject_trace_context(self, event_dict: MutableMapping[str, Any]) -> MutableMapping[str, Any]: - """Inject AWS X-Ray trace context into log event. - - Adds X-Ray trace ID in the format expected by CloudWatch Logs for - correlation with X-Ray traces. - - Args: - event_dict: The structlog event dictionary. - - Returns: - The event dictionary with trace context added. - """ - span = trace.get_current_span() - if span == trace.INVALID_SPAN: - return event_dict - - ctx = span.get_span_context() - if not ctx.is_valid: - return event_dict - - # Format trace ID for X-Ray (first 8 chars are timestamp, rest is random) - # X-Ray trace ID format: 1-{timestamp}-{random} - trace_id_hex = f'{ctx.trace_id:032x}' - xray_trace_id = f'1-{trace_id_hex[:8]}-{trace_id_hex[8:]}' - - # Add X-Ray trace header format for CloudWatch correlation - sampled = '1' if ctx.trace_flags.sampled else '0' - event_dict['_X_AMZN_TRACE_ID'] = f'Root={xray_trace_id};Parent={ctx.span_id:016x};Sampled={sampled}' - - return event_dict - - -class SigV4SigningAdapter(HTTPAdapter): - """HTTP adapter that signs requests with AWS SigV4 authentication. - - This adapter intercepts all HTTP requests and signs them with AWS SigV4 - before sending. This enables the OTLP exporter to authenticate with - AWS services like X-Ray that require SigV4. - - Example: - ```python - session = requests.Session() - adapter = SigV4SigningAdapter( - credentials=botocore_session.get_credentials(), - region='us-west-2', - service='xray', - ) - session.mount('https://', adapter) - ``` - """ - - def __init__( - self, - credentials: Any, # noqa: ANN401 - botocore credentials type - region: str, - service: str = 'xray', - **kwargs: Any, # noqa: ANN401 - ) -> None: - """Initialize the SigV4 signing adapter. - - Args: - credentials: Botocore credentials object. - region: AWS region for signing. - service: AWS service name for signing (default: 'xray'). - **kwargs: Additional arguments passed to HTTPAdapter. - """ - super().__init__(**kwargs) - self._credentials = credentials - self._region = region - self._service = service - - def send( # type: ignore[override] - self, - request: requests.PreparedRequest, - **kwargs: Any, # noqa: ANN401 - ) -> requests.Response: - """Send the request after signing it with SigV4. - - Args: - request: The prepared request to send. - **kwargs: Additional arguments passed to the parent send method. - - Returns: - The response from the server. - """ - if self._credentials is not None: - # Sign the request - aws_request = AWSRequest( - method=request.method or 'POST', - url=request.url or '', - headers=dict(request.headers) if request.headers else {}, - data=request.body or b'', - ) - SigV4Auth(self._credentials, self._service, self._region).add_auth(aws_request) - - # Update the original request with signed headers - if request.headers is None: - request.headers = cast(Any, {}) - request.headers.update(dict(aws_request.headers)) - - return super().send(request, **kwargs) - - -def _create_sigv4_session( - credentials: Any, # noqa: ANN401 - botocore credentials type - region: str, - service: str = 'xray', -) -> requests.Session: - """Create a requests Session that signs all requests with SigV4. - - Args: - credentials: Botocore credentials object. - region: AWS region for signing. - service: AWS service name for signing. - - Returns: - A configured requests Session with SigV4 signing. - """ - session = requests.Session() - - # Configure retry logic - retry = Retry( - total=3, - backoff_factor=0.5, - status_forcelist=[500, 502, 503, 504], - ) - - # Create signing adapter with retry - adapter = SigV4SigningAdapter( - credentials=credentials, - region=region, - service=service, - max_retries=retry, - ) - - # Mount the signing adapter for HTTPS requests - session.mount('https://', adapter) - - return session - - -class AwsXRayOtlpExporter(SpanExporter): - """OTLP/HTTP exporter with AWS SigV4 authentication for X-Ray. - - This exporter sends spans via OTLP/HTTP to the AWS X-Ray endpoint, - signing each request with AWS SigV4 authentication using botocore. - - The SigV4 signing is implemented via a custom requests Session adapter - that intercepts all outgoing HTTP requests and adds the required - Authorization, X-Amz-Date, and X-Amz-Security-Token headers. - - Args: - region: AWS region for the X-Ray endpoint. - error_handler: Optional callback for export errors. - - Example: - ```python - exporter = AwsXRayOtlpExporter(region='us-west-2') - ``` - - Note: - Uses standard AWS credential chain (environment variables, IAM role, - credential file, etc.) via botocore. - """ - - def __init__( - self, - region: str, - error_handler: Callable[[Exception], None] | None = None, - ) -> None: - """Initialize the X-Ray OTLP exporter. - - Args: - region: AWS region for the X-Ray endpoint. - error_handler: Optional callback invoked when export errors occur. - """ - self._region = region - self._error_handler = error_handler - self._endpoint = XRAY_OTLP_ENDPOINT_PATTERN.format(region=region) - - # Initialize botocore session for SigV4 signing - self._botocore_session = BotocoreSession() - self._credentials = self._botocore_session.get_credentials() - - if self._credentials is None: - logger.warning( - 'No AWS credentials found for SigV4 signing. X-Ray export will fail without valid credentials.' - ) - - # Create a session with SigV4 signing adapter - signing_session = _create_sigv4_session( - credentials=self._credentials, - region=region, - service='xray', - ) - - # Create the underlying OTLP exporter with signing session - self._otlp_exporter = OTLPSpanExporter( - endpoint=self._endpoint, - session=signing_session, - ) - - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: - """Export spans to AWS X-Ray via OTLP/HTTP with SigV4 authentication. - - All requests are automatically signed with AWS SigV4 using the - credentials from the botocore session. - - Args: - spans: A sequence of OpenTelemetry ReadableSpan objects to export. - - Returns: - SpanExportResult indicating success or failure. - """ - try: - return self._otlp_exporter.export(spans) - except Exception as ex: - logger.error('Error while writing to X-Ray', exc_info=ex) - if self._error_handler: - self._error_handler(ex) - return SpanExportResult.FAILURE - - def shutdown(self) -> None: - """Shutdown the exporter.""" - self._otlp_exporter.shutdown() - - def force_flush(self, timeout_millis: int = 30000) -> bool: - """Force flush pending spans. - - Args: - timeout_millis: Timeout in milliseconds. - - Returns: - True if flush succeeded. - """ - return self._otlp_exporter.force_flush(timeout_millis) - - -class TimeAdjustedSpan(RedactedSpan): - """Wraps a span to ensure non-zero duration. - - X-Ray requires end_time > start_time. This wrapper ensures spans have - at least 1 microsecond duration. - """ - - @property - def end_time(self) -> int | None: - """Return the span end time, adjusted to be > start_time. - - Returns: - The end time, with minimum 1 microsecond after start if needed. - """ - start = self._span.start_time - end = self._span.end_time - - # X-Ray requires end_time > start_time. - # If the span is unfinished (end_time is None) or has zero duration, - # we provide a minimum 1 microsecond duration. - if start is not None and (end is None or end <= start): - return start + 1000 # 1 microsecond in nanoseconds - - return end - - -class AwsAdjustingTraceExporter(AdjustingTraceExporter): - """AWS-specific span exporter that adds X-Ray compatibility. - - This extends the base AdjustingTraceExporter to handle AWS X-Ray - specific requirements before spans are adjusted and exported. - - Example: - ```python - exporter = AwsAdjustingTraceExporter( - exporter=AwsXRayOtlpExporter(region='us-west-2'), - log_input_and_output=False, - region='us-west-2', - ) - ``` - """ - - def __init__( - self, - exporter: SpanExporter, - log_input_and_output: bool = False, - region: str | None = None, - error_handler: Callable[[Exception], None] | None = None, - ) -> None: - """Initialize the AWS adjusting trace exporter. - - Args: - exporter: The underlying SpanExporter to wrap. - log_input_and_output: If True, preserve input/output in spans. - Defaults to False (redact for privacy). - region: AWS region (for future use in telemetry). - error_handler: Optional callback invoked when export errors occur. - """ - super().__init__( - exporter=exporter, - log_input_and_output=log_input_and_output, - error_handler=error_handler, - ) - # Store region for potential future AWS-specific telemetry - self._region = region - - def _adjust(self, span: ReadableSpan) -> ReadableSpan: - """Apply all adjustments to a span including time adjustment. - - This overrides the base method to add time adjustment for X-Ray - compatibility (end_time must be > start_time). - - Args: - span: The span to adjust. - - Returns: - The adjusted span with guaranteed non-zero duration. - """ - # Apply standard adjustments from base class - span = super()._adjust(span) - - # Fix start/end times for X-Ray (must be end > start) - return TimeAdjustedSpan(span, dict(span.attributes) if span.attributes else {}) - - -def add_aws_telemetry( - region: str | None = None, - sampler: Sampler | None = None, - log_input_and_output: bool = False, - force_dev_export: bool = True, - disable_traces: bool = False, -) -> None: - """Configure AWS telemetry export for traces to X-Ray. - - This function sets up OpenTelemetry export to AWS X-Ray. By default, - model inputs and outputs are redacted for privacy protection. - - Args: - region: AWS region for X-Ray endpoint. If not provided, uses - AWS_REGION or AWS_DEFAULT_REGION environment variables. - sampler: OpenTelemetry trace sampler. Controls which traces are - collected and exported. Defaults to AlwaysOnSampler. Common options: - - AlwaysOnSampler: Collect all traces - - AlwaysOffSampler: Collect no traces - - TraceIdRatioBasedSampler: Sample a percentage of traces - log_input_and_output: If True, preserve model input/output in traces. - Defaults to False (redact for privacy). Only enable this in - trusted environments where PII exposure is acceptable. - force_dev_export: If True, export telemetry even in dev environment. - Defaults to True. Set to False for production-only telemetry. - disable_traces: If True, traces will not be exported. - Defaults to False. - - Raises: - ValueError: If region cannot be resolved from parameters or environment. - - Example: - ```python - # Default: PII redaction enabled, uses AWS_REGION env var - add_aws_telemetry() - - # Enable input/output logging (disable PII redaction) - add_aws_telemetry(log_input_and_output=True) - - # Force export in dev environment with specific region - add_aws_telemetry(force_dev_export=True, region='us-west-2') - ``` - - Note: - This implementation currently sends traces to an OTLP endpoint. - For collector-less export with SigV4 authentication, either: - 1. Run an ADOT collector locally that handles authentication - 2. Use ADOT auto-instrumentation with environment variables: - - OTEL_PYTHON_DISTRO=aws_distro - - OTEL_PYTHON_CONFIGURATOR=aws_configurator - - OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=https://xray.{region}.amazonaws.com/v1/traces - - See Also: - - AWS X-Ray: https://docs.aws.amazon.com/xray/ - - ADOT Python: https://aws-otel.github.io/docs/getting-started/python-sdk - """ - manager = AwsTelemetry( - region=region, - sampler=sampler, - log_input_and_output=log_input_and_output, - force_dev_export=force_dev_export, - disable_traces=disable_traces, - ) - - manager.initialize() - - -# Error handling helpers -_tracing_error_logged = False - - -def _handle_tracing_error(error: Exception) -> None: - """Handle trace export errors with helpful messages. - - Only logs detailed instructions once to avoid spam. - - Args: - error: The export error. - """ - global _tracing_error_logged - if _tracing_error_logged: - return - - error_str = str(error).lower() - if 'permission' in error_str or 'denied' in error_str or '403' in error_str: - _tracing_error_logged = True - logger.error( - 'Unable to send traces to AWS X-Ray. ' - 'Ensure the IAM role/user has the "AWSXrayWriteOnlyPolicy" policy ' - 'or xray:PutTraceSegments permission. ' - f'Error: {error}' - ) - elif 'credential' in error_str or 'unauthorized' in error_str or '401' in error_str: - _tracing_error_logged = True - logger.error( - 'AWS credentials not found or invalid. ' - 'Configure credentials via environment variables (AWS_ACCESS_KEY_ID, ' - 'AWS_SECRET_ACCESS_KEY), IAM role, or ~/.aws/credentials file. ' - f'Error: {error}' - ) - else: - logger.error('Error exporting traces to AWS X-Ray', error=str(error)) diff --git a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/typing.py b/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/typing.py deleted file mode 100644 index 379cb937cd..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/plugins/amazon_bedrock/typing.py +++ /dev/null @@ -1,951 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""AWS Bedrock configuration types for Genkit. - -This module defines configuration schemas that align with the AWS Bedrock -Converse API and model-specific parameters for each supported provider. - -Architecture Overview:: - - +------------------+ +----------------------+ +------------------+ - | Genkit App | | AWS Bedrock Plugin | | AWS Bedrock | - | | | | | Service | - | +------------+ | | +----------------+ | | | - | | ai.generate|------->| | BedrockModel |------->| Converse API | - | +------------+ | | +----------------+ | | ConverseStream | - | | | | | | - | +------------+ | | +----------------+ | | +------------+ | - | | ai.embed |------->| | EmbedAction |------->| | Embeddings | | - | +------------+ | | +----------------+ | | +------------+ | - +------------------+ +----------------------+ +------------------+ - | - v - +------------------+ - | Model Configs | - | (this module) | - | | - | - BedrockConfig | - | - AnthropicConfig| - | - MetaLlamaConfig| - | - CohereConfig | - | - MistralConfig | - | - etc. | - +------------------+ - -Design Rationale: - We use static configuration classes with `extra='allow'` rather than dynamic - parameter discovery for several reasons: - - 1. **Type Safety**: Static configs provide IDE autocompletion, type checking, - and validation for known parameters with documented constraints. - - 2. **Forward Compatibility**: The `extra='allow'` Pydantic setting allows any - additional parameters to pass through, supporting new/undocumented params. - - 3. **DevUI Compatibility**: Model-specific configs enable the Genkit DevUI to - show relevant configuration options for each model family. - -Common Parameters NOT in Converse API: - When using models through the Converse API, some native model parameters - may need to be passed via `additionalModelRequestFields`: - - +-------------------+------------------+--------------------------------------+ - | Parameter | Native API | Converse Alternative | - +-------------------+------------------+--------------------------------------+ - | top_k | Anthropic, | Pass via additionalModelRequestFields| - | | Cohere | or use top_p instead. | - +-------------------+------------------+--------------------------------------+ - | random_seed | Mistral | Pass via additionalModelRequestFields| - +-------------------+------------------+--------------------------------------+ - | safe_prompt | Mistral | Pass via additionalModelRequestFields| - +-------------------+------------------+--------------------------------------+ - | documents | Cohere | Pass via additionalModelRequestFields| - +-------------------+------------------+--------------------------------------+ - | thinking | Anthropic | Pass via additionalModelRequestFields| - +-------------------+------------------+--------------------------------------+ - -Supported Model Providers: - This plugin supports foundation models available through the Bedrock service: - - +------------------+--------------------------------------------------+ - | Provider | Models | - +------------------+--------------------------------------------------+ - | Amazon | Nova Pro, Nova Lite, Nova Micro, Nova Premier, | - | | Nova Canvas, Nova Reel, Nova Sonic, Titan | - | Anthropic | Claude Sonnet 4.5, Claude Opus 4.5, Claude 4, | - | | Claude Haiku 4.5, Claude 3.5 Haiku, Claude 3 | - | AI21 Labs | Jamba 1.5 Large, Jamba 1.5 Mini | - | Cohere | Command R, Command R+, Embed v4, Rerank 3.5 | - | DeepSeek | DeepSeek-R1, DeepSeek-V3.1 | - | Google | Gemma 3 4B IT, Gemma 3 12B IT, Gemma 3 27B IT | - | Luma AI | Ray v2 (video generation) | - | Meta | Llama 4 Maverick, Llama 4 Scout, Llama 3.3, | - | | Llama 3.2, Llama 3.1, Llama 3 | - | MiniMax | MiniMax M2 | - | Mistral AI | Mistral Large 3, Pixtral Large, Magistral Small, | - | | Ministral, Mixtral, Voxtral | - | Moonshot AI | Kimi K2 Thinking | - | NVIDIA | Nemotron Nano 9B v2, Nemotron Nano 12B v2 | - | OpenAI | GPT OSS 120B, GPT OSS 20B, GPT OSS Safeguard | - | Qwen | Qwen3 32B, Qwen3 235B, Qwen3 Coder, Qwen3 VL | - | Stability AI | Stable Diffusion 3.5, Stable Image Core/Ultra | - | TwelveLabs | Pegasus, Marengo Embed (video understanding) | - | Writer | Palmyra X4, Palmyra X5 | - +------------------+--------------------------------------------------+ - -See Also: - AWS Bedrock Documentation: - - Model Parameters: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html - - Converse API: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html - - InferenceConfiguration: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InferenceConfiguration.html - - Model-Specific Documentation: - - Anthropic: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html - - Meta Llama: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html - - Mistral: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html - - Cohere: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html - - DeepSeek: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-deepseek.html - - Amazon Nova: https://docs.aws.amazon.com/nova/latest/userguide/getting-started-schema.html - -Trademark Notice: - This is a community plugin and is not officially supported by Amazon Web Services. - "Amazon", "AWS", "Amazon Bedrock", and related marks are trademarks of - Amazon.com, Inc. or its affiliates. Model names (Claude, Llama, Mistral, etc.) - are trademarks of their respective owners. -""" - -import sys -from typing import Any, ClassVar, Literal - -if sys.version_info < (3, 11): - from strenum import StrEnum -else: - from enum import StrEnum - -from pydantic import BaseModel, ConfigDict, Field -from pydantic.alias_generators import to_camel - - -class CohereSafetyMode(StrEnum): - """Safety mode for Cohere models. - - Controls the safety instruction inserted into the prompt. - - See: https://docs.cohere.com/v2/docs/safety-modes - """ - - CONTEXTUAL = 'CONTEXTUAL' - STRICT = 'STRICT' - OFF = 'OFF' - - -class CohereToolChoice(StrEnum): - """Tool choice for Cohere models. - - Controls whether the model is forced to use a tool. - - See: https://docs.cohere.com/v2/reference/chat - """ - - REQUIRED = 'REQUIRED' - NONE = 'NONE' - - -class GenkitCommonConfigMixin(BaseModel): - """Genkit common configuration parameters mixin. - - These parameters match the Genkit core GenerationCommonConfigSchema and are - expected by the Genkit DevUI for proper rendering of the config pane. - - Reference: - - JS Schema: js/ai/src/model-types.ts (GenerationCommonConfigSchema) - - Python Schema: genkit/core/typing.py (GenerationCommonConfig) - - When creating model configs, inherit from this mixin (via BedrockConfig) - to ensure DevUI compatibility. - - Parameters: - version: A specific version of the model family. - temperature: Controls randomness in token selection (0.0-1.0). - max_output_tokens: Maximum number of tokens to generate. - top_k: Maximum number of tokens to consider when sampling. - top_p: Nucleus sampling probability mass (0.0-1.0). - stop_sequences: Strings that will stop output generation. - """ - - model_config: ClassVar[ConfigDict] = ConfigDict( - extra='allow', - populate_by_name=True, - alias_generator=to_camel, - ) - - version: str | None = Field( - default=None, - description='A specific version of the model family.', - ) - temperature: float | None = Field( - default=None, - ge=0.0, - le=1.0, - description='Controls randomness in token selection (0.0-1.0).', - ) - max_output_tokens: int | None = Field( - default=None, - description='Maximum number of tokens to generate.', - ) - top_k: int | None = Field( - default=None, - description='Maximum number of tokens to consider when sampling.', - ) - top_p: float | None = Field( - default=None, - ge=0.0, - le=1.0, - description='Nucleus sampling probability mass (0.0-1.0).', - ) - stop_sequences: list[str] | None = Field( - default=None, - description='Strings that will stop output generation.', - ) - - -class BedrockConfig(GenkitCommonConfigMixin): - """Base AWS Bedrock configuration for Genkit. - - Combines: - - **GenkitCommonConfigMixin**: Standard Genkit parameters for DevUI compatibility - - **Bedrock Converse API parameters**: For AWS Bedrock API compatibility - - Use model-specific configs (AnthropicConfig, MetaLlamaConfig, etc.) for additional - model-specific parameters. All model configs inherit from this base. - - Official Documentation: - - Converse API: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html - - InferenceConfiguration: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InferenceConfiguration.html - """ - - # Bedrock Converse API standard parameters - # Note: temperature, top_p, stop_sequences are inherited from GenkitCommonConfigMixin - max_tokens: int | None = Field( - default=None, - description='Maximum tokens (Bedrock-style). Use max_output_tokens for Genkit compatibility.', - ) - - -class AmazonNovaConfig(BedrockConfig): - """Configuration for Amazon Nova models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig. - - Supports: - - Nova Pro (text, image, video input) - - Nova Lite (text, image, video input) - - Nova Micro (text only) - - Nova Premier (text, image, video input) - - Nova Canvas (image generation) - - Nova Reel (video generation) - - Nova Sonic / Nova 2 Sonic (speech) - - Official Documentation: - - Amazon Nova: https://docs.aws.amazon.com/nova/latest/userguide/what-is-nova.html - - Request Schema: https://docs.aws.amazon.com/nova/latest/userguide/getting-started-schema.html - - Note: The timeout period for inference calls to Amazon Nova is 60 minutes. - """ - - pass # Nova uses standard Converse API parameters from BedrockConfig - - -class AnthropicToolChoice(StrEnum): - """Tool choice mode for Anthropic models. - - Controls how the model uses tools. - - See: https://docs.anthropic.com/en/api/messages - """ - - AUTO = 'auto' - ANY = 'any' - TOOL = 'tool' - - -class AnthropicEffort(StrEnum): - """Effort level for Claude Opus 4.5 (beta). - - Controls how liberally Claude spends tokens for the best result. - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html - """ - - HIGH = 'high' - MEDIUM = 'medium' - LOW = 'low' - - -class AnthropicConfig(BedrockConfig): - """Configuration for Anthropic Claude models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus Anthropic-specific params. - - Supports: - - Claude Sonnet 4.5 (anthropic.claude-sonnet-4-5-20250929-v1:0) - - Claude Opus 4.5 (anthropic.claude-opus-4-5-20251101-v1:0) - - Claude Opus 4.1 (anthropic.claude-opus-4-1-20250805-v1:0) - - Claude Sonnet 4 (anthropic.claude-sonnet-4-20250514-v1:0) - - Claude Haiku 4.5 (anthropic.claude-haiku-4-5-20251001-v1:0) - - Claude 3.5 Haiku (anthropic.claude-3-5-haiku-20241022-v1:0) - - Claude 3 Haiku (anthropic.claude-3-haiku-20240307-v1:0) - - Official Documentation: - - Anthropic on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html - - Messages API: https://docs.anthropic.com/en/api/messages - - Request/Response: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html - - Note: - - Claude Sonnet 4.5 and Haiku 4.5 support specifying either temperature OR top_p, not both. - - The timeout period for Claude 4 models is 60 minutes. - """ - - # Anthropic-specific parameters - anthropic_version: str | None = Field( - default=None, - description='Anthropic API version. Use "bedrock-2023-05-31" for Messages API.', - ) - anthropic_beta: list[str] | None = Field( - default=None, - description='Beta headers for opt-in features (e.g., "computer-use-2025-01-24", "effort-2025-11-24").', - ) - system: str | None = Field( - default=None, - description='System prompt for context and instructions. Requires Claude 2.1+.', - ) - thinking: dict[str, Any] | None = Field( - default=None, - description='Configuration for enabling Claude extended thinking capability.', - ) - effort: AnthropicEffort | None = Field( - default=None, - description='Effort level for Claude Opus 4.5 (beta). Requires "effort-2025-11-24" beta header.', - ) - tool_choice: dict[str, str] | None = Field( - default=None, - description='How to use tools: {"type": "auto"|"any"|"tool", "name": "tool_name"}.', - ) - - -class AI21JambaConfig(BedrockConfig): - """Configuration for AI21 Labs Jamba models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus Jamba-specific params. - - Supports: Jamba 1.5 Large, Jamba 1.5 Mini (256K context window) - - Official Documentation: - - AI21 on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-jamba.html - - AI21 Jamba API: https://docs.ai21.com/reference/jamba-1-6-api-ref - - Features: - - 256K token context window - - Structured JSON output - - Function calling support - - Multilingual (EN, ES, FR, PT, IT, NL, DE, AR, HE) - - Default Values: - - temperature: 1.0 (range: 0.0-2.0, note: higher max than most models) - - top_p: 1.0 (range: 0.01-1.0) - """ - - # Note: Jamba supports temperature up to 2.0, unlike most models - # The base BedrockConfig limits to 1.0, but extra='allow' lets higher values through - - n: int | None = Field( - default=None, - ge=1, - le=16, - description='Number of chat responses to generate. Default: 1, Max: 16. Note: n > 1 requires temperature > 0.', - ) - frequency_penalty: float | None = Field( - default=None, - description='Reduce frequency of repeated words. Higher values produce fewer repeated words.', - ) - presence_penalty: float | None = Field( - default=None, - description='Reduce repetition of any repeated tokens, applied equally regardless of frequency.', - ) - stop: list[str] | None = Field( - default=None, - description='Stop sequences (up to 64K each). Supports newlines as \\n.', - ) - - -class CohereConfig(BedrockConfig): - """Configuration for Cohere models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus Cohere-specific params. - - Supports: - - Command R+ (cohere.command-r-plus-v1:0) - - Command R (cohere.command-r-v1:0) - - Embed English (cohere.embed-english-v3) - - Embed Multilingual (cohere.embed-multilingual-v3) - - Embed v4 (cohere.embed-v4:0) - text and image - - Rerank 3.5 (cohere.rerank-v3-5:0) - - Official Documentation: - - Cohere on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html - - Cohere Chat API: https://docs.cohere.com/v2/reference/chat - """ - - k: int | None = Field( - default=None, - ge=0, - le=500, - description='Top-k sampling (Cohere-specific). When k=0, k-sampling is disabled. Default: 0.', - ) - p: float | None = Field( - default=None, - ge=0.01, - le=0.99, - description='Nucleus sampling probability mass (Cohere-specific). Default: 0.75.', - ) - safety_mode: CohereSafetyMode | None = Field( - default=None, - description='Safety instruction mode: CONTEXTUAL, STRICT, or OFF.', - ) - tool_choice: CohereToolChoice | None = Field( - default=None, - description='Force tool use: REQUIRED or NONE.', - ) - documents: list[str | dict[str, Any]] | None = Field( - default=None, - description='Documents for RAG-based generation with citations.', - ) - search_queries_only: bool | None = Field( - default=None, - description='Only return search queries without model response.', - ) - preamble: str | None = Field( - default=None, - description='Override default preamble for search query generation.', - ) - prompt_truncation: Literal['OFF', 'AUTO_PRESERVE_ORDER'] | None = Field( - default=None, - description='How to handle prompt truncation when exceeding context.', - ) - frequency_penalty: float | None = Field( - default=None, - ge=0.0, - le=1.0, - description='Reduce repetitiveness proportional to frequency.', - ) - presence_penalty: float | None = Field( - default=None, - ge=0.0, - le=1.0, - description='Reduce repetitiveness for any repeated tokens.', - ) - seed: int | None = Field( - default=None, - description='Seed for deterministic sampling.', - ) - return_prompt: bool | None = Field( - default=None, - description='Return full prompt in response.', - ) - raw_prompting: bool | None = Field( - default=None, - description='Send message without preprocessing.', - ) - - -class DeepSeekConfig(BedrockConfig): - """Configuration for DeepSeek models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus DeepSeek-specific params. - - Supports: DeepSeek-R1, DeepSeek-V3.1 - - Official Documentation: - - DeepSeek on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-deepseek.html - - DeepSeek API: https://api-docs.deepseek.com/api/create-chat-completion - - Reasoning Guide: https://api-docs.deepseek.com/guides/reasoning_model.html - - Important Notes: - - For optimal response quality with DeepSeek-R1, limit max_tokens to 8,192 or fewer. - - While API accepts up to 32,768 tokens, quality degrades above 8,192. - - Response includes reasoning_content for reasoning models (Converse API). - - Must use cross-region inference profile ID (e.g., us.deepseek.r1-v1:0). - - Response Stop Reasons: - - 'stop': Model finished generating - - 'length': Hit max_tokens limit (increase max_tokens) - """ - - stop: list[str] | None = Field( - default=None, - max_length=10, - description='Stop sequences (max 10 items).', - ) - - -class GoogleGemmaConfig(BedrockConfig): - """Configuration for Google Gemma models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig. - - Supports: - - Gemma 3 4B IT (google.gemma-3-4b-it) - text + image - - Gemma 3 12B IT (google.gemma-3-12b-it) - text + image - - Gemma 3 27B IT (google.gemma-3-27b-it) - text + image - - Official Documentation: - - Gemma on Bedrock: Available in model catalog - - Google AI: https://ai.google.dev/gemma - - Note: "IT" = Instruction Tuned variants. - """ - - pass # Gemma uses standard parameters from BedrockConfig - - -class MetaLlamaConfig(BedrockConfig): - """Configuration for Meta Llama models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus Llama-specific params. - - Supports: - - Llama 4 Maverick 17B Instruct (meta.llama4-maverick-17b-instruct-v1:0) - - Llama 4 Scout 17B Instruct (meta.llama4-scout-17b-instruct-v1:0) - - Llama 3.3 70B Instruct (meta.llama3-3-70b-instruct-v1:0) - - Llama 3.2 90B/11B/3B/1B Instruct (multimodal for 90B/11B) - - Llama 3.1 405B/70B/8B Instruct - - Llama 3 70B/8B Instruct - - Official Documentation: - - Meta on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html - - Llama Guide: https://ai.meta.com/llama/get-started/#prompting - - Note: - - Llama 3.2 Instruct and Llama 3.3 Instruct models use geofencing. - - Llama 4 models support streaming. - - Default Values: - - temperature: 0.5 (range: 0-1) - - top_p: 0.9 (range: 0-1) - - max_gen_len: 512 (range: 1-2048) - """ - - max_gen_len: int | None = Field( - default=None, - ge=1, - le=2048, - description='Maximum tokens to generate (Llama naming). Default: 512, Max: 2048.', - ) - images: list[str] | None = Field( - default=None, - description='List of base64-encoded images for Llama 3.2+ multimodal models.', - ) - - -class MiniMaxConfig(BedrockConfig): - """Configuration for MiniMax models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig. - - Supports: MiniMax M2 - - Official Documentation: - - MiniMax on Bedrock: Available in model catalog - """ - - pass # MiniMax uses standard parameters from BedrockConfig - - -class MistralToolChoice(StrEnum): - """Tool choice mode for Mistral models. - - Controls how the model uses tools. - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral-chat-completion.html - """ - - AUTO = 'auto' - ANY = 'any' - NONE = 'none' - - -class MistralConfig(BedrockConfig): - """Configuration for Mistral AI models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus Mistral-specific params. - - Supports: - - Mistral Large 3 (mistral.mistral-large-3-675b-instruct) - - Pixtral Large 25.02 (mistral.pixtral-large-2502-v1:0) - multimodal - - Magistral Small 2509 (mistral.magistral-small-2509) - multimodal - - Ministral 3B/8B/14B (mistral.ministral-3-*-instruct) - - Mixtral 8x7B Instruct (mistral.mixtral-8x7b-instruct-v0:1) - - Mistral 7B Instruct (mistral.mistral-7b-instruct-v0:2) - - Voxtral Mini/Small (mistral.voxtral-*) - speech models - - Official Documentation: - - Mistral on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html - - Chat Completion: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral-chat-completion.html - - Mistral API: https://docs.mistral.ai/capabilities/completion/ - - Note: Mistral models support both text completion and chat completion formats. - - Default Values (Mistral Large): - - temperature: 0.7 (range: 0-1) - - top_p: 1.0 (range: 0-1) - - max_tokens: 8192 - """ - - random_seed: int | None = Field( - default=None, - description='Seed for deterministic sampling. Mistral uses random_seed instead of seed.', - ) - safe_prompt: bool | None = Field( - default=None, - description='Enable safety additions to reduce risky outputs.', - ) - tool_choice: MistralToolChoice | None = Field( - default=None, - description='How to use tools: auto (model decides), any (forced), none (disabled).', - ) - - -class MoonshotConfig(BedrockConfig): - """Configuration for Moonshot AI models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig. - - Supports: Kimi K2 Thinking - - Official Documentation: - - Moonshot on Bedrock: Available in model catalog - """ - - pass # Moonshot uses standard parameters from BedrockConfig - - -class NvidiaConfig(BedrockConfig): - """Configuration for NVIDIA models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus NVIDIA-specific params. - - Supports: - - NVIDIA Nemotron Nano 9B v2 (nvidia.nemotron-nano-9b-v2) - text - - NVIDIA Nemotron Nano 12B v2 VL BF16 (nvidia.nemotron-nano-12b-v2) - text + image - - Official Documentation: - - NVIDIA NIM API: https://docs.nvidia.com/nim/large-language-models/latest/api-reference.html - - Thinking Budget: https://docs.nvidia.com/nim/large-language-models/latest/thinking-budget-control.html - - Note: NVIDIA NIM API is compatible with OpenAI's API format. - """ - - max_thinking_tokens: int | None = Field( - default=None, - description='Max reasoning tokens before final answer. Controls thinking budget for reflection models.', - ) - - -class OpenAIConfig(BedrockConfig): - """Configuration for OpenAI models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig. - - Supports: - - GPT OSS 120B (openai.gpt-oss-120b-1:0) - - GPT OSS 20B (openai.gpt-oss-20b-1:0) - - GPT OSS Safeguard 120B (openai.gpt-oss-safeguard-120b) - - GPT OSS Safeguard 20B (openai.gpt-oss-safeguard-20b) - - Official Documentation: - - OpenAI on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-openai.html - - Note: These are open-source models, not OpenAI's proprietary GPT-4/GPT-3.5. - """ - - pass # OpenAI uses standard parameters from BedrockConfig - - -class QwenConfig(BedrockConfig): - """Configuration for Alibaba Qwen models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus Qwen-specific params. - - Supports: - - Qwen3 32B dense (qwen.qwen3-32b-v1:0) - - Qwen3 235B A22B 2507 (qwen.qwen3-235b-a22b-2507-v1:0) - - Qwen3 Next 80B A3B (qwen.qwen3-next-80b-a3b) - - Qwen3 Coder 480B A35B (qwen.qwen3-coder-480b-a35b-v1:0) - - Qwen3 Coder 30B A3B (qwen.qwen3-coder-30b-a3b-v1:0) - - Qwen3 VL 235B A22B (qwen.qwen3-vl-235b-a22b) - text + image - - Official Documentation: - - Qwen on Bedrock: Available in model catalog - - Qwen API: https://www.alibabacloud.com/help/en/model-studio/qwen-api-reference - - Thinking Mode: https://www.alibabacloud.com/help/en/model-studio/use-qwen-by-calling-api - - Key Features: - - OpenAI-compatible API format - - Advanced reasoning with Thinking mode (Qwen3) - - Multimodal support (text + images for VL models) - - Note: Qwen API is OpenAI-compatible, so standard parameters work. - """ - - enable_thinking: bool | None = Field( - default=None, - description='Enable Thinking mode for enhanced reasoning (Qwen3 models).', - ) - - -class WriterConfig(BedrockConfig): - """Configuration for Writer AI models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig. - - Supports: Palmyra X4 (128K context), Palmyra X5 (1M context) - - Official Documentation: - - Writer on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-writer-palmyra.html - - Writer API: https://dev.writer.com/api-guides/chat-completion - - Bedrock Integration: https://dev.writer.com/providers/aws-bedrock - - Key Features: - - Advanced reasoning and multi-step tool-calling - - Code generation and structured outputs - - Built-in RAG (Retrieval-Augmented Generation) - - Multilingual support (30+ languages) - - Adaptive reasoning for context-based strategy adjustment - - Note: Writer API uses standard OpenAI-compatible parameters. - Temperature defaults to 1 and controls response randomness. - """ - - -class StabilityAspectRatio(StrEnum): - """Aspect ratio options for Stability AI image generation. - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-diffusion-3-5-large.html - """ - - RATIO_16_9 = '16:9' - RATIO_1_1 = '1:1' - RATIO_21_9 = '21:9' - RATIO_2_3 = '2:3' - RATIO_3_2 = '3:2' - RATIO_4_5 = '4:5' - RATIO_5_4 = '5:4' - RATIO_9_16 = '9:16' - RATIO_9_21 = '9:21' - - -class StabilityOutputFormat(StrEnum): - """Output format options for Stability AI image generation. - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-diffusion-3-5-large.html - """ - - JPEG = 'jpeg' - PNG = 'png' - WEBP = 'webp' - - -class StabilityMode(StrEnum): - """Generation mode for Stability AI models. - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-diffusion-3-5-large.html - """ - - TEXT_TO_IMAGE = 'text-to-image' - IMAGE_TO_IMAGE = 'image-to-image' - - -class StabilityConfig(BedrockConfig): - """Configuration for Stability AI models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig. - - Supports: - - Stable Diffusion 3.5 Large (stability.sd3-5-large-v1:0) - - Stable Image Core 1.0 (stability.stable-image-core-v1:1) - - Stable Image Ultra 1.0 (stability.stable-image-ultra-v1:1) - - Stable Image Control (Sketch, Structure) - - Stable Image Editing (Inpaint, Outpaint, Erase, Search and Replace/Recolor) - - Stable Image Upscale (Fast, Creative, Conservative) - - Official Documentation: - - Stability on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-stability-diffusion.html - - SD 3.5 Large: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-diffusion-3-5-large.html - - Stable Image Services: https://docs.aws.amazon.com/bedrock/latest/userguide/stable-image-services.html - - Note: These are image generation/editing models, not text generation. - - Generation Modes: - - text-to-image: Requires only prompt - - image-to-image: Requires prompt, image, and strength - """ - - # Stable Diffusion 3.5 Large parameters - seed: int | None = Field( - default=None, - ge=0, - le=4294967294, - description='Random seed for reproducible generation. 0 = random seed.', - ) - aspect_ratio: StabilityAspectRatio | None = Field( - default=None, - description='Aspect ratio for text-to-image. Default: 1:1.', - ) - mode: StabilityMode | None = Field( - default=None, - description='Generation mode: text-to-image or image-to-image.', - ) - negative_prompt: str | None = Field( - default=None, - max_length=10000, - description='Text describing elements to exclude from the output image.', - ) - output_format: StabilityOutputFormat | None = Field( - default=None, - description='Output image format: jpeg, png, or webp. Default: png.', - ) - image: str | None = Field( - default=None, - description='Base64-encoded input image for image-to-image mode. Min 64px per side.', - ) - strength: float | None = Field( - default=None, - ge=0.0, - le=1.0, - description='Influence of input image (image-to-image). 0 = preserve input, 1 = ignore input.', - ) - - # Legacy parameters for older Stability models (Stable Diffusion XL, etc.) - cfg_scale: float | None = Field( - default=None, - description='(Legacy) How strongly the image should conform to prompt. Range: 0-35.', - ) - steps: int | None = Field( - default=None, - description='(Legacy) Number of diffusion steps. Range: 10-150.', - ) - style_preset: str | None = Field( - default=None, - description='(Legacy) Style preset for image generation.', - ) - - -class TitanConfig(BedrockConfig): - r"""Configuration for Amazon Titan models on AWS Bedrock. - - Inherits all Genkit common parameters from BedrockConfig plus Titan-specific params. - - Supports: - - Titan Text Large (amazon.titan-tg1-large) - - Titan Embeddings G1 - Text (amazon.titan-embed-text-v1) - - Titan Text Embeddings V2 (amazon.titan-embed-text-v2:0) - - Titan Multimodal Embeddings G1 (amazon.titan-embed-image-v1) - - Titan Image Generator G1 v2 (amazon.titan-image-generator-v2:0) - - Official Documentation: - - Titan on Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html - - Default Values: - - temperature: 0.7 (range: 0.0-1.0) - - top_p: 0.9 (range: 0.0-1.0) - - max_token_count: 512 - - Prompt Format: - For conversational responses, use: "User: \nBot:" - """ - - max_token_count: int | None = Field( - default=None, - description='Maximum tokens to generate (Titan naming). Default: 512.', - ) - - -class TextEmbeddingConfig(BaseModel): - """Configuration for text embedding requests. - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html - - Attributes: - dimensions: Output embedding dimensions (model-dependent). - normalize: Whether to normalize the output vector. - input_type: Type of input for Cohere models: search_document, search_query, etc. - """ - - model_config: ClassVar[ConfigDict] = ConfigDict( - extra='allow', - populate_by_name=True, - alias_generator=to_camel, - ) - - dimensions: int | None = Field( - default=None, - description='Output embedding dimensions.', - ) - normalize: bool | None = Field( - default=None, - description='Whether to normalize the output vector.', - ) - input_type: str | None = Field( - default=None, - description='Input type for Cohere: search_document, search_query, classification, clustering.', - ) - - -__all__ = [ - 'AI21JambaConfig', - # Model-Specific Configs (16 providers) - 'AmazonNovaConfig', - 'AnthropicConfig', - # Enums - 'AnthropicEffort', - 'AnthropicToolChoice', - # Base Config - 'BedrockConfig', - 'CohereConfig', - 'CohereSafetyMode', - 'CohereToolChoice', - 'DeepSeekConfig', - # Mixins - 'GenkitCommonConfigMixin', - 'GoogleGemmaConfig', - 'MetaLlamaConfig', - 'MiniMaxConfig', - 'MistralConfig', - 'MistralToolChoice', - 'MoonshotConfig', - 'NvidiaConfig', - 'OpenAIConfig', - 'QwenConfig', - 'StabilityAspectRatio', - 'StabilityConfig', - 'StabilityMode', - 'StabilityOutputFormat', - # Embedding Config - 'TextEmbeddingConfig', - 'TitanConfig', - 'WriterConfig', -] diff --git a/py/plugins/amazon-bedrock/src/genkit/py.typed b/py/plugins/amazon-bedrock/src/genkit/py.typed deleted file mode 100644 index 8b13789179..0000000000 --- a/py/plugins/amazon-bedrock/src/genkit/py.typed +++ /dev/null @@ -1 +0,0 @@ - diff --git a/py/plugins/amazon-bedrock/tests/amazon_bedrock_converters_test.py b/py/plugins/amazon-bedrock/tests/amazon_bedrock_converters_test.py deleted file mode 100644 index 7165f954ea..0000000000 --- a/py/plugins/amazon-bedrock/tests/amazon_bedrock_converters_test.py +++ /dev/null @@ -1,742 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Amazon Bedrock format conversion utilities. - -Covers finish reason mapping, role conversion, system message extraction, -content block conversion (to/from Bedrock), tool definitions, config -normalization, JSON instructions, media handling, and inference profile -ID resolution. -""" - -import base64 - -from genkit.plugins.amazon_bedrock.models.converters import ( - FINISH_REASON_MAP, - INFERENCE_PROFILE_PREFIXES, - INFERENCE_PROFILE_SUPPORTED_PROVIDERS, - StreamingFenceStripper, - build_json_instruction, - build_media_block, - build_usage, - convert_media_data_uri, - from_bedrock_content, - get_effective_model_id, - is_image_media, - map_finish_reason, - maybe_strip_fences, - normalize_config, - parse_tool_call_args, - separate_system_messages, - strip_markdown_fences, - to_bedrock_content, - to_bedrock_role, - to_bedrock_tool, -) -from genkit.plugins.amazon_bedrock.typing import BedrockConfig -from genkit.types import ( - FinishReason, - GenerateRequest, - GenerationCommonConfig, - Media, - MediaPart, - Message, - OutputConfig, - Part, - Role, - TextPart, - ToolDefinition, - ToolRequest, - ToolRequestPart, - ToolResponse, - ToolResponsePart, -) - - -class TestMapFinishReason: - """Tests for finish reason mapping.""" - - def test_end_turn_maps_to_stop(self) -> None: - """Test End turn maps to stop.""" - got = map_finish_reason('end_turn') - assert got == FinishReason.STOP, f'map_finish_reason("end_turn") = {got}, want STOP' - - def test_stop_sequence_maps_to_stop(self) -> None: - """Test Stop sequence maps to stop.""" - got = map_finish_reason('stop_sequence') - assert got == FinishReason.STOP, f'map_finish_reason("stop_sequence") = {got}, want STOP' - - def test_max_tokens_maps_to_length(self) -> None: - """Test Max tokens maps to length.""" - got = map_finish_reason('max_tokens') - assert got == FinishReason.LENGTH, f'map_finish_reason("max_tokens") = {got}, want LENGTH' - - def test_tool_use_maps_to_stop(self) -> None: - """Test Tool use maps to stop.""" - got = map_finish_reason('tool_use') - assert got == FinishReason.STOP, f'map_finish_reason("tool_use") = {got}, want STOP' - - def test_content_filtered_maps_to_blocked(self) -> None: - """Test Content filtered maps to blocked.""" - got = map_finish_reason('content_filtered') - assert got == FinishReason.BLOCKED, f'map_finish_reason("content_filtered") = {got}, want BLOCKED' - - def test_guardrail_intervened_maps_to_blocked(self) -> None: - """Test Guardrail intervened maps to blocked.""" - got = map_finish_reason('guardrail_intervened') - assert got == FinishReason.BLOCKED, f'map_finish_reason("guardrail_intervened") = {got}, want BLOCKED' - - def test_unknown_reason_maps_to_unknown(self) -> None: - """Test Unknown reason maps to unknown.""" - got = map_finish_reason('something_new') - assert got == FinishReason.UNKNOWN, f'map_finish_reason("something_new") = {got}, want UNKNOWN' - - def test_empty_string_maps_to_unknown(self) -> None: - """Test Empty string maps to unknown.""" - got = map_finish_reason('') - assert got == FinishReason.UNKNOWN, f'map_finish_reason("") = {got}, want UNKNOWN' - - def test_finish_reason_map_is_complete(self) -> None: - """Ensure the constant covers all expected Bedrock stop reasons.""" - expected_keys = { - 'end_turn', - 'stop_sequence', - 'max_tokens', - 'tool_use', - 'content_filtered', - 'guardrail_intervened', - } - assert FINISH_REASON_MAP.keys() == expected_keys, ( - f'FINISH_REASON_MAP keys = {set(FINISH_REASON_MAP.keys())}, want {expected_keys}' - ) - - -class TestToBedrockRole: - """Tests for Genkit β†’ Bedrock role conversion.""" - - def test_user_role_enum(self) -> None: - """Test User role enum.""" - assert to_bedrock_role(Role.USER) == 'user' - - def test_model_role_enum(self) -> None: - """Test Model role enum.""" - assert to_bedrock_role(Role.MODEL) == 'assistant' - - def test_tool_role_enum(self) -> None: - """Test Tool role enum.""" - assert to_bedrock_role(Role.TOOL) == 'user' - - def test_user_string(self) -> None: - """Test User string.""" - assert to_bedrock_role('user') == 'user' - - def test_model_string(self) -> None: - """Test Model string.""" - assert to_bedrock_role('model') == 'assistant' - - def test_assistant_string(self) -> None: - """Test Assistant string.""" - assert to_bedrock_role('assistant') == 'assistant' - - def test_tool_string(self) -> None: - """Test Tool string.""" - assert to_bedrock_role('tool') == 'user' - - def test_unknown_string_defaults_to_user(self) -> None: - """Test Unknown string defaults to user.""" - assert to_bedrock_role('admin') == 'user' - - def test_case_insensitive_string(self) -> None: - """Test Case insensitive string.""" - assert to_bedrock_role('MODEL') == 'assistant' - - -class TestSeparateSystemMessages: - """Tests for system message extraction.""" - - def test_no_messages(self) -> None: - """Test No messages.""" - system, conv = separate_system_messages([]) - assert not (system or conv) - - def test_no_system_messages(self) -> None: - """Test No system messages.""" - msgs = [Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))])] - system, conv = separate_system_messages(msgs) - assert not (system), f'Expected no system messages, got {system}' - assert len(conv) == 1, f'Expected 1 conversation message, got {len(conv)}' - - def test_single_system_message(self) -> None: - """Test Single system message.""" - msgs = [ - Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='Be helpful.'))]), - Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))]), - ] - system, conv = separate_system_messages(msgs) - assert system == ['Be helpful.'], f'system = {system}, want ["Be helpful."]' - assert len(conv) == 1, f'Expected 1 conversation message, got {len(conv)}' - - def test_multiple_system_messages(self) -> None: - """Test Multiple system messages.""" - msgs = [ - Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='Rule 1'))]), - Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))]), - Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='Rule 2'))]), - ] - system, conv = separate_system_messages(msgs) - assert system == ['Rule 1', 'Rule 2'], f'system = {system}' - assert len(conv) == 1, f'Expected 1 conversation message, got {len(conv)}' - - def test_multi_part_system_message(self) -> None: - """Test Multi part system message.""" - msgs = [ - Message( - role=Role.SYSTEM, - content=[ - Part(root=TextPart(text='Part A')), - Part(root=TextPart(text='Part B')), - ], - ), - ] - system, conv = separate_system_messages(msgs) - assert system == ['Part APart B'], f'system = {system}' - - def test_preserves_conversation_order(self) -> None: - """Test Preserves conversation order.""" - msgs = [ - Message(role=Role.USER, content=[Part(root=TextPart(text='Q1'))]), - Message(role=Role.MODEL, content=[Part(root=TextPart(text='A1'))]), - Message(role=Role.USER, content=[Part(root=TextPart(text='Q2'))]), - ] - system, conv = separate_system_messages(msgs) - assert len(conv) == 3, f'Expected 3 conversation messages, got {len(conv)}' - roles = [m.role for m in conv] - assert roles == [Role.USER, Role.MODEL, Role.USER], f'Conversation roles = {roles}' - - -class TestToBedrockTool: - """Tests for Genkit β†’ Bedrock tool definition conversion.""" - - def test_basic_tool(self) -> None: - """Test Basic tool.""" - tool = ToolDefinition( - name='get_weather', - description='Fetch current weather', - input_schema={'type': 'object', 'properties': {'city': {'type': 'string'}}}, - ) - got = to_bedrock_tool(tool) - want = { - 'toolSpec': { - 'name': 'get_weather', - 'description': 'Fetch current weather', - 'inputSchema': { - 'json': {'type': 'object', 'properties': {'city': {'type': 'string'}}}, - }, - }, - } - assert got == want, f'got {got}, want {want}' - - def test_tool_without_schema(self) -> None: - """Test Tool without schema.""" - tool = ToolDefinition(name='ping', description='Ping service') - got = to_bedrock_tool(tool) - expected = {'type': 'object', 'properties': {}} - assert got['toolSpec']['inputSchema']['json'] == expected, ( - f'Expected default schema, got {got["toolSpec"]["inputSchema"]}' - ) - - def test_tool_with_empty_description(self) -> None: - """Test Tool with empty description.""" - tool = ToolDefinition(name='noop', description='') - got = to_bedrock_tool(tool) - assert got['toolSpec']['description'] == '', f'Expected empty description, got {got["toolSpec"]["description"]}' - - -class TestToBedrockContent: - """Tests for Genkit Part β†’ Bedrock content block conversion.""" - - def test_text_part(self) -> None: - """Test Text part.""" - part = Part(root=TextPart(text='Hello world')) - got = to_bedrock_content(part) - assert got == {'text': 'Hello world'}, f'got {got}' - - def test_tool_request_part(self) -> None: - """Test Tool request part.""" - part = Part( - root=ToolRequestPart(tool_request=ToolRequest(ref='call-1', name='get_weather', input={'city': 'London'})) - ) - got = to_bedrock_content(part) - want = { - 'toolUse': { - 'toolUseId': 'call-1', - 'name': 'get_weather', - 'input': {'city': 'London'}, - }, - } - assert got == want, f'got {got}, want {want}' - - def test_tool_request_part_without_ref(self) -> None: - """Test Tool request part without ref.""" - part = Part(root=ToolRequestPart(tool_request=ToolRequest(name='ping', input={}))) - got = to_bedrock_content(part) - assert got is not None, 'Expected non-None result' - assert got['toolUse']['toolUseId'] == '', f'Expected empty toolUseId, got {got["toolUse"]["toolUseId"]}' - - def test_tool_response_part_string_output(self) -> None: - """Test Tool response part string output.""" - part = Part(root=ToolResponsePart(tool_response=ToolResponse(ref='call-1', name='get_weather', output='Sunny'))) - got = to_bedrock_content(part) - want = { - 'toolResult': { - 'toolUseId': 'call-1', - 'content': [{'text': 'Sunny'}], - }, - } - assert got == want, f'got {got}, want {want}' - - def test_tool_response_part_dict_output(self) -> None: - """Test Tool response part dict output.""" - part = Part( - root=ToolResponsePart(tool_response=ToolResponse(ref='call-1', name='get_weather', output={'temp': 20})) - ) - got = to_bedrock_content(part) - assert got is not None, 'Expected non-None result' - assert got['toolResult']['content'] == [{'json': {'temp': 20}}], f'got {got}' - - def test_media_part_returns_none(self) -> None: - """Test Media part returns none.""" - part = Part(root=MediaPart(media=Media(url='https://example.com/img.png', content_type='image/png'))) - got = to_bedrock_content(part) - assert got is None, f'Expected None for MediaPart, got {got}' - - -class TestFromBedrockContent: - """Tests for Bedrock content block β†’ Genkit Part conversion.""" - - def test_text_block(self) -> None: - """Test Text block.""" - parts = from_bedrock_content([{'text': 'Hello'}]) - assert len(parts) == 1, f'Expected 1 part, got {len(parts)}' - root = parts[0].root - assert isinstance(root, TextPart), f'Expected TextPart, got {type(root)}' - assert root.text == 'Hello' - - def test_tool_use_block(self) -> None: - """Test Tool use block.""" - parts = from_bedrock_content([ - { - 'toolUse': { - 'toolUseId': 'abc-123', - 'name': 'search', - 'input': {'query': 'test'}, - } - } - ]) - assert len(parts) == 1, f'Expected 1 part, got {len(parts)}' - root = parts[0].root - assert isinstance(root, ToolRequestPart), f'Expected ToolRequestPart, got {type(root)}' - assert root.tool_request.name == 'search', f'tool name = {root.tool_request.name}' - assert root.tool_request.ref == 'abc-123', f'tool ref = {root.tool_request.ref}' - - def test_reasoning_content_string(self) -> None: - """Test Reasoning content string.""" - parts = from_bedrock_content([ - { - 'reasoningContent': { - 'reasoningText': 'Let me think...', - } - } - ]) - assert len(parts) == 1 - root = parts[0].root - assert isinstance(root, TextPart), f'Expected TextPart, got {type(root)}' - assert '[Reasoning]' in root.text or 'Let me think' not in root.text, f'text = {root.text}' - - def test_reasoning_content_dict(self) -> None: - """Test Reasoning content dict.""" - parts = from_bedrock_content([ - { - 'reasoningContent': { - 'reasoningText': {'text': 'Step 1: analyze'}, - } - } - ]) - assert len(parts) == 1 - root = parts[0].root - assert isinstance(root, TextPart), f'Expected TextPart, got {type(root)}' - assert 'Step 1: analyze' in root.text, f'text = {root.text}' - - def test_multiple_blocks(self) -> None: - """Test Multiple blocks.""" - parts = from_bedrock_content([ - {'text': 'Result:'}, - {'toolUse': {'toolUseId': 'x', 'name': 'calc', 'input': {}}}, - ]) - assert len(parts) == 2, f'Expected 2 parts, got {len(parts)}' - - def test_empty_blocks(self) -> None: - """Test Empty blocks.""" - parts = from_bedrock_content([]) - assert len(parts) == 0 - - -class TestParseToolCallArgs: - """Tests for tool call argument JSON parsing.""" - - def test_valid_json(self) -> None: - """Test Valid json.""" - got = parse_tool_call_args('{"x": 1}') - assert got == {'x': 1}, f'got {got}' - - def test_invalid_json_returns_string(self) -> None: - """Test Invalid json returns string.""" - got = parse_tool_call_args('not json') - assert got == 'not json', f'got {got}' - - def test_empty_string_returns_empty_dict(self) -> None: - """Test Empty string returns empty dict.""" - got = parse_tool_call_args('') - assert got == {}, f'got {got}' - - def test_nested_json(self) -> None: - """Test Nested json.""" - got = parse_tool_call_args('{"a": {"b": [1, 2]}}') - assert got == {'a': {'b': [1, 2]}}, f'got {got}' - - -class TestBuildUsage: - """Tests for usage statistics construction.""" - - def test_full_usage(self) -> None: - """Test Full usage.""" - got = build_usage({'inputTokens': 10, 'outputTokens': 20, 'totalTokens': 30}) - assert got.input_tokens == 10 or got.output_tokens != 20 or got.total_tokens != 30, f'got {got}' - - def test_missing_fields_default_to_zero(self) -> None: - """Test Missing fields default to zero.""" - got = build_usage({}) - assert got.input_tokens == 0 or got.output_tokens != 0 or got.total_tokens != 0, f'got {got}' - - def test_partial_usage(self) -> None: - """Test Partial usage.""" - got = build_usage({'inputTokens': 5}) - assert got.input_tokens == 5 or got.output_tokens != 0, f'got {got}' - - -class TestNormalizeConfig: - """Tests for config normalization.""" - - def test_none_returns_default(self) -> None: - """Test None returns default.""" - got = normalize_config(None) - assert isinstance(got, BedrockConfig), f'Expected BedrockConfig, got {type(got)}' - - def test_bedrock_config_passthrough(self) -> None: - """Test Bedrock config passthrough.""" - config = BedrockConfig(temperature=0.5) - got = normalize_config(config) - assert got is config, 'Expected same instance' - - def test_generation_common_config(self) -> None: - """Test Generation common config.""" - config = GenerationCommonConfig(temperature=0.7, max_output_tokens=100, top_p=0.9) - got = normalize_config(config) - assert got.temperature == 0.7, f'temperature = {got.temperature}' - assert got.max_tokens == 100, f'max_tokens = {got.max_tokens}' - assert got.top_p == 0.9, f'top_p = {got.top_p}' - - def test_dict_with_camel_case_keys(self) -> None: - """Test Dict with camel case keys.""" - config = {'maxOutputTokens': 200, 'topP': 0.8} - got = normalize_config(config) - assert got.max_tokens == 200, f'max_tokens = {got.max_tokens}' - assert got.top_p == 0.8, f'top_p = {got.top_p}' - - def test_dict_with_snake_case_keys(self) -> None: - """Test Dict with snake case keys.""" - config = {'temperature': 0.5, 'stop_sequences': ['END']} - got = normalize_config(config) - assert got.temperature == 0.5, f'temperature = {got.temperature}' - assert got.stop_sequences == ['END'], f'stop_sequences = {got.stop_sequences}' - - def test_unknown_type_returns_default(self) -> None: - """Test Unknown type returns default.""" - got = normalize_config(42) - assert isinstance(got, BedrockConfig), f'Expected BedrockConfig, got {type(got)}' - - -class TestBuildJsonInstruction: - """Tests for JSON output instruction generation.""" - - def test_no_output_returns_none(self) -> None: - """Test No output returns none.""" - request = GenerateRequest(messages=[]) - got = build_json_instruction(request) - assert got is None, f'Expected None, got {got}' - - def test_text_format_returns_none(self) -> None: - """Test Text format returns none.""" - request = GenerateRequest(messages=[], output=OutputConfig(format='text')) - got = build_json_instruction(request) - assert got is None, f'Expected None, got {got}' - - def test_json_format_without_schema(self) -> None: - """Test Json format without schema.""" - request = GenerateRequest(messages=[], output=OutputConfig(format='json')) - got = build_json_instruction(request) - assert got is not None, 'Expected non-None instruction' - assert 'valid JSON' in got, f'Missing JSON instruction: {got}' - - def test_json_format_with_schema(self) -> None: - """Test Json format with schema.""" - schema = {'type': 'object', 'properties': {'name': {'type': 'string'}}} - request = GenerateRequest(messages=[], output=OutputConfig(format='json', schema=schema)) - got = build_json_instruction(request) - assert got is not None, 'Expected non-None instruction' - assert 'name' in got, f'Schema not in instruction: {got}' - - -class TestConvertMediaDataUri: - """Tests for data URI media parsing.""" - - def test_png_data_uri(self) -> None: - """Test Png data uri.""" - png_data = base64.b64encode(b'\x89PNG').decode('ascii') - media = Media(url=f'data:image/png;base64,{png_data}', content_type='image/png') - media_bytes, format_str, is_data = convert_media_data_uri(media) - assert is_data, 'Expected is_data_uri=True' - assert format_str == 'png', f'format = {format_str}' - assert media_bytes == b'\x89PNG', f'bytes = {media_bytes}' - - def test_http_url_returns_false(self) -> None: - """Test Http url returns false.""" - media = Media(url='https://example.com/img.jpg', content_type='image/jpeg') - _, _, is_data = convert_media_data_uri(media) - assert not (is_data), 'Expected is_data_uri=False for HTTP URL' - - def test_data_uri_without_comma(self) -> None: - """Test Data uri without comma.""" - media = Media(url='data:image/png;base64', content_type='image/png') - _, _, is_data = convert_media_data_uri(media) - assert not (is_data), 'Expected is_data_uri=False for malformed data URI' - - -class TestIsImageMedia: - """Tests for image vs video classification.""" - - def test_image_content_type(self) -> None: - """Test Image content type.""" - assert is_image_media('image/png', '') - - def test_video_content_type(self) -> None: - """Test Video content type.""" - assert not (is_image_media('video/mp4', '')) - - def test_image_url_extension(self) -> None: - """Test Image url extension.""" - assert is_image_media('', 'photo.jpg') - - def test_video_url_no_image_ext(self) -> None: - """Test Video url no image ext.""" - assert not (is_image_media('', 'video.mp4')), 'mp4 URL without content type should not be image' - - def test_no_content_type_no_ext(self) -> None: - """Test No content type no ext.""" - # No image extension β†’ defaults to False - assert not (is_image_media('', 'blob')) - - -class TestBuildMediaBlock: - """Tests for Bedrock media block construction.""" - - def test_image_block(self) -> None: - """Test Image block.""" - got = build_media_block(b'\x89PNG', 'png', is_image=True) - assert 'image' in got, f'Expected image key, got {got}' - assert got['image']['format'] == 'png' - - def test_video_block(self) -> None: - """Test Video block.""" - got = build_media_block(b'\x00', 'mp4', is_image=False) - assert 'video' in got, f'Expected video key, got {got}' - assert got['video']['format'] == 'mp4' - - -class TestGetEffectiveModelId: - """Tests for inference profile ID resolution.""" - - def test_already_prefixed_returns_unchanged(self) -> None: - """Test Already prefixed returns unchanged.""" - got = get_effective_model_id('us.anthropic.claude-v3', bearer_token='tok', aws_region='us-east-1') - assert got == 'us.anthropic.claude-v3', f'got {got}' - - def test_no_bearer_token_returns_unchanged(self) -> None: - """Test No bearer token returns unchanged.""" - got = get_effective_model_id('anthropic.claude-v3', bearer_token=None, aws_region='us-east-1') - assert got == 'anthropic.claude-v3', f'got {got}' - - def test_unsupported_provider_returns_unchanged(self) -> None: - """Test Unsupported provider returns unchanged.""" - got = get_effective_model_id('stability.sd3', bearer_token='tok', aws_region='us-east-1') - assert got == 'stability.sd3', f'got {got}' - - def test_no_region_returns_unchanged(self) -> None: - """Test No region returns unchanged.""" - got = get_effective_model_id('anthropic.claude-v3', bearer_token='tok', aws_region=None) - assert got == 'anthropic.claude-v3', f'got {got}' - - def test_us_region_adds_us_prefix(self) -> None: - """Test Us region adds us prefix.""" - got = get_effective_model_id('anthropic.claude-v3', bearer_token='tok', aws_region='us-east-1') - assert got == 'us.anthropic.claude-v3', f'got {got}' - - def test_eu_region_adds_eu_prefix(self) -> None: - """Test Eu region adds eu prefix.""" - got = get_effective_model_id('meta.llama3', bearer_token='tok', aws_region='eu-west-1') - assert got == 'eu.meta.llama3', f'got {got}' - - def test_ap_region_adds_apac_prefix(self) -> None: - """Test Ap region adds apac prefix.""" - got = get_effective_model_id('cohere.command-r', bearer_token='tok', aws_region='ap-southeast-1') - assert got == 'apac.cohere.command-r', f'got {got}' - - def test_unknown_region_defaults_to_us(self) -> None: - """Test Unknown region defaults to us.""" - got = get_effective_model_id('anthropic.claude-v3', bearer_token='tok', aws_region='xx-central-1') - assert got == 'us.anthropic.claude-v3', f'got {got}' - - def test_inference_profile_prefixes_constant(self) -> None: - """Test Inference profile prefixes constant.""" - expected_prefixes = ('us.', 'eu.', 'apac.') - assert INFERENCE_PROFILE_PREFIXES == expected_prefixes, ( - f'INFERENCE_PROFILE_PREFIXES = {INFERENCE_PROFILE_PREFIXES}' - ) - - def test_supported_providers_includes_anthropic(self) -> None: - """Test Supported providers includes anthropic.""" - assert 'anthropic.' in INFERENCE_PROFILE_SUPPORTED_PROVIDERS - - -class TestStripMarkdownFences: - """Tests for strip_markdown_fences.""" - - def test_strips_json_fences(self) -> None: - """Strips ```json ... ``` fences.""" - text = '```json\n{"name": "John"}\n```' - assert strip_markdown_fences(text) == '{"name": "John"}' - - def test_strips_plain_fences(self) -> None: - """Strips ``` ... ``` fences without language tag.""" - text = '```\n{"a": 1}\n```' - assert strip_markdown_fences(text) == '{"a": 1}' - - def test_preserves_plain_json(self) -> None: - """Does not alter valid JSON without fences.""" - text = '{"name": "John"}' - assert strip_markdown_fences(text) == text - - def test_preserves_non_json_text(self) -> None: - """Does not alter plain text.""" - text = 'Hello, world!' - assert strip_markdown_fences(text) == text - - def test_strips_multiline_json(self) -> None: - """Strips fences around multiline JSON.""" - text = '```json\n{\n "a": 1\n}\n```' - assert strip_markdown_fences(text) == '{\n "a": 1\n}' - - -class TestMaybeStripFences: - """Tests for maybe_strip_fences.""" - - def test_strips_fences_for_json_output(self) -> None: - """Strips markdown fences when JSON output is requested.""" - request = GenerateRequest( - messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='json', schema={'type': 'object'}), - ) - parts = [Part(root=TextPart(text='```json\n{"a": 1}\n```'))] - result = maybe_strip_fences(request, parts) - assert result[0].root.text == '{"a": 1}' - - def test_no_op_for_text_output(self) -> None: - """Does not modify responses when output format is not json.""" - request = GenerateRequest( - messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='text'), - ) - fenced = '```json\n{"a": 1}\n```' - parts = [Part(root=TextPart(text=fenced))] - result = maybe_strip_fences(request, parts) - assert result[0].root.text == fenced - - def test_no_op_when_no_fences(self) -> None: - """Does not modify clean JSON responses.""" - request = GenerateRequest( - messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='json', schema={'type': 'object'}), - ) - text = '{"name": "John"}' - parts = [Part(root=TextPart(text=text))] - result = maybe_strip_fences(request, parts) - assert result is parts - - -class TestStreamingFenceStripper: - """Tests for StreamingFenceStripper.""" - - def test_strips_fences_across_chunks(self) -> None: - """Strips opening and closing fences split across chunks.""" - stripper = StreamingFenceStripper(json_mode=True) - chunks = ['```json\n', '{"name":', ' "John"}\n', '```'] - out = [stripper.process(c) for c in chunks] - out.append(stripper.flush()) - combined = ''.join(out) - assert combined == '{"name": "John"}' - - def test_strips_fence_in_single_chunk(self) -> None: - """Strips fences when entire response is one chunk.""" - stripper = StreamingFenceStripper(json_mode=True) - out = stripper.process('```json\n{"a": 1}\n```') - out += stripper.flush() - assert out == '{"a": 1}' - - def test_no_op_when_not_json_mode(self) -> None: - """Passes text through unchanged when json_mode is False.""" - stripper = StreamingFenceStripper(json_mode=False) - text = '```json\n{"a": 1}\n```' - assert stripper.process(text) == text - assert stripper.flush() == '' - - def test_no_fence_passes_through(self) -> None: - """Passes text through when no fence is detected.""" - stripper = StreamingFenceStripper(json_mode=True) - chunks = ['{"name":', ' "John"}'] - out = [stripper.process(c) for c in chunks] - out.append(stripper.flush()) - combined = ''.join(out) - assert combined == '{"name": "John"}' - - def test_buffers_small_prefix(self) -> None: - """Buffers small initial chunks until fence can be detected.""" - stripper = StreamingFenceStripper(json_mode=True) - # First chunk is small β€” should be buffered. - assert stripper.process('```') == '' - # Second chunk triggers flush with fence detection. - assert stripper.process('json\n{"a":') == '{"a":' - assert stripper.process(' 1}\n```') == ' 1}' - assert stripper.flush() == '' diff --git a/py/plugins/amazon-bedrock/tests/amazon_bedrock_plugin_test.py b/py/plugins/amazon-bedrock/tests/amazon_bedrock_plugin_test.py deleted file mode 100644 index edb569f6a8..0000000000 --- a/py/plugins/amazon-bedrock/tests/amazon_bedrock_plugin_test.py +++ /dev/null @@ -1,1059 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Unit tests for AWS Bedrock plugin. - -This module tests the AWS Bedrock plugin functionality including: -- Plugin initialization -- Model naming utilities -- Configuration schema mapping -- Message conversion -- Model info registry -""" - -from collections.abc import AsyncIterator -from typing import Any -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from genkit.core.action import ActionRunContext -from genkit.plugins.amazon_bedrock import ( - AMAZON_BEDROCK_PLUGIN_NAME, - AnthropicConfig, - BedrockConfig, - CohereConfig, - DeepSeekConfig, - GenkitCommonConfigMixin, - MetaLlamaConfig, - MistralConfig, - bedrock_model, - bedrock_name, - claude_sonnet_4_5, - deepseek_r1, - get_config_schema_for_model, - get_inference_profile_prefix, - inference_profile, - llama_3_3_70b, - mistral_large_3, - nova_pro, -) -from genkit.plugins.amazon_bedrock.models.model import BedrockModel -from genkit.plugins.amazon_bedrock.models.model_info import ( - SUPPORTED_BEDROCK_MODELS, - SUPPORTED_EMBEDDING_MODELS, - get_model_info, -) -from genkit.plugins.amazon_bedrock.plugin import _strip_inference_profile_prefix -from genkit.plugins.amazon_bedrock.typing import ( - AI21JambaConfig, - AmazonNovaConfig, - CohereSafetyMode, - CohereToolChoice, - StabilityAspectRatio, - StabilityConfig, - StabilityMode, - StabilityOutputFormat, -) -from genkit.types import GenerateRequest, GenerateResponseChunk, Message, Part, Role, TextPart, ToolRequest - - -class TestBedrockNaming: - """Tests for model naming utilities.""" - - def test_bedrock_name_basic(self) -> None: - """Test bedrock_name creates fully qualified names.""" - result = bedrock_name('anthropic.claude-sonnet-4-5-20250929-v1:0') - assert result == 'amazon-bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_bedrock_model_alias(self) -> None: - """Test bedrock_model is an alias for bedrock_name.""" - result = bedrock_model('meta.llama3-3-70b-instruct-v1:0') - assert result == 'amazon-bedrock/meta.llama3-3-70b-instruct-v1:0' - - def test_predefined_model_references(self) -> None: - """Test predefined model references are correctly formatted with direct model IDs.""" - # Pre-defined references use direct model IDs (work with IAM credentials) - assert claude_sonnet_4_5 == 'amazon-bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0' - assert nova_pro == 'amazon-bedrock/amazon.nova-pro-v1:0' - assert llama_3_3_70b == 'amazon-bedrock/meta.llama3-3-70b-instruct-v1:0' - assert mistral_large_3 == 'amazon-bedrock/mistral.mistral-large-3-675b-instruct' - assert deepseek_r1 == 'amazon-bedrock/deepseek.r1-v1:0' - - def test_plugin_name_constant(self) -> None: - """Test plugin name constant.""" - assert AMAZON_BEDROCK_PLUGIN_NAME == 'amazon-bedrock' - - -class TestConfigSchemaMapping: - """Tests for configuration schema mapping.""" - - def test_anthropic_model_gets_anthropic_config(self) -> None: - """Test Anthropic models get AnthropicConfig.""" - config_class = get_config_schema_for_model('anthropic.claude-sonnet-4-5-20250929-v1:0') - assert config_class == AnthropicConfig - - def test_anthropic_inference_profile_gets_anthropic_config(self) -> None: - """Test Anthropic inference profiles (with us. prefix) get AnthropicConfig.""" - config_class = get_config_schema_for_model('us.anthropic.claude-sonnet-4-5-20250929-v1:0') - assert config_class == AnthropicConfig - - def test_meta_model_gets_llama_config(self) -> None: - """Test Meta models get MetaLlamaConfig.""" - config_class = get_config_schema_for_model('meta.llama3-3-70b-instruct-v1:0') - assert config_class == MetaLlamaConfig - - def test_meta_inference_profile_gets_llama_config(self) -> None: - """Test Meta inference profiles get MetaLlamaConfig.""" - config_class = get_config_schema_for_model('us.meta.llama3-3-70b-instruct-v1:0') - assert config_class == MetaLlamaConfig - - def test_mistral_model_gets_mistral_config(self) -> None: - """Test Mistral models get MistralConfig.""" - config_class = get_config_schema_for_model('mistral.mistral-large-3-675b-instruct') - assert config_class == MistralConfig - - def test_cohere_model_gets_cohere_config(self) -> None: - """Test Cohere models get CohereConfig.""" - config_class = get_config_schema_for_model('cohere.command-r-plus-v1:0') - assert config_class == CohereConfig - - def test_deepseek_model_gets_deepseek_config(self) -> None: - """Test DeepSeek models get DeepSeekConfig.""" - config_class = get_config_schema_for_model('deepseek.r1-v1:0') - assert config_class == DeepSeekConfig - - def test_deepseek_inference_profile_gets_deepseek_config(self) -> None: - """Test DeepSeek inference profiles get DeepSeekConfig.""" - config_class = get_config_schema_for_model('us.deepseek.r1-v1:0') - assert config_class == DeepSeekConfig - - def test_amazon_nova_gets_nova_config(self) -> None: - """Test Amazon Nova models get AmazonNovaConfig.""" - config_class = get_config_schema_for_model('amazon.nova-pro-v1:0') - assert config_class == AmazonNovaConfig - - def test_amazon_nova_inference_profile_gets_nova_config(self) -> None: - """Test Amazon Nova inference profiles get AmazonNovaConfig.""" - config_class = get_config_schema_for_model('us.amazon.nova-pro-v1:0') - assert config_class == AmazonNovaConfig - - def test_unknown_model_gets_base_config(self) -> None: - """Test unknown models get base BedrockConfig.""" - config_class = get_config_schema_for_model('unknown.model-v1:0') - assert config_class == BedrockConfig - - -class TestModelInfo: - """Tests for model info registry.""" - - def test_supported_models_not_empty(self) -> None: - """Test that supported models registry is populated.""" - assert len(SUPPORTED_BEDROCK_MODELS) > 0 - - def test_supported_embeddings_not_empty(self) -> None: - """Test that embedding models registry is populated.""" - assert len(SUPPORTED_EMBEDDING_MODELS) > 0 - - def test_get_known_model_info(self) -> None: - """Test getting info for a known model.""" - model_info = get_model_info('anthropic.claude-sonnet-4-5-20250929-v1:0') - assert model_info.label == 'Claude Sonnet 4.5' - assert model_info.supports is not None - assert model_info.supports.multiturn is True - assert model_info.supports.tools is True - - def test_get_unknown_model_info(self) -> None: - """Test getting info for an unknown model returns default.""" - model_info = get_model_info('unknown.model-v1:0') - assert model_info.label is not None - assert 'unknown.model-v1:0' in model_info.label - assert model_info.supports is not None - - def test_claude_models_support_media(self) -> None: - """Test Claude models support media (images).""" - model_info = get_model_info('anthropic.claude-sonnet-4-5-20250929-v1:0') - assert model_info.supports is not None - assert model_info.supports.media is True - - def test_nova_models_support_media(self) -> None: - """Test Nova Pro/Lite models support media.""" - model_info = get_model_info('amazon.nova-pro-v1:0') - assert model_info.supports is not None - assert model_info.supports.media is True - - def test_nova_micro_text_only(self) -> None: - """Test Nova Micro is text-only.""" - model_info = get_model_info('amazon.nova-micro-v1:0') - assert model_info.supports is not None - assert model_info.supports.media is False - - def test_deepseek_r1_no_tools(self) -> None: - """Test DeepSeek R1 doesn't support tools in Bedrock.""" - model_info = get_model_info('deepseek.r1-v1:0') - assert model_info.supports is not None - assert model_info.supports.tools is False - - -class TestConfigTypes: - """Tests for configuration type classes.""" - - def test_bedrock_config_inherits_genkit_mixin(self) -> None: - """Test BedrockConfig inherits from GenkitCommonConfigMixin.""" - assert issubclass(BedrockConfig, GenkitCommonConfigMixin) - - def test_bedrock_config_default_values(self) -> None: - """Test BedrockConfig has expected default values.""" - config = BedrockConfig() - assert config.temperature is None - assert config.max_tokens is None - assert config.top_p is None - assert config.stop_sequences is None - - def test_bedrock_config_with_values(self) -> None: - """Test BedrockConfig accepts values.""" - config = BedrockConfig( - temperature=0.7, - max_tokens=1000, - top_p=0.9, - stop_sequences=['END'], - ) - assert config.temperature == 0.7 - assert config.max_tokens == 1000 - assert config.top_p == 0.9 - assert config.stop_sequences == ['END'] - - def test_anthropic_config_has_thinking(self) -> None: - """Test AnthropicConfig has thinking parameter.""" - config = AnthropicConfig(thinking={'enabled': True}) - assert config.thinking == {'enabled': True} - - def test_meta_llama_config_has_max_gen_len(self) -> None: - """Test MetaLlamaConfig has max_gen_len parameter.""" - config = MetaLlamaConfig(max_gen_len=1024) - assert config.max_gen_len == 1024 - - def test_mistral_config_has_safe_prompt(self) -> None: - """Test MistralConfig has safe_prompt parameter.""" - config = MistralConfig(safe_prompt=True, random_seed=42) - assert config.safe_prompt is True - assert config.random_seed == 42 - - def test_cohere_config_has_documents(self) -> None: - """Test CohereConfig has documents parameter.""" - config = CohereConfig( - k=50, - p=0.9, - safety_mode=CohereSafetyMode.CONTEXTUAL, - documents=['doc1', 'doc2'], - ) - assert config.k == 50 - assert config.p == 0.9 - assert config.safety_mode == CohereSafetyMode.CONTEXTUAL - assert config.documents == ['doc1', 'doc2'] - - def test_cohere_safety_mode_enum(self) -> None: - """Test CohereSafetyMode enum values.""" - assert CohereSafetyMode.CONTEXTUAL == 'CONTEXTUAL' - assert CohereSafetyMode.STRICT == 'STRICT' - assert CohereSafetyMode.OFF == 'OFF' - - def test_cohere_tool_choice_enum(self) -> None: - """Test CohereToolChoice enum values.""" - assert CohereToolChoice.REQUIRED == 'REQUIRED' - assert CohereToolChoice.NONE == 'NONE' - - def test_config_allows_extra_fields(self) -> None: - """Test configs allow extra fields for forward compatibility.""" - # Use model_validate to test that extra fields are allowed - config = BedrockConfig.model_validate({ - 'temperature': 0.5, - 'unknown_future_param': 'value', - }) - assert config.temperature == 0.5 - # Extra fields should be allowed due to extra='allow' - - def test_ai21_jamba_config_has_all_params(self) -> None: - """Test AI21 Jamba config has all documented parameters.""" - config = AI21JambaConfig( - n=3, - frequency_penalty=0.5, - presence_penalty=0.3, - stop=['###', '\n'], - ) - assert config.n == 3 - assert config.frequency_penalty == 0.5 - assert config.presence_penalty == 0.3 - assert config.stop == ['###', '\n'] - - def test_ai21_jamba_n_validation(self) -> None: - """Test AI21 Jamba n parameter has valid range.""" - config = AI21JambaConfig(n=1) - assert config.n == 1 - - config = AI21JambaConfig(n=16) - assert config.n == 16 - - with pytest.raises(ValueError): - AI21JambaConfig(n=0) # Below min - - with pytest.raises(ValueError): - AI21JambaConfig(n=17) # Above max - - def test_stability_config_text_to_image(self) -> None: - """Test Stability config for text-to-image generation.""" - config = StabilityConfig( - mode=StabilityMode.TEXT_TO_IMAGE, - aspect_ratio=StabilityAspectRatio.RATIO_16_9, - seed=12345, - negative_prompt='blurry, low quality', - output_format=StabilityOutputFormat.PNG, - ) - assert config.mode == StabilityMode.TEXT_TO_IMAGE - assert config.aspect_ratio == StabilityAspectRatio.RATIO_16_9 - assert config.seed == 12345 - assert config.negative_prompt == 'blurry, low quality' - assert config.output_format == StabilityOutputFormat.PNG - - def test_stability_config_image_to_image(self) -> None: - """Test Stability config for image-to-image generation.""" - config = StabilityConfig( - mode=StabilityMode.IMAGE_TO_IMAGE, - image='base64encodedimage', - strength=0.7, - seed=42, - ) - assert config.mode == StabilityMode.IMAGE_TO_IMAGE - assert config.image == 'base64encodedimage' - assert config.strength == 0.7 - assert config.seed == 42 - - def test_stability_strength_validation(self) -> None: - """Test Stability strength parameter has valid range.""" - config = StabilityConfig(strength=0.0) - assert config.strength == 0.0 - - config = StabilityConfig(strength=1.0) - assert config.strength == 1.0 - - with pytest.raises(ValueError): - StabilityConfig(strength=-0.1) # Below min - - with pytest.raises(ValueError): - StabilityConfig(strength=1.1) # Above max - - def test_stability_aspect_ratio_enum(self) -> None: - """Test StabilityAspectRatio enum values.""" - assert StabilityAspectRatio.RATIO_1_1 == '1:1' - assert StabilityAspectRatio.RATIO_16_9 == '16:9' - assert StabilityAspectRatio.RATIO_9_16 == '9:16' - assert StabilityAspectRatio.RATIO_21_9 == '21:9' - - def test_stability_output_format_enum(self) -> None: - """Test StabilityOutputFormat enum values.""" - assert StabilityOutputFormat.JPEG == 'jpeg' - assert StabilityOutputFormat.PNG == 'png' - assert StabilityOutputFormat.WEBP == 'webp' - - def test_stability_mode_enum(self) -> None: - """Test StabilityMode enum values.""" - assert StabilityMode.TEXT_TO_IMAGE == 'text-to-image' - assert StabilityMode.IMAGE_TO_IMAGE == 'image-to-image' - - def test_temperature_validation(self) -> None: - """Test temperature is validated within range.""" - # Valid temperature - config = BedrockConfig(temperature=0.5) - assert config.temperature == 0.5 - - # Invalid temperature should raise - with pytest.raises(ValueError): - BedrockConfig(temperature=1.5) # > 1.0 - - def test_top_p_validation(self) -> None: - """Test top_p is validated within range.""" - # Valid top_p - config = BedrockConfig(top_p=0.9) - assert config.top_p == 0.9 - - # Invalid top_p should raise - with pytest.raises(ValueError): - BedrockConfig(top_p=1.5) # > 1.0 - - -class TestEmbeddingModels: - """Tests for embedding model registry.""" - - def test_titan_embeddings_present(self) -> None: - """Test Amazon Titan embedding models are registered.""" - assert 'amazon.titan-embed-text-v2:0' in SUPPORTED_EMBEDDING_MODELS - assert 'amazon.titan-embed-text-v1' in SUPPORTED_EMBEDDING_MODELS - - def test_cohere_embeddings_present(self) -> None: - """Test Cohere embedding models are registered.""" - assert 'cohere.embed-english-v3' in SUPPORTED_EMBEDDING_MODELS - assert 'cohere.embed-multilingual-v3' in SUPPORTED_EMBEDDING_MODELS - - def test_embedding_model_has_dimensions(self) -> None: - """Test embedding models have dimensions specified.""" - titan_embed = SUPPORTED_EMBEDDING_MODELS['amazon.titan-embed-text-v2:0'] - assert 'dimensions' in titan_embed - assert titan_embed['dimensions'] > 0 - - def test_embedding_model_has_input_types(self) -> None: - """Test embedding models have input types specified.""" - titan_embed = SUPPORTED_EMBEDDING_MODELS['amazon.titan-embed-text-v2:0'] - assert 'supports' in titan_embed - assert 'input' in titan_embed['supports'] - assert 'text' in titan_embed['supports']['input'] - - -class TestInferenceProfileHelpers: - """Tests for inference profile helper functions.""" - - def test_get_inference_profile_prefix_us_regions(self) -> None: - """Test US regions return 'us' prefix.""" - assert get_inference_profile_prefix('us-east-1') == 'us' - assert get_inference_profile_prefix('us-east-2') == 'us' - assert get_inference_profile_prefix('us-west-1') == 'us' - assert get_inference_profile_prefix('us-west-2') == 'us' - - def test_get_inference_profile_prefix_eu_regions(self) -> None: - """Test EU regions return 'eu' prefix.""" - assert get_inference_profile_prefix('eu-west-1') == 'eu' - assert get_inference_profile_prefix('eu-west-2') == 'eu' - assert get_inference_profile_prefix('eu-central-1') == 'eu' - assert get_inference_profile_prefix('eu-north-1') == 'eu' - - def test_get_inference_profile_prefix_apac_regions(self) -> None: - """Test APAC regions return 'apac' prefix.""" - assert get_inference_profile_prefix('ap-northeast-1') == 'apac' - assert get_inference_profile_prefix('ap-southeast-1') == 'apac' - assert get_inference_profile_prefix('ap-south-1') == 'apac' - - def test_get_inference_profile_prefix_other_regions(self) -> None: - """Test other regions are routed appropriately.""" - # Canada routed through US - assert get_inference_profile_prefix('ca-central-1') == 'us' - # South America routed through US - assert get_inference_profile_prefix('sa-east-1') == 'us' - # Middle East routed through EU - assert get_inference_profile_prefix('me-south-1') == 'eu' - - def test_inference_profile_us(self) -> None: - """Test inference_profile with US region.""" - result = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0', 'us-east-1') - assert result == 'amazon-bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_inference_profile_eu(self) -> None: - """Test inference_profile with EU region.""" - result = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0', 'eu-west-1') - assert result == 'amazon-bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_inference_profile_apac(self) -> None: - """Test inference_profile with APAC region.""" - result = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0', 'ap-northeast-1') - assert result == 'amazon-bedrock/apac.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_inference_profile_default_region(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Test inference_profile uses AWS_REGION env var when no region specified.""" - monkeypatch.setenv('AWS_REGION', 'eu-central-1') - result = inference_profile('amazon.nova-pro-v1:0') - assert result == 'amazon-bedrock/eu.amazon.nova-pro-v1:0' - - def test_inference_profile_no_region_raises_error(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Test inference_profile raises error when no region is available.""" - monkeypatch.delenv('AWS_REGION', raising=False) - monkeypatch.delenv('AWS_DEFAULT_REGION', raising=False) - with pytest.raises(ValueError, match='AWS region is required'): - inference_profile('amazon.nova-pro-v1:0') - - def test_get_inference_profile_prefix_no_region_raises_error(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Test get_inference_profile_prefix raises error when no region is available.""" - monkeypatch.delenv('AWS_REGION', raising=False) - monkeypatch.delenv('AWS_DEFAULT_REGION', raising=False) - with pytest.raises(ValueError, match='AWS region is required'): - get_inference_profile_prefix() - - -class TestStripInferenceProfilePrefix: - """Tests for _strip_inference_profile_prefix. - - The InvokeModel API (used for embeddings) does NOT support inference - profile IDs. This helper strips regional prefixes so we always pass - base model IDs to InvokeModel. - """ - - def test_strip_us_prefix(self) -> None: - """Test stripping 'us.' prefix from model ID.""" - assert _strip_inference_profile_prefix('us.amazon.titan-embed-text-v2:0') == 'amazon.titan-embed-text-v2:0' - - def test_strip_eu_prefix(self) -> None: - """Test stripping 'eu.' prefix from model ID.""" - assert _strip_inference_profile_prefix('eu.cohere.embed-english-v3') == 'cohere.embed-english-v3' - - def test_strip_apac_prefix(self) -> None: - """Test stripping 'apac.' prefix from model ID.""" - assert _strip_inference_profile_prefix('apac.amazon.titan-embed-text-v1') == 'amazon.titan-embed-text-v1' - - def test_no_prefix_unchanged(self) -> None: - """Test model IDs without prefix are returned unchanged.""" - assert _strip_inference_profile_prefix('amazon.titan-embed-text-v2:0') == 'amazon.titan-embed-text-v2:0' - - def test_no_prefix_cohere(self) -> None: - """Test Cohere model ID without prefix is unchanged.""" - assert _strip_inference_profile_prefix('cohere.embed-english-v3') == 'cohere.embed-english-v3' - - def test_no_prefix_nova(self) -> None: - """Test Nova embed model ID without prefix is unchanged.""" - assert _strip_inference_profile_prefix('amazon.nova-embed-text-v1:0') == 'amazon.nova-embed-text-v1:0' - - def test_round_trip_with_inference_profile(self) -> None: - """Test that inference_profile + strip recovers the original model ID. - - This verifies the round-trip property: a base model ID passed through - inference_profile() and then _strip_inference_profile_prefix() should - yield the original base model ID (without the plugin prefix). - """ - base_id = 'amazon.titan-embed-text-v2:0' - # inference_profile adds 'amazon-bedrock/{prefix}.' prefixed name - full_ref = inference_profile(base_id, 'us-east-1') - # Strip 'amazon-bedrock/' prefix (as the plugin does) - model_id = full_ref.split('/', 1)[1] # 'us.amazon.titan-embed-text-v2:0' - # Strip inference profile prefix - stripped = _strip_inference_profile_prefix(model_id) - assert stripped == base_id - - -class TestAutoInferenceProfileConversion: - """Tests for automatic inference profile conversion in BedrockModel. - - When using API key authentication (AWS_BEARER_TOKEN_BEDROCK), the plugin - should automatically convert direct model IDs to inference profile IDs - by adding the appropriate regional prefix. - """ - - @pytest.fixture - def mock_client(self) -> object: - """Create a mock boto3 client.""" - return MagicMock() - - @pytest.fixture - def clear_env(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Clear all relevant environment variables before each test.""" - monkeypatch.delenv('AWS_BEARER_TOKEN_BEDROCK', raising=False) - monkeypatch.delenv('AWS_REGION', raising=False) - monkeypatch.delenv('AWS_DEFAULT_REGION', raising=False) - - def test_iam_auth_returns_direct_model_id( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test IAM auth (no API key) returns direct model ID unchanged.""" - # No AWS_BEARER_TOKEN_BEDROCK set = IAM auth - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - assert model._get_effective_model_id() == 'anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_api_key_auth_us_region_adds_us_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth with US region adds 'us.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-east-1') - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - assert model._get_effective_model_id() == 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_api_key_auth_us_west_region_adds_us_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth with us-west region adds 'us.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-west-2') - - model = BedrockModel('mistral.mistral-large-3-675b-instruct', mock_client) - assert model._get_effective_model_id() == 'us.mistral.mistral-large-3-675b-instruct' - - def test_api_key_auth_eu_region_adds_eu_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth with EU region adds 'eu.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'eu-west-1') - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - assert model._get_effective_model_id() == 'eu.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_api_key_auth_eu_central_region_adds_eu_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth with eu-central region adds 'eu.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'eu-central-1') - - model = BedrockModel('meta.llama3-3-70b-instruct-v1:0', mock_client) - assert model._get_effective_model_id() == 'eu.meta.llama3-3-70b-instruct-v1:0' - - def test_api_key_auth_apac_region_adds_apac_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth with APAC region adds 'apac.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'ap-northeast-1') - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - assert model._get_effective_model_id() == 'apac.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_api_key_auth_ap_southeast_region_adds_apac_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth with ap-southeast region adds 'apac.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'ap-southeast-1') - - model = BedrockModel('deepseek.r1-v1:0', mock_client) - assert model._get_effective_model_id() == 'apac.deepseek.r1-v1:0' - - def test_api_key_auth_uses_aws_default_region( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth falls back to AWS_DEFAULT_REGION.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_DEFAULT_REGION', 'eu-west-2') - - model = BedrockModel('amazon.nova-pro-v1:0', mock_client) - assert model._get_effective_model_id() == 'eu.amazon.nova-pro-v1:0' - - def test_api_key_auth_aws_region_takes_precedence( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test AWS_REGION takes precedence over AWS_DEFAULT_REGION.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-east-1') - monkeypatch.setenv('AWS_DEFAULT_REGION', 'eu-west-1') - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - # Should use AWS_REGION (us), not AWS_DEFAULT_REGION (eu) - assert model._get_effective_model_id() == 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_model_already_has_us_prefix_unchanged( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test model with existing 'us.' prefix is unchanged.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'eu-west-1') # Different region - - model = BedrockModel('us.anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - # Should NOT add another prefix, even though region is EU - assert model._get_effective_model_id() == 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_model_already_has_eu_prefix_unchanged( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test model with existing 'eu.' prefix is unchanged.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-east-1') # Different region - - model = BedrockModel('eu.anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - # Should NOT add another prefix - assert model._get_effective_model_id() == 'eu.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_model_already_has_apac_prefix_unchanged( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test model with existing 'apac.' prefix is unchanged.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-east-1') # Different region - - model = BedrockModel('apac.amazon.nova-pro-v1:0', mock_client) - # Should NOT add another prefix - assert model._get_effective_model_id() == 'apac.amazon.nova-pro-v1:0' - - def test_api_key_auth_no_region_returns_direct_id( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test API key auth without region returns direct model ID with warning.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - # No AWS_REGION or AWS_DEFAULT_REGION set - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - # Should return direct ID (will likely fail at API call, but that's expected) - assert model._get_effective_model_id() == 'anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_unknown_region_defaults_to_us( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test unknown region prefix defaults to 'us.'.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'unknown-region-1') # Unrecognized - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - assert model._get_effective_model_id() == 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_canada_region_uses_us_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test Canada region uses 'us.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'ca-central-1') - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - # Canada is routed through US prefix (defaults to us for unknown) - assert model._get_effective_model_id() == 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_south_america_region_uses_us_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test South America region uses 'us.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'sa-east-1') - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - # SA mapped to apac in _get_effective_model_id (starts with sa-) - assert model._get_effective_model_id() == 'apac.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_middle_east_region_uses_apac_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test Middle East region uses 'apac.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'me-south-1') - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - # ME mapped to apac in _get_effective_model_id (starts with me-) - assert model._get_effective_model_id() == 'apac.anthropic.claude-sonnet-4-5-20250929-v1:0' - - def test_africa_region_uses_apac_prefix( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test Africa region uses 'apac.' prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'af-south-1') - - model = BedrockModel('amazon.nova-pro-v1:0', mock_client) - # Africa mapped to apac in _get_effective_model_id (starts with af-) - assert model._get_effective_model_id() == 'apac.amazon.nova-pro-v1:0' - - def test_deepseek_model_auto_conversion( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test DeepSeek model auto-converts with API key auth.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-west-2') - - model = BedrockModel('deepseek.r1-v1:0', mock_client) - assert model._get_effective_model_id() == 'us.deepseek.r1-v1:0' - - def test_nova_model_auto_conversion( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test Nova model auto-converts with API key auth.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'ap-south-1') - - model = BedrockModel('amazon.nova-pro-v1:0', mock_client) - assert model._get_effective_model_id() == 'apac.amazon.nova-pro-v1:0' - - def test_cohere_model_auto_conversion( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test Cohere model auto-converts with API key auth.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'eu-north-1') - - model = BedrockModel('cohere.command-r-plus-v1:0', mock_client) - assert model._get_effective_model_id() == 'eu.cohere.command-r-plus-v1:0' - - def test_ai21_model_no_conversion_with_api_key( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test AI21 models do NOT get inference profile prefix (not supported).""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-east-1') - - model = BedrockModel('ai21.jamba-1-5-large-v1:0', mock_client) - # AI21 doesn't support cross-region inference profiles - should use direct ID - assert model._get_effective_model_id() == 'ai21.jamba-1-5-large-v1:0' - - def test_ai21_mini_model_no_conversion_with_api_key( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test AI21 Jamba Mini model does NOT get inference profile prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'eu-west-1') - - model = BedrockModel('ai21.jamba-1-5-mini-v1:0', mock_client) - # AI21 doesn't support cross-region inference profiles - assert model._get_effective_model_id() == 'ai21.jamba-1-5-mini-v1:0' - - def test_stability_model_no_conversion_with_api_key( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test Stability AI models do NOT get inference profile prefix (not supported).""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-west-2') - - model = BedrockModel('stability.sd3-5-large-v1:0', mock_client) - # Stability doesn't support cross-region inference profiles - assert model._get_effective_model_id() == 'stability.sd3-5-large-v1:0' - - def test_unsupported_provider_no_conversion( - self, mock_client: object, clear_env: None, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Test unknown/unsupported providers do NOT get inference profile prefix.""" - monkeypatch.setenv('AWS_BEARER_TOKEN_BEDROCK', 'test-token') - monkeypatch.setenv('AWS_REGION', 'us-east-1') - - model = BedrockModel('unknown-provider.some-model-v1:0', mock_client) - # Unknown provider - should not add prefix - assert model._get_effective_model_id() == 'unknown-provider.some-model-v1:0' - - -class _AsyncIter(AsyncIterator[Any]): - """Wraps a list into an async iterator for mocking aioboto3 streams.""" - - def __init__(self, items: list[Any]) -> None: - self._items = iter(items) - - def __aiter__(self) -> '_AsyncIter': - return self - - async def __anext__(self) -> Any: # noqa: ANN401 - try: - return next(self._items) - except StopIteration: - raise StopAsyncIteration # noqa: B904 - - -class TestStreamingToolUseParsing: - """Regression tests for streaming tool use assembly. - - Bedrock ConverseStream sends tool calls across multiple events: - 1. contentBlockStart β€” contains toolUseId and name - 2. contentBlockDelta(s) β€” contains input JSON fragments - 3. contentBlockStop - - A previous bug compared contentBlockIndex as int (from delta) vs - string key (from start), causing the delta handler to create a - *new* entry with empty name/ref instead of appending to the - existing one. This resulted in RuntimeError: 'failed not found'. - """ - - @pytest.fixture() - def mock_client(self) -> MagicMock: - """Create a mock async bedrock-runtime client.""" - client = MagicMock() - # converse_stream is async with aioboto3 - client.converse_stream = AsyncMock() - return client - - @pytest.mark.asyncio - async def test_tool_use_name_preserved_across_stream_events(self, mock_client: MagicMock) -> None: - """Tool name and ref from contentBlockStart survive delta assembly.""" - # Simulate the ConverseStream event sequence for a tool call - mock_client.converse_stream.return_value = { - 'stream': _AsyncIter([ - # 1. contentBlockStart: carries toolUseId and name - { - 'contentBlockStart': { - 'contentBlockIndex': 1, - 'start': { - 'toolUse': { - 'toolUseId': 'call_abc123', - 'name': 'get_weather', - } - }, - } - }, - # 2. contentBlockDelta: carries input JSON fragment - { - 'contentBlockDelta': { - 'contentBlockIndex': 1, - 'delta': { - 'toolUse': { - 'input': '{"location"', - } - }, - } - }, - # 3. Another delta with more input - { - 'contentBlockDelta': { - 'contentBlockIndex': 1, - 'delta': { - 'toolUse': { - 'input': ': "London"}', - } - }, - } - }, - # 4. contentBlockStop - { - 'contentBlockStop': { - 'contentBlockIndex': 1, - } - }, - # 5. messageStop - { - 'messageStop': { - 'stopReason': 'tool_use', - } - }, - # 5. metadata - { - 'metadata': { - 'usage': { - 'inputTokens': 100, - 'outputTokens': 50, - 'totalTokens': 150, - } - } - }, - ]) - } - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - request = GenerateRequest( - messages=[ - Message( - role=Role.USER, - content=[Part(root=TextPart(text='What is the weather in London?'))], - ) - ], - tools=[], - ) - - chunks: list[GenerateResponseChunk] = [] - ctx = MagicMock(spec=ActionRunContext) - ctx.send_chunk = MagicMock(side_effect=lambda c: chunks.append(c)) - - response = await model._generate_streaming( - {'modelId': 'anthropic.claude-sonnet-4-5-20250929-v1:0', 'messages': []}, - ctx, - request, - ) - - # Find the tool request part in the response - assert response.message is not None - tool_parts = [p for p in response.message.content if p.root.tool_request is not None] - assert len(tool_parts) == 1, f'Expected 1 tool request, got {len(tool_parts)}' - - tool_req = tool_parts[0].root.tool_request - assert tool_req is not None - assert isinstance(tool_req, ToolRequest) - assert tool_req.name == 'get_weather', f'Tool name should be "get_weather", got "{tool_req.name}"' - assert tool_req.ref == 'call_abc123', f'Tool ref should be "call_abc123", got "{tool_req.ref}"' - assert tool_req.input == {'location': 'London'}, f'Tool input should be parsed JSON, got {tool_req.input!r}' - - # Verify a tool request chunk was emitted via ctx.send_chunk. - tool_chunks = [ - c - for c in chunks - if any(hasattr(p.root, 'tool_request') and p.root.tool_request is not None for p in c.content) - ] - assert len(tool_chunks) == 1, f'Expected 1 tool request chunk, got {len(tool_chunks)}' - - @pytest.mark.asyncio - async def test_multiple_tool_calls_in_stream(self, mock_client: MagicMock) -> None: - """Multiple tool calls in a single stream are assembled correctly.""" - mock_client.converse_stream.return_value = { - 'stream': _AsyncIter([ - # First tool call at block index 0 - { - 'contentBlockStart': { - 'contentBlockIndex': 0, - 'start': { - 'toolUse': { - 'toolUseId': 'call_001', - 'name': 'get_weather', - } - }, - } - }, - { - 'contentBlockDelta': { - 'contentBlockIndex': 0, - 'delta': {'toolUse': {'input': '{"location": "London"}'}}, - } - }, - # Second tool call at block index 1 - { - 'contentBlockStart': { - 'contentBlockIndex': 1, - 'start': { - 'toolUse': { - 'toolUseId': 'call_002', - 'name': 'get_time', - } - }, - } - }, - { - 'contentBlockDelta': { - 'contentBlockIndex': 1, - 'delta': {'toolUse': {'input': '{"timezone": "UTC"}'}}, - } - }, - {'messageStop': {'stopReason': 'tool_use'}}, - { - 'metadata': { - 'usage': { - 'inputTokens': 200, - 'outputTokens': 80, - 'totalTokens': 280, - } - } - }, - ]) - } - - model = BedrockModel('anthropic.claude-sonnet-4-5-20250929-v1:0', mock_client) - request = GenerateRequest( - messages=[ - Message( - role=Role.USER, - content=[Part(root=TextPart(text='Weather and time?'))], - ) - ], - tools=[], - ) - ctx = MagicMock(spec=ActionRunContext) - ctx.send_chunk = MagicMock() - - response = await model._generate_streaming( - {'modelId': 'anthropic.claude-sonnet-4-5-20250929-v1:0', 'messages': []}, - ctx, - request, - ) - - assert response.message is not None - tool_parts = [p for p in response.message.content if p.root.tool_request is not None] - assert len(tool_parts) == 2, f'Expected 2 tool requests, got {len(tool_parts)}' - - # Verify first tool - tool_req_0 = tool_parts[0].root.tool_request - assert tool_req_0 is not None - assert isinstance(tool_req_0, ToolRequest) - assert tool_req_0.name == 'get_weather' - assert tool_req_0.ref == 'call_001' - assert tool_req_0.input == {'location': 'London'} - - # Verify second tool - tool_req_1 = tool_parts[1].root.tool_request - assert tool_req_1 is not None - assert isinstance(tool_req_1, ToolRequest) - assert tool_req_1.name == 'get_time' - assert tool_req_1.ref == 'call_002' - assert tool_req_1.input == {'timezone': 'UTC'} diff --git a/py/plugins/amazon-bedrock/tests/amazon_bedrock_telemetry_test.py b/py/plugins/amazon-bedrock/tests/amazon_bedrock_telemetry_test.py deleted file mode 100644 index 9c0102f030..0000000000 --- a/py/plugins/amazon-bedrock/tests/amazon_bedrock_telemetry_test.py +++ /dev/null @@ -1,423 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for AWS X-Ray telemetry integration in the amazon-bedrock plugin. - -Covers: -- Region resolution from environment variables -- X-Ray OTLP endpoint configuration -- AwsXRayOtlpExporter with SigV4 authentication -- SigV4SigningAdapter for HTTPS requests -- AwsTelemetry manager class -- TimeAdjustedSpan for zero-duration spans -- Log-trace correlation -- Error handling for missing credentials -""" - -import os -from unittest import mock - -import pytest -import requests - -from genkit.plugins.amazon_bedrock.telemetry.tracing import ( - XRAY_OTLP_ENDPOINT_PATTERN, - AwsAdjustingTraceExporter, - AwsTelemetry, - AwsXRayOtlpExporter, - SigV4SigningAdapter, - TimeAdjustedSpan, - _create_sigv4_session, - _resolve_region, - add_aws_telemetry, -) - - -class TestResolveRegion: - """Tests for the _resolve_region function.""" - - def test_explicit_region_takes_precedence(self) -> None: - """Explicit region parameter should override environment variables.""" - with mock.patch.dict(os.environ, {'AWS_REGION': 'us-east-1'}): - result = _resolve_region(region='eu-west-1') - assert result == 'eu-west-1' - - def test_aws_region_env_var(self) -> None: - """AWS_REGION environment variable should be used when no explicit region.""" - with mock.patch.dict(os.environ, {'AWS_REGION': 'us-west-2'}, clear=True): - result = _resolve_region() - assert result == 'us-west-2' - - def test_aws_default_region_fallback(self) -> None: - """AWS_DEFAULT_REGION should be used as fallback.""" - env = {'AWS_DEFAULT_REGION': 'ap-southeast-1'} - with mock.patch.dict(os.environ, env, clear=True): - result = _resolve_region() - assert result == 'ap-southeast-1' - - def test_aws_region_priority_over_default(self) -> None: - """AWS_REGION should take priority over AWS_DEFAULT_REGION.""" - env = { - 'AWS_REGION': 'us-east-1', - 'AWS_DEFAULT_REGION': 'us-west-2', - } - with mock.patch.dict(os.environ, env, clear=True): - result = _resolve_region() - assert result == 'us-east-1' - - def test_no_region_returns_none(self) -> None: - """Should return None when no region is configured.""" - with mock.patch.dict(os.environ, {}, clear=True): - result = _resolve_region() - assert result is None - - -class TestXRayOtlpEndpoint: - """Tests for X-Ray OTLP endpoint configuration.""" - - def test_endpoint_pattern_format(self) -> None: - """Endpoint pattern should produce correct URLs.""" - endpoint = XRAY_OTLP_ENDPOINT_PATTERN.format(region='us-west-2') - assert endpoint == 'https://xray.us-west-2.amazonaws.com/v1/traces' - - def test_endpoint_pattern_different_regions(self) -> None: - """Endpoint should work for various AWS regions.""" - regions = ['us-east-1', 'eu-west-1', 'ap-southeast-1', 'sa-east-1'] - for region in regions: - endpoint = XRAY_OTLP_ENDPOINT_PATTERN.format(region=region) - assert region in endpoint - assert endpoint.startswith('https://xray.') - assert endpoint.endswith('/v1/traces') - - -class TestAwsXRayOtlpExporter: - """Tests for the AwsXRayOtlpExporter class.""" - - def test_exporter_initialization(self) -> None: - """Exporter should initialize with region.""" - exporter = AwsXRayOtlpExporter(region='us-west-2') - assert exporter._region == 'us-west-2' - assert 'us-west-2' in exporter._endpoint - - def test_exporter_with_error_handler(self) -> None: - """Exporter should accept error handler callback.""" - errors: list[Exception] = [] - exporter = AwsXRayOtlpExporter( - region='us-west-2', - error_handler=lambda e: errors.append(e), - ) - assert exporter._error_handler is not None - - def test_exporter_uses_sigv4_session(self) -> None: - """Exporter should use a session with SigV4 signing adapter mounted.""" - exporter = AwsXRayOtlpExporter(region='us-west-2') - # The OTLP exporter should have a session configured - assert exporter._otlp_exporter._session is not None - - -class TestSigV4SigningAdapter: - """Tests for the SigV4SigningAdapter class.""" - - def test_adapter_initialization(self) -> None: - """Adapter should initialize with credentials and region.""" - mock_credentials = mock.MagicMock() - adapter = SigV4SigningAdapter( - credentials=mock_credentials, - region='us-west-2', - service='xray', - ) - assert adapter._credentials == mock_credentials - assert adapter._region == 'us-west-2' - assert adapter._service == 'xray' - - def test_adapter_default_service(self) -> None: - """Adapter should default to xray service.""" - mock_credentials = mock.MagicMock() - adapter = SigV4SigningAdapter( - credentials=mock_credentials, - region='eu-west-1', - ) - assert adapter._service == 'xray' - - def test_adapter_signs_request(self) -> None: - """Adapter should add SigV4 headers to request.""" - # Create mock credentials - mock_credentials = mock.MagicMock() - mock_credentials.access_key = 'AKIAIOSFODNN7EXAMPLE' - mock_credentials.secret_key = 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY' - mock_credentials.token = None - - adapter = SigV4SigningAdapter( - credentials=mock_credentials, - region='us-west-2', - service='xray', - ) - - # Create a mock request - request = requests.PreparedRequest() - request.prepare( - method='POST', - url='https://xray.us-west-2.amazonaws.com/v1/traces', - headers={'Content-Type': 'application/x-protobuf'}, - data=b'test-payload', - ) - - # Mock the parent send method - with mock.patch.object( - adapter.__class__.__bases__[0], - 'send', - return_value=mock.MagicMock(status_code=200), - ): - adapter.send(request) - - # Verify that the request now has Authorization header - assert request.headers is not None - assert 'Authorization' in request.headers - assert 'AWS4-HMAC-SHA256' in request.headers['Authorization'] - - def test_adapter_handles_none_credentials(self) -> None: - """Adapter should handle None credentials gracefully.""" - adapter = SigV4SigningAdapter( - credentials=None, - region='us-west-2', - ) - - request = requests.PreparedRequest() - request.prepare( - method='POST', - url='https://xray.us-west-2.amazonaws.com/v1/traces', - headers={}, - data=b'test', - ) - - # Should not raise, just skip signing - with mock.patch.object( - adapter.__class__.__bases__[0], - 'send', - return_value=mock.MagicMock(status_code=200), - ): - adapter.send(request) - - # No Authorization header should be set - assert request.headers is not None - assert 'Authorization' not in request.headers - - -class TestCreateSigV4Session: - """Tests for the _create_sigv4_session function.""" - - def test_session_has_adapter_mounted(self) -> None: - """Session should have SigV4 adapter mounted for HTTPS.""" - mock_credentials = mock.MagicMock() - session = _create_sigv4_session( - credentials=mock_credentials, - region='us-west-2', - ) - - # Get the adapter for an HTTPS URL - adapter = session.get_adapter('https://xray.us-west-2.amazonaws.com') - assert isinstance(adapter, SigV4SigningAdapter) - - def test_session_adapter_has_correct_region(self) -> None: - """Session adapter should be configured with correct region.""" - mock_credentials = mock.MagicMock() - session = _create_sigv4_session( - credentials=mock_credentials, - region='eu-west-1', - service='xray', - ) - - adapter = session.get_adapter('https://xray.eu-west-1.amazonaws.com') - assert isinstance(adapter, SigV4SigningAdapter) - assert adapter._region == 'eu-west-1' - assert adapter._service == 'xray' - - -class TestAwsAdjustingTraceExporter: - """Tests for the AwsAdjustingTraceExporter class.""" - - def test_exporter_initialization(self) -> None: - """Adjusting exporter should wrap base exporter.""" - base_exporter = mock.MagicMock() - exporter = AwsAdjustingTraceExporter( - exporter=base_exporter, - log_input_and_output=False, - region='us-west-2', - ) - assert exporter._exporter is base_exporter - assert exporter._log_input_and_output is False - assert exporter._region == 'us-west-2' - - def test_exporter_with_logging_enabled(self) -> None: - """Adjusting exporter should respect log_input_and_output flag.""" - base_exporter = mock.MagicMock() - exporter = AwsAdjustingTraceExporter( - exporter=base_exporter, - log_input_and_output=True, - ) - assert exporter._log_input_and_output is True - - -class TestTimeAdjustedSpan: - """Tests for the TimeAdjustedSpan class.""" - - def test_zero_duration_span_adjusted(self) -> None: - """Spans with zero duration should get minimum 1 microsecond.""" - mock_span = mock.MagicMock() - mock_span.start_time = 1000000000 # 1 second in nanoseconds - mock_span.end_time = 1000000000 # Same as start (zero duration) - mock_span.attributes = {} - - adjusted = TimeAdjustedSpan(mock_span, {}) - # Should add 1000 nanoseconds (1 microsecond) - assert adjusted.end_time == 1000001000 - - def test_none_end_time_adjusted(self) -> None: - """Spans with None end_time should get adjusted.""" - mock_span = mock.MagicMock() - mock_span.start_time = 1000000000 - mock_span.end_time = None - mock_span.attributes = {} - - adjusted = TimeAdjustedSpan(mock_span, {}) - assert adjusted.end_time == 1000001000 - - def test_valid_duration_unchanged(self) -> None: - """Spans with valid duration should remain unchanged.""" - mock_span = mock.MagicMock() - mock_span.start_time = 1000000000 - mock_span.end_time = 2000000000 # 1 second later - mock_span.attributes = {} - - adjusted = TimeAdjustedSpan(mock_span, {}) - assert adjusted.end_time == 2000000000 - - -class TestAwsTelemetry: - """Tests for the AwsTelemetry manager class.""" - - def test_initialization_with_explicit_region(self) -> None: - """Manager should accept explicit region.""" - with mock.patch.dict(os.environ, {}, clear=True): - telemetry = AwsTelemetry(region='us-west-2') - assert telemetry.region == 'us-west-2' - - def test_initialization_with_env_region(self) -> None: - """Manager should use AWS_REGION env var.""" - with mock.patch.dict(os.environ, {'AWS_REGION': 'eu-west-1'}, clear=True): - telemetry = AwsTelemetry() - assert telemetry.region == 'eu-west-1' - - def test_initialization_raises_without_region(self) -> None: - """Manager should raise ValueError without region.""" - with mock.patch.dict(os.environ, {}, clear=True): - with pytest.raises(ValueError, match='AWS region is required'): - AwsTelemetry() - - def test_default_configuration(self) -> None: - """Manager should have correct defaults.""" - with mock.patch.dict(os.environ, {'AWS_REGION': 'us-west-2'}): - telemetry = AwsTelemetry() - assert telemetry.log_input_and_output is False - assert telemetry.force_dev_export is True - assert telemetry.disable_traces is False - - def test_inject_trace_context_no_span(self) -> None: - """Trace context injection should handle no active span.""" - with mock.patch.dict(os.environ, {'AWS_REGION': 'us-west-2'}): - telemetry = AwsTelemetry() - event_dict: dict[str, str] = {'message': 'test'} - - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.trace') as mock_trace: - mock_trace.get_current_span.return_value = mock_trace.INVALID_SPAN - result = telemetry._inject_trace_context(event_dict) - - assert '_X_AMZN_TRACE_ID' not in result - - -class TestAddAwsTelemetry: - """Tests for the add_aws_telemetry function.""" - - def test_raises_without_region(self) -> None: - """Should raise ValueError when region is not configured.""" - with mock.patch.dict(os.environ, {}, clear=True): - with pytest.raises(ValueError, match='AWS region is required'): - add_aws_telemetry() - - def test_accepts_explicit_region(self) -> None: - """Should accept explicit region parameter.""" - with mock.patch.dict(os.environ, {}, clear=True): - # Mock the exporter to avoid actual export - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.add_custom_exporter') as mock_add: - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.propagate'): - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.trace'): - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.structlog'): - add_aws_telemetry(region='us-west-2') - mock_add.assert_called_once() - - def test_skips_in_dev_without_force(self) -> None: - """Should skip telemetry in dev environment without force_dev_export.""" - with ( - mock.patch.dict(os.environ, {'AWS_REGION': 'us-west-2'}), - mock.patch( - 'genkit.plugins.amazon_bedrock.telemetry.tracing.is_dev_environment', - return_value=True, - ), - mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.add_custom_exporter') as mock_add, - ): - add_aws_telemetry(force_dev_export=False) - mock_add.assert_not_called() - - def test_exports_in_dev_with_force(self) -> None: - """Should export telemetry in dev environment with force_dev_export=True.""" - with ( - mock.patch.dict(os.environ, {'AWS_REGION': 'us-west-2'}), - mock.patch( - 'genkit.plugins.amazon_bedrock.telemetry.tracing.is_dev_environment', - return_value=True, - ), - mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.add_custom_exporter') as mock_add, - ): - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.propagate'): - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.trace'): - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.structlog'): - add_aws_telemetry(force_dev_export=True) - mock_add.assert_called_once() - - def test_disable_traces(self) -> None: - """Should not add exporter when disable_traces=True.""" - with mock.patch.dict(os.environ, {'AWS_REGION': 'us-west-2'}): - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.add_custom_exporter') as mock_add: - with mock.patch('genkit.plugins.amazon_bedrock.telemetry.tracing.structlog'): - add_aws_telemetry(disable_traces=True) - mock_add.assert_not_called() - - -class TestTelemetryExportFromPlugin: - """Tests for the add_aws_telemetry export from the main plugin.""" - - def test_add_aws_telemetry_exported(self) -> None: - """add_aws_telemetry should be exported from the main plugin.""" - from genkit.plugins.amazon_bedrock import add_aws_telemetry as exported_fn - - # Verify it's the same function - assert exported_fn is add_aws_telemetry - - def test_add_aws_telemetry_in_all(self) -> None: - """add_aws_telemetry should be in __all__.""" - from genkit.plugins import amazon_bedrock - - assert 'add_aws_telemetry' in amazon_bedrock.__all__ diff --git a/py/plugins/amazon-bedrock/tests/amazon_bedrock_typing_test.py b/py/plugins/amazon-bedrock/tests/amazon_bedrock_typing_test.py deleted file mode 100644 index d4cee92b2a..0000000000 --- a/py/plugins/amazon-bedrock/tests/amazon_bedrock_typing_test.py +++ /dev/null @@ -1,165 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Amazon Bedrock typing and config schemas.""" - -import pytest -from pydantic import ValidationError - -from genkit.plugins.amazon_bedrock.typing import ( - AnthropicConfig, - BedrockConfig, - CohereConfig, - CohereSafetyMode, - CohereToolChoice, - MetaLlamaConfig, -) - - -class TestStrEnums: - """Tests for Bedrock StrEnum types.""" - - def test_cohere_safety_mode_values(self) -> None: - """Test Cohere safety mode values.""" - assert CohereSafetyMode.CONTEXTUAL == 'CONTEXTUAL' - assert CohereSafetyMode.STRICT == 'STRICT' - assert CohereSafetyMode.OFF == 'OFF' - - def test_cohere_tool_choice_values(self) -> None: - """Test Cohere tool choice values.""" - assert CohereToolChoice.REQUIRED == 'REQUIRED' - assert CohereToolChoice.NONE == 'NONE' - - -class TestBedrockConfig: - """Tests for the base BedrockConfig.""" - - def test_defaults(self) -> None: - """Test Defaults.""" - cfg = BedrockConfig() - assert cfg.max_tokens is None - assert cfg.temperature is None - assert cfg.top_p is None - - def test_genkit_common_params(self) -> None: - """Test Genkit common params.""" - cfg = BedrockConfig(temperature=0.5, max_output_tokens=1024, top_p=0.9) - assert cfg.temperature == 0.5 - assert cfg.max_output_tokens == 1024 - assert cfg.top_p == 0.9 - - def test_temperature_bounds(self) -> None: - """Test Temperature bounds.""" - cfg = BedrockConfig(temperature=0.0) - assert cfg.temperature == 0.0 - with pytest.raises(ValidationError): - BedrockConfig(temperature=-0.1) - - def test_top_p_bounds(self) -> None: - """Test Top p bounds.""" - cfg = BedrockConfig(top_p=1.0) - assert cfg.top_p == 1.0 - with pytest.raises(ValidationError): - BedrockConfig(top_p=1.1) - - def test_extra_fields_allowed(self) -> None: - """Test Extra fields allowed.""" - cfg = BedrockConfig.model_validate({'custom_param': 42}) - assert cfg.model_extra is not None - assert cfg.model_extra['custom_param'] == 42 - - -class TestAnthropicConfig: - """Tests for AnthropicConfig.""" - - def test_defaults(self) -> None: - """Test Defaults.""" - cfg = AnthropicConfig() - assert cfg.top_k is None - assert cfg.thinking is None - - def test_top_k_valid(self) -> None: - """Test Top k valid.""" - cfg = AnthropicConfig(top_k=40) - assert cfg.top_k == 40 - - def test_thinking_dict(self) -> None: - """Test Thinking dict.""" - cfg = AnthropicConfig.model_validate({'thinking': {'type': 'enabled', 'budget_tokens': 1024}}) - assert cfg.thinking is not None - assert cfg.thinking['type'] == 'enabled' - - def test_inherits_bedrock_config(self) -> None: - """Test Inherits bedrock config.""" - cfg = AnthropicConfig(temperature=0.7, max_output_tokens=2048) - assert cfg.temperature == 0.7 - assert cfg.max_output_tokens == 2048 - - -class TestCohereConfig: - """Tests for CohereConfig.""" - - def test_defaults(self) -> None: - """Test Defaults.""" - cfg = CohereConfig() - assert cfg.safety_mode is None - assert cfg.tool_choice is None - - def test_safety_mode_enum(self) -> None: - """Test Safety mode enum.""" - cfg = CohereConfig(safety_mode=CohereSafetyMode.STRICT) - assert cfg.safety_mode == 'STRICT' - - def test_tool_choice_enum(self) -> None: - """Test Tool choice enum.""" - cfg = CohereConfig(tool_choice=CohereToolChoice.REQUIRED) - assert cfg.tool_choice == 'REQUIRED' - - -class TestMetaLlamaConfig: - """Tests for MetaLlamaConfig.""" - - def test_defaults(self) -> None: - """Test Defaults.""" - cfg = MetaLlamaConfig() - assert cfg.max_tokens is None - - def test_inherits_bedrock(self) -> None: - """Test Inherits bedrock.""" - cfg = MetaLlamaConfig(temperature=0.3, top_p=0.8) - assert cfg.temperature == 0.3 - assert cfg.top_p == 0.8 - - -class TestCamelCaseAliases: - """Tests for camelCase alias generation.""" - - def test_max_output_tokens_alias(self) -> None: - """Test Max output tokens alias.""" - cfg = BedrockConfig.model_validate({'maxOutputTokens': 512}) - assert cfg.max_output_tokens == 512 - - def test_stop_sequences_alias(self) -> None: - """Test Stop sequences alias.""" - cfg = BedrockConfig.model_validate({'stopSequences': ['END']}) - assert cfg.stop_sequences == ['END'] - - def test_round_trip_by_alias(self) -> None: - """Test Round trip by alias.""" - cfg = BedrockConfig(max_output_tokens=100, stop_sequences=['STOP']) - data = cfg.model_dump(by_alias=True, exclude_none=True) - assert 'maxOutputTokens' in data - assert 'stopSequences' in data diff --git a/py/plugins/anthropic/src/genkit/plugins/anthropic/model_info.py b/py/plugins/anthropic/src/genkit/plugins/anthropic/model_info.py index ea7d3b13af..223a02d777 100644 --- a/py/plugins/anthropic/src/genkit/plugins/anthropic/model_info.py +++ b/py/plugins/anthropic/src/genkit/plugins/anthropic/model_info.py @@ -16,7 +16,7 @@ """Anthropic Models for Genkit.""" -from genkit.types import ( +from genkit import ( Constrained, ModelInfo, Supports, diff --git a/py/plugins/anthropic/src/genkit/plugins/anthropic/models.py b/py/plugins/anthropic/src/genkit/plugins/anthropic/models.py index 5c99abc33b..9ab64a2914 100644 --- a/py/plugins/anthropic/src/genkit/plugins/anthropic/models.py +++ b/py/plugins/anthropic/src/genkit/plugins/anthropic/models.py @@ -27,26 +27,18 @@ import json from typing import Any +import structlog + from anthropic import AsyncAnthropic from anthropic.types import Message as AnthropicMessage -from genkit.ai import ActionRunContext -from genkit.blocks.model import get_basic_usage_stats -from genkit.core.logging import get_logger -from genkit.plugins.anthropic.model_info import get_model_info -from genkit.plugins.anthropic.utils import ( - build_cache_usage, - get_cache_control, - maybe_strip_fences, - to_anthropic_media, -) -from genkit.types import ( +from genkit import ( FinishReason, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationUsage, MediaPart, Message, + ModelRequest, + ModelResponse, + ModelResponseChunk, + ModelUsage, Part, Role, TextPart, @@ -54,8 +46,17 @@ ToolRequestPart, ToolResponsePart, ) +from genkit.model import get_basic_usage_stats +from genkit.plugin_api import ActionRunContext +from genkit.plugins.anthropic.model_info import get_model_info +from genkit.plugins.anthropic.utils import ( + build_cache_usage, + get_cache_control, + maybe_strip_fences, + to_anthropic_media, +) -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) DEFAULT_MAX_OUTPUT_TOKENS = 4096 @@ -107,7 +108,7 @@ def __init__(self, model_name: str, client: AsyncAnthropic) -> None: self.model_name = model_info.versions[0] if model_info.versions else model_name self.client = client - async def generate(self, request: GenerateRequest, ctx: ActionRunContext | None = None) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext | None = None) -> ModelResponse: """Generate response from Anthropic. Args: @@ -155,13 +156,13 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext | None # Build usage with cache-aware token counts. usage = self._build_usage(response, basic_usage) - return GenerateResponse( + return ModelResponse( message=response_message, usage=usage, finish_reason=finish_reason, ) - def _build_usage(self, response: AnthropicMessage, basic_usage: GenerationUsage) -> GenerationUsage: + def _build_usage(self, response: AnthropicMessage, basic_usage: ModelUsage) -> ModelUsage: """Build usage stats including cache read/write token counts. Delegates to :func:`utils.build_cache_usage` for the actual @@ -172,7 +173,7 @@ def _build_usage(self, response: AnthropicMessage, basic_usage: GenerationUsage) basic_usage: Basic character/image usage from message content. Returns: - GenerationUsage with token and character counts. + ModelUsage with token and character counts. """ return build_cache_usage( input_tokens=response.usage.input_tokens, @@ -182,7 +183,7 @@ def _build_usage(self, response: AnthropicMessage, basic_usage: GenerationUsage) cache_read_input_tokens=getattr(response.usage, 'cache_read_input_tokens', None) or 0, ) - def _build_params(self, request: GenerateRequest) -> dict[str, Any]: + def _build_params(self, request: ModelRequest) -> dict[str, Any]: """Build Anthropic API parameters.""" config = request.config params: dict[str, Any] = {} @@ -191,7 +192,7 @@ def _build_params(self, request: GenerateRequest) -> dict[str, Any]: params = config.copy() elif config: if hasattr(config, 'model_dump'): - params = config.model_dump(exclude_none=True) + params = config.model_dump(exclude_none=True, by_alias=False) else: params = {k: v for k, v in vars(config).items() if v is not None} @@ -233,21 +234,21 @@ def _build_params(self, request: GenerateRequest) -> dict[str, Any]: system = self._extract_system(request.messages) # Handle JSON output constraint - if request.output and request.output.format == 'json': + if request.output_format == 'json': supports_json = 'json' in (self._model_info.supports.output or []) if self._model_info.supports else False - if request.output.schema and supports_json: + if request.output_schema and supports_json: # Use native structured outputs via output_config. params['output_config'] = { 'format': { 'type': 'json_schema', - 'schema': _to_anthropic_schema(request.output.schema), + 'schema': _to_anthropic_schema(request.output_schema), } } else: # Fall back to system prompt instruction. instruction = '\n\nOutput valid JSON. Do not wrap the JSON in markdown code fences.' - if request.output.schema: - schema_str = json.dumps(request.output.schema, indent=2) + if request.output_schema: + schema_str = json.dumps(request.output_schema, indent=2) instruction += f'\n\nFollow this JSON schema:\n{schema_str}' system = (system or '') + instruction @@ -285,7 +286,7 @@ async def _generate_streaming(self, params: dict[str, Any], ctx: ActionRunContex 3. ``content_block_stop`` We track in-progress tool calls and emit a - :class:`GenerateResponseChunk` containing the tool request when + :class:`ModelResponseChunk` containing the tool request when the block finishes. """ # Track in-progress tool-use blocks by index. @@ -308,7 +309,7 @@ async def _generate_streaming(self, params: dict[str, Any], ctx: ActionRunContex delta = chunk.delta if getattr(delta, 'type', None) == 'text_delta' and hasattr(delta, 'text'): ctx.send_chunk( - GenerateResponseChunk( + ModelResponseChunk( role=Role.MODEL, index=0, content=[Part(root=TextPart(text=str(delta.text)))], @@ -330,7 +331,7 @@ async def _generate_streaming(self, params: dict[str, Any], ctx: ActionRunContex except (json.JSONDecodeError, TypeError): tool_input = tool_info['input_json'] ctx.send_chunk( - GenerateResponseChunk( + ModelResponseChunk( role=Role.MODEL, index=0, content=[ diff --git a/py/plugins/anthropic/src/genkit/plugins/anthropic/plugin.py b/py/plugins/anthropic/src/genkit/plugins/anthropic/plugin.py index 3a41e74bc6..86330b8f48 100644 --- a/py/plugins/anthropic/src/genkit/plugins/anthropic/plugin.py +++ b/py/plugins/anthropic/src/genkit/plugins/anthropic/plugin.py @@ -19,15 +19,19 @@ from typing import Any, cast from anthropic import AsyncAnthropic -from genkit.ai import ActionRunContext, Plugin -from genkit.blocks.model import model_action_metadata -from genkit.core._loop_local import _loop_local_client -from genkit.core.action import Action, ActionMetadata -from genkit.core.registry import ActionKind -from genkit.core.schema import to_json_schema +from genkit import ModelConfig, ModelRequest, ModelResponse +from genkit.model import model_action_metadata +from genkit.plugin_api import ( + Action, + ActionKind, + ActionMetadata, + ActionRunContext, + Plugin, + loop_local_client, + to_json_schema, +) from genkit.plugins.anthropic.model_info import SUPPORTED_ANTHROPIC_MODELS, get_model_info from genkit.plugins.anthropic.models import AnthropicModel -from genkit.types import GenerateRequest, GenerateResponse, GenerationCommonConfig ANTHROPIC_PLUGIN_NAME = 'anthropic' @@ -67,9 +71,7 @@ def __init__( """ self.models = models or list(SUPPORTED_ANTHROPIC_MODELS.keys()) self._anthropic_params = anthropic_params - self._runtime_client = _loop_local_client( - lambda: AsyncAnthropic(**cast(dict[str, Any], self._anthropic_params)) - ) + self._runtime_client = loop_local_client(lambda: AsyncAnthropic(**cast(dict[str, Any], self._anthropic_params))) async def init(self) -> list[Action]: """Initialize plugin. @@ -108,7 +110,7 @@ def _create_model_action(self, name: str) -> Action: model_info = get_model_info(clean_name) - async def _generate(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def _generate(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: model = AnthropicModel(model_name=clean_name, client=self._runtime_client()) return await model.generate(request, ctx) @@ -121,7 +123,7 @@ async def _generate(request: GenerateRequest, ctx: ActionRunContext) -> Generate 'supports': ( model_info.supports.model_dump(by_alias=True, exclude_none=True) if model_info.supports else {} ), - 'customOptions': to_json_schema(GenerationCommonConfig), + 'customOptions': to_json_schema(ModelConfig), }, }, ) @@ -138,7 +140,7 @@ async def list_actions(self) -> list[ActionMetadata]: model_action_metadata( name=anthropic_name(model_name), info=model_info.model_dump(by_alias=True, exclude_none=True), - config_schema=GenerationCommonConfig, + config_schema=ModelConfig, ) ) return actions diff --git a/py/plugins/anthropic/src/genkit/plugins/anthropic/utils.py b/py/plugins/anthropic/src/genkit/plugins/anthropic/utils.py index ccc7b0a010..8b994a6a23 100644 --- a/py/plugins/anthropic/src/genkit/plugins/anthropic/utils.py +++ b/py/plugins/anthropic/src/genkit/plugins/anthropic/utils.py @@ -27,10 +27,11 @@ import re from typing import Any -from genkit.core.logging import get_logger -from genkit.types import GenerateRequest, GenerationUsage, MediaPart, Part, TextPart +import structlog -logger = get_logger(__name__) +from genkit import MediaPart, ModelRequest, ModelUsage, Part, TextPart + +logger = structlog.get_logger(__name__) # PDF MIME type for document handling. PDF_MIME_TYPE = 'application/pdf' @@ -76,7 +77,7 @@ def strip_markdown_fences(text: str) -> str: return text -def maybe_strip_fences(request: GenerateRequest, parts: list[Part]) -> list[Part]: +def maybe_strip_fences(request: ModelRequest, parts: list[Part]) -> list[Part]: """Strip markdown fences from text parts when JSON output is expected. Args: @@ -86,7 +87,7 @@ def maybe_strip_fences(request: GenerateRequest, parts: list[Part]) -> list[Part Returns: Parts with fences stripped from text if JSON was requested. """ - if not request.output or request.output.format != 'json': + if request.output_format != 'json': return parts cleaned: list[Part] = [] @@ -234,11 +235,11 @@ def to_anthropic_media(media_part: MediaPart) -> dict[str, Any]: def build_cache_usage( input_tokens: int, output_tokens: int, - basic_usage: GenerationUsage, + basic_usage: ModelUsage, cache_creation_input_tokens: int = 0, cache_read_input_tokens: int = 0, -) -> GenerationUsage: - """Build GenerationUsage with cache-aware token counts. +) -> ModelUsage: + """Build ModelUsage with cache-aware token counts. Args: input_tokens: Number of input tokens from the API response. @@ -248,7 +249,7 @@ def build_cache_usage( cache_read_input_tokens: Tokens read from existing cache entries. Returns: - GenerationUsage with token, character, and cache counts. + ModelUsage with token, character, and cache counts. """ custom: dict[str, float] = {} if cache_creation_input_tokens: @@ -256,7 +257,7 @@ def build_cache_usage( if cache_read_input_tokens: custom['cache_read_input_tokens'] = cache_read_input_tokens - return GenerationUsage( + return ModelUsage( input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=input_tokens + output_tokens, diff --git a/py/plugins/anthropic/tests/anthropic_models_test.py b/py/plugins/anthropic/tests/anthropic_models_test.py index 2388a9414f..583daad8f8 100644 --- a/py/plugins/anthropic/tests/anthropic_models_test.py +++ b/py/plugins/anthropic/tests/anthropic_models_test.py @@ -21,35 +21,34 @@ import pytest -from genkit.plugins.anthropic.models import AnthropicModel -from genkit.plugins.anthropic.utils import maybe_strip_fences, strip_markdown_fences -from genkit.types import ( - GenerateRequest, - GenerateResponseChunk, - GenerationCommonConfig, +from genkit import ( Media, MediaPart, Message, Metadata, - OutputConfig, + ModelConfig, + ModelRequest, + ModelResponseChunk, Part, Role, TextPart, ToolDefinition, ToolRequestPart, ) +from genkit.plugins.anthropic.models import AnthropicModel +from genkit.plugins.anthropic.utils import maybe_strip_fences, strip_markdown_fences -def _create_sample_request() -> GenerateRequest: +def _create_sample_request() -> ModelRequest: """Create a sample generation request for testing.""" - return GenerateRequest( + return ModelRequest( messages=[ Message( role=Role.USER, content=[Part(root=TextPart(text='Hello, how are you?'))], ) ], - config=GenerationCommonConfig(), + config=ModelConfig(), tools=[ ToolDefinition( name='get_weather', @@ -139,9 +138,9 @@ async def test_generate_with_config() -> None: model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Test'))])], - config=GenerationCommonConfig( + config=ModelConfig( temperature=0.7, max_output_tokens=100, top_p=0.9, @@ -244,9 +243,9 @@ async def test_streaming_generation() -> None: ctx = MagicMock() ctx.is_streaming = True - collected_chunks: list[GenerateResponseChunk] = [] + collected_chunks: list[ModelResponseChunk] = [] - def send_chunk(chunk: GenerateResponseChunk) -> None: + def send_chunk(chunk: ModelResponseChunk) -> None: collected_chunks.append(chunk) ctx.send_chunk = send_chunk @@ -318,7 +317,7 @@ async def test_streaming_tool_request() -> None: ctx = MagicMock() ctx.is_streaming = True - collected_chunks: list[GenerateResponseChunk] = [] + collected_chunks: list[ModelResponseChunk] = [] ctx.send_chunk = lambda chunk: collected_chunks.append(chunk) response = await model.generate(sample_request, ctx) @@ -381,9 +380,10 @@ class TestMaybeStripFences: def test_strips_fences_for_json_output(self) -> None: """Strips markdown fences when JSON output is requested.""" - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='json', schema={'type': 'object'}), + output_format='json', + output_schema={'type': 'object'}, ) parts = [Part(root=TextPart(text='```json\n{"a": 1}\n```'))] result = maybe_strip_fences(request, parts) @@ -391,9 +391,9 @@ def test_strips_fences_for_json_output(self) -> None: def test_no_op_for_text_output(self) -> None: """Does not modify responses when output format is not json.""" - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='text'), + output_format='text', ) fenced = '```json\n{"a": 1}\n```' parts = [Part(root=TextPart(text=fenced))] @@ -402,7 +402,7 @@ def test_no_op_for_text_output(self) -> None: def test_no_op_for_no_output(self) -> None: """Does not modify responses when no output config is set.""" - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], ) fenced = '```json\n{"a": 1}\n```' @@ -412,9 +412,10 @@ def test_no_op_for_no_output(self) -> None: def test_no_op_when_no_fences(self) -> None: """Does not modify clean JSON responses.""" - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='json', schema={'type': 'object'}), + output_format='json', + output_schema={'type': 'object'}, ) text = '{"name": "John"}' parts = [Part(root=TextPart(text=text))] @@ -487,7 +488,7 @@ async def test_cache_token_tracking_in_usage() -> None: mock_client.messages.create = AsyncMock(return_value=mock_response) model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Test'))])], ) @@ -516,7 +517,7 @@ async def test_no_cache_tokens_when_caching_not_used() -> None: mock_client.messages.create = AsyncMock(return_value=mock_response) model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Test'))])], ) @@ -635,12 +636,10 @@ def test_structured_output_uses_native_output_config() -> None: mock_client = MagicMock() model = AnthropicModel(model_name='claude-opus-4-6', client=mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a cat'))])], - output=OutputConfig( - format='json', - schema={'type': 'object', 'properties': {'name': {'type': 'string'}}}, - ), + output_format='json', + output_schema={'type': 'object', 'properties': {'name': {'type': 'string'}}}, ) params = model._build_params(request) @@ -655,9 +654,9 @@ def test_structured_output_falls_back_to_system_prompt() -> None: mock_client = MagicMock() model = AnthropicModel(model_name='claude-sonnet-4', client=mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Generate JSON'))])], - output=OutputConfig(format='json'), + output_format='json', ) params = model._build_params(request) @@ -673,12 +672,10 @@ def test_structured_output_falls_back_for_unsupported_models() -> None: # Claude 3.5 Haiku is marked as not supporting JSON natively in model_info.py model = AnthropicModel(model_name='claude-3-5-haiku', client=mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a cat'))])], - output=OutputConfig( - format='json', - schema={'type': 'object', 'properties': {'name': {'type': 'string'}}}, - ), + output_format='json', + output_schema={'type': 'object', 'properties': {'name': {'type': 'string'}}}, ) params = model._build_params(request) diff --git a/py/plugins/anthropic/tests/anthropic_plugin_test.py b/py/plugins/anthropic/tests/anthropic_plugin_test.py index 3437b312a2..d4edf65217 100644 --- a/py/plugins/anthropic/tests/anthropic_plugin_test.py +++ b/py/plugins/anthropic/tests/anthropic_plugin_test.py @@ -23,21 +23,21 @@ import pytest -from genkit.core.registry import ActionKind -from genkit.plugins.anthropic import Anthropic, anthropic_name -from genkit.plugins.anthropic.model_info import ( - SUPPORTED_ANTHROPIC_MODELS as SUPPORTED_MODELS, - get_model_info, -) -from genkit.types import ( - GenerateRequest, - GenerationCommonConfig, +from genkit import ( + ActionKind, Message, + ModelConfig, + ModelRequest, Part, Role, TextPart, ToolDefinition, ) +from genkit.plugins.anthropic import Anthropic, anthropic_name +from genkit.plugins.anthropic.model_info import ( + SUPPORTED_ANTHROPIC_MODELS as SUPPORTED_MODELS, + get_model_info, +) def test_anthropic_name() -> None: @@ -186,16 +186,16 @@ def test_get_model_info_unknown() -> None: assert info.supports.tools is True -def _create_sample_request() -> GenerateRequest: +def _create_sample_request() -> ModelRequest: """Create a sample generation request for testing.""" - return GenerateRequest( + return ModelRequest( messages=[ Message( role=Role.USER, content=[Part(root=TextPart(text='Hello, how are you?'))], ) ], - config=GenerationCommonConfig(), + config=ModelConfig(), tools=[ ToolDefinition( name='get_weather', diff --git a/py/plugins/anthropic/tests/anthropic_utils_test.py b/py/plugins/anthropic/tests/anthropic_utils_test.py index d1321217c2..adc8f690b8 100644 --- a/py/plugins/anthropic/tests/anthropic_utils_test.py +++ b/py/plugins/anthropic/tests/anthropic_utils_test.py @@ -23,6 +23,13 @@ import base64 +from genkit import ( + Media, + MediaPart, + Metadata, + ModelUsage, + TextPart, +) from genkit.plugins.anthropic.utils import ( DOCUMENT_MIME_TYPES, PDF_MIME_TYPE, @@ -33,13 +40,6 @@ to_anthropic_image, to_anthropic_media, ) -from genkit.types import ( - GenerationUsage, - Media, - MediaPart, - Metadata, - TextPart, -) # --------------------------------------------------------------------------- # get_cache_control tests @@ -232,7 +232,7 @@ class TestBuildCacheUsage: def test_basic_usage_without_cache(self) -> None: """Builds usage without cache tokens.""" - basic = GenerationUsage(input_characters=10, output_characters=20) + basic = ModelUsage(input_characters=10, output_characters=20) result = build_cache_usage( input_tokens=100, output_tokens=50, @@ -247,7 +247,7 @@ def test_basic_usage_without_cache(self) -> None: def test_usage_with_cache_creation(self) -> None: """Includes cache_creation_input_tokens in custom.""" - basic = GenerationUsage() + basic = ModelUsage() result = build_cache_usage( input_tokens=100, output_tokens=50, @@ -260,7 +260,7 @@ def test_usage_with_cache_creation(self) -> None: def test_usage_with_cache_read(self) -> None: """Includes cache_read_input_tokens in custom.""" - basic = GenerationUsage() + basic = ModelUsage() result = build_cache_usage( input_tokens=100, output_tokens=50, @@ -273,7 +273,7 @@ def test_usage_with_cache_read(self) -> None: def test_usage_with_both_cache_fields(self) -> None: """Includes both cache token fields when both are present.""" - basic = GenerationUsage() + basic = ModelUsage() result = build_cache_usage( input_tokens=100, output_tokens=50, @@ -287,7 +287,7 @@ def test_usage_with_both_cache_fields(self) -> None: def test_zero_cache_tokens_are_excluded(self) -> None: """Zero cache tokens don't appear in custom field.""" - basic = GenerationUsage() + basic = ModelUsage() result = build_cache_usage( input_tokens=100, output_tokens=50, diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/audio.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/audio.py index 93f2d86852..9317b33928 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/audio.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/audio.py @@ -25,7 +25,7 @@ Data Flow (TTS):: β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ GenerateRequest (text input) β”‚ + β”‚ ModelRequest (text input) β”‚ β”‚ β”‚ β”‚ β”‚ β–Ό β”‚ β”‚ to_tts_params() ──► SpeechCreateParams β”‚ @@ -34,13 +34,13 @@ β”‚ client.audio.speech.create() β”‚ β”‚ β”‚ β”‚ β”‚ β–Ό β”‚ - β”‚ to_tts_response() ──► GenerateResponse (audio media part) β”‚ + β”‚ to_tts_response() ──► ModelResponse (audio media part) β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ Data Flow (STT):: β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ GenerateRequest (audio media input) β”‚ + β”‚ ModelRequest (audio media input) β”‚ β”‚ β”‚ β”‚ β”‚ β–Ό β”‚ β”‚ to_stt_params() ──► TranscriptionCreateParams β”‚ @@ -49,7 +49,7 @@ β”‚ client.audio.transcriptions.create() β”‚ β”‚ β”‚ β”‚ β”‚ β–Ό β”‚ - β”‚ to_stt_response() ──► GenerateResponse (text part) β”‚ + β”‚ to_stt_response() ──► ModelResponse (text part) β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ """ @@ -62,27 +62,27 @@ from openai._legacy_response import HttpxBinaryResponseContent from openai.types.audio import Transcription -from genkit.ai import ActionRunContext -from genkit.core.typing import FinishReason -from genkit.plugins.compat_oai.models.utils import ( - _extract_media, - _extract_text, - _find_text, - decode_data_uri_bytes, - extract_config_dict, -) -from genkit.types import ( - GenerateRequest, - GenerateResponse, +from genkit import ( Media, MediaPart, Message, ModelInfo, + ModelRequest, + ModelResponse, Part, Role, Supports, TextPart, ) +from genkit.model import FinishReason +from genkit.plugin_api import ActionRunContext +from genkit.plugins.compat_oai.models.utils import ( + _extract_media, + _extract_text, + _find_text, + decode_data_uri_bytes, + extract_config_dict, +) # Maps audio response formats to their MIME types. RESPONSE_FORMAT_MEDIA_TYPES: dict[str, str] = { @@ -176,9 +176,9 @@ def _to_tts_params( model_name: str, - request: GenerateRequest, + request: ModelRequest, ) -> dict[str, Any]: - """Convert a GenerateRequest into OpenAI TTS parameters. + """Convert a ModelRequest into OpenAI TTS parameters. Args: model_name: The TTS model name (e.g., 'tts-1'). @@ -202,7 +202,7 @@ def _to_tts_params( params[key] = config.pop(key) # Strip standard GenAI config keys. - for key in ('temperature', 'maxOutputTokens', 'stopSequences', 'topK', 'topP'): + for key in ('temperature', 'max_output_tokens', 'stop_sequences', 'top_k', 'top_p'): config.pop(key, None) return {k: v for k, v in params.items() if v is not None} @@ -211,8 +211,8 @@ def _to_tts_params( def _to_tts_response( response: HttpxBinaryResponseContent, response_format: str = 'mp3', -) -> GenerateResponse: - """Convert an OpenAI speech response to a Genkit GenerateResponse. +) -> ModelResponse: + """Convert an OpenAI speech response to a Genkit ModelResponse. The response body is read as bytes and encoded as a base64 data URI. @@ -221,7 +221,7 @@ def _to_tts_response( response_format: The audio format used (determines MIME type). Returns: - A GenerateResponse with a media part containing the audio data. + A ModelResponse with a media part containing the audio data. """ # The response from speech.create() is an HttpxBinaryResponseContent # which supports .read() to get raw bytes. @@ -229,7 +229,7 @@ def _to_tts_response( media_type = RESPONSE_FORMAT_MEDIA_TYPES.get(response_format, 'audio/mpeg') b64_data = base64.b64encode(audio_bytes).decode('ascii') - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=[ @@ -249,9 +249,9 @@ def _to_tts_response( def _to_stt_params( model_name: str, - request: GenerateRequest, + request: ModelRequest, ) -> dict[str, Any]: - """Convert a GenerateRequest into OpenAI transcription parameters. + """Convert a ModelRequest into OpenAI transcription parameters. Extracts the audio media from the first message and converts it into a file-like object suitable for the transcriptions API. @@ -290,19 +290,19 @@ def _to_stt_params( # Determine response format: config override > output format > default. response_format = config.pop('response_format', None) - if not response_format and request.output and request.output.format in ('json', 'text'): - response_format = request.output.format + if not response_format and request.output_format and request.output_format in ('json', 'text'): + response_format = request.output_format params['response_format'] = response_format or 'text' # Strip standard GenAI config keys. - for key in ('maxOutputTokens', 'stopSequences', 'topK', 'topP'): + for key in ('max_output_tokens', 'stop_sequences', 'top_k', 'top_p'): config.pop(key, None) return {k: v for k, v in params.items() if v is not None} -def _to_stt_response(result: Transcription | str) -> GenerateResponse: - """Convert an OpenAI transcription result to a Genkit GenerateResponse. +def _to_stt_response(result: Transcription | str) -> ModelResponse: + """Convert an OpenAI transcription result to a Genkit ModelResponse. Handles the full union of types returned by transcriptions.create(). All non-str result types (Transcription, TranscriptionVerbose, @@ -313,7 +313,7 @@ def _to_stt_response(result: Transcription | str) -> GenerateResponse: object with a .text attribute, or a plain string). Returns: - A GenerateResponse with a text part containing the transcription. + A ModelResponse with a text part containing the transcription. """ if isinstance(result, str): text = result @@ -321,7 +321,7 @@ def _to_stt_response(result: Transcription | str) -> GenerateResponse: text = result.text else: text = str(result) - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=[Part(root=TextPart(text=text))], @@ -353,7 +353,7 @@ def name(self) -> str: """The name of the TTS model.""" return self._model_name - async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: """Generate speech audio from the request. Args: @@ -361,7 +361,7 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> Gen ctx: The action run context. Returns: - A GenerateResponse containing audio media parts. + A ModelResponse containing audio media parts. """ params = _to_tts_params(self._model_name, request) response_format = params.get('response_format', 'mp3') @@ -392,7 +392,7 @@ def name(self) -> str: """The name of the STT model.""" return self._model_name - async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: """Transcribe audio from the request. Args: @@ -400,7 +400,7 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> Gen ctx: The action run context. Returns: - A GenerateResponse containing the transcribed text. + A ModelResponse containing the transcribed text. """ params = _to_stt_params(self._model_name, request) result = await self._client.audio.transcriptions.create( diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/handler.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/handler.py index c096ba56f2..4f165233f3 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/handler.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/handler.py @@ -20,18 +20,18 @@ from openai import AsyncOpenAI -from genkit.ai import ActionRunContext +from genkit import ( + ModelInfo, + ModelRequest, + ModelResponse, +) +from genkit.plugin_api import ActionRunContext from genkit.plugins.compat_oai.models.model import OpenAIModel from genkit.plugins.compat_oai.models.model_info import ( SUPPORTED_OPENAI_COMPAT_MODELS, SUPPORTED_OPENAI_MODELS, PluginSource, ) -from genkit.types import ( - GenerateRequest, - GenerateResponse, - ModelInfo, -) class OpenAIModelHandler: @@ -66,14 +66,14 @@ def _get_supported_models(source: PluginSource) -> dict[str, ModelInfo]: @classmethod def get_model_handler( cls, model: str, client: AsyncOpenAI, source: PluginSource = PluginSource.OPENAI - ) -> Callable[[GenerateRequest, ActionRunContext], Awaitable[GenerateResponse]]: + ) -> Callable[[ModelRequest, ActionRunContext], Awaitable[ModelResponse]]: """Factory method to initialize the model handler for the specified OpenAI model. OpenAI models in this context are not instantiated as traditional classes but rather as Actions. This method returns a callable that serves as an action handler, conforming to the structure of: - Action[GenerateRequest, GenerateResponse, GenerateResponseChunk] + Action[ModelRequest, ModelResponse, ModelResponseChunk] Args: model: The OpenAI model name. @@ -109,7 +109,7 @@ def _validate_version(self, version: str) -> None: if model_info.versions is not None and version not in model_info.versions: raise ValueError(f"Model version '{version}' is not supported.") - async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: """Processes the request using OpenAI's chat completion API. Args: @@ -117,7 +117,7 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> Gen ctx: The context of the action run. Returns: - A GenerateResponse containing the model's response. + A ModelResponse containing the model's response. Raises: ValueError: If the specified model version is not supported. diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/image.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/image.py index f42ae44816..033ead196f 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/image.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/image.py @@ -22,7 +22,7 @@ Data Flow:: β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ GenerateRequest (text prompt) β”‚ + β”‚ ModelRequest (text prompt) β”‚ β”‚ β”‚ β”‚ β”‚ β–Ό β”‚ β”‚ to_image_generate_params() ──► ImageGenerateParams β”‚ @@ -31,7 +31,7 @@ β”‚ client.images.generate() β”‚ β”‚ β”‚ β”‚ β”‚ β–Ό β”‚ - β”‚ to_generate_response() ──► GenerateResponse (media parts) β”‚ + β”‚ to_generate_response() ──► ModelResponse (media parts) β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ """ @@ -42,20 +42,20 @@ from openai import AsyncOpenAI from openai.types.images_response import ImagesResponse -from genkit.ai import ActionRunContext -from genkit.core.typing import FinishReason -from genkit.plugins.compat_oai.models.utils import _extract_text, extract_config_dict -from genkit.types import ( - GenerateRequest, - GenerateResponse, +from genkit import ( Media, MediaPart, Message, ModelInfo, + ModelRequest, + ModelResponse, Part, Role, Supports, ) +from genkit.model import FinishReason +from genkit.plugin_api import ActionRunContext +from genkit.plugins.compat_oai.models.utils import _extract_text, extract_config_dict # Supported image generation models with their metadata. SUPPORTED_IMAGE_MODELS: dict[str, ModelInfo] = { @@ -88,9 +88,9 @@ def _to_image_generate_params( model_name: str, - request: GenerateRequest, + request: ModelRequest, ) -> dict[str, Any]: - """Convert a GenerateRequest into OpenAI image generation parameters. + """Convert a ModelRequest into OpenAI image generation parameters. Extracts the text prompt and maps Genkit config options to OpenAI's image generation API parameters. @@ -113,7 +113,7 @@ def _to_image_generate_params( } # Strip standard GenAI config keys that don't apply to image generation. - for key in ('temperature', 'maxOutputTokens', 'stopSequences', 'topK', 'topP'): + for key in ('temperature', 'max_output_tokens', 'stop_sequences', 'top_k', 'top_p'): config.pop(key, None) # Pass remaining config through (size, quality, style, n, etc.). @@ -123,8 +123,8 @@ def _to_image_generate_params( return {k: v for k, v in params.items() if v is not None} -def _to_generate_response(result: ImagesResponse) -> GenerateResponse: - """Convert an OpenAI ImagesResponse to a Genkit GenerateResponse. +def _to_generate_response(result: ImagesResponse) -> ModelResponse: + """Convert an OpenAI ImagesResponse to a Genkit ModelResponse. Each generated image becomes a media part in the response message. @@ -132,11 +132,11 @@ def _to_generate_response(result: ImagesResponse) -> GenerateResponse: result: The OpenAI images.generate() response object. Returns: - A GenerateResponse with media parts for each generated image. + A ModelResponse with media parts for each generated image. """ images = result.data if not images: - return GenerateResponse( + return ModelResponse( message=Message(role=Role.MODEL, content=[]), finish_reason=FinishReason.STOP, ) @@ -150,7 +150,7 @@ def _to_generate_response(result: ImagesResponse) -> GenerateResponse: if url: content.append(Part(root=MediaPart(media=Media(content_type='image/png', url=url)))) - return GenerateResponse( + return ModelResponse( message=Message(role=Role.MODEL, content=content), finish_reason=FinishReason.STOP, ) @@ -179,7 +179,7 @@ def name(self) -> str: """The name of the image model.""" return self._model_name - async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: """Generate images from the request. Args: @@ -187,7 +187,7 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> Gen ctx: The action run context. Returns: - A GenerateResponse containing generated image media parts. + A ModelResponse containing generated image media parts. """ params = _to_image_generate_params(self._model_name, request) result = await self._client.images.generate(**params) diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py index 9adb1b56c2..c66b8d3a5d 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py @@ -20,12 +20,23 @@ from collections.abc import Callable from typing import Any, cast +import structlog from openai import AsyncOpenAI from openai.lib._pydantic import _ensure_strict_json_schema -from genkit.core.action._action import ActionRunContext -from genkit.core.logging import get_logger -from genkit.core.typing import GenerationCommonConfig as CoreGenerationCommonConfig +from genkit import ( + Message, + ModelConfig, + ModelRequest, + ModelResponse, + ModelResponseChunk, + Part, + ReasoningPart, + Role, + TextPart, + ToolDefinition, +) +from genkit.plugin_api import ActionRunContext from genkit.plugins.compat_oai.models.model_info import SUPPORTED_OPENAI_MODELS from genkit.plugins.compat_oai.models.utils import ( DictMessageAdapter, @@ -34,21 +45,8 @@ strip_markdown_fences, ) from genkit.plugins.compat_oai.typing import OpenAIConfig, SupportedOutputFormat -from genkit.types import ( - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationCommonConfig, - Message, - OutputConfig, - Part, - ReasoningPart, - Role, - TextPart, - ToolDefinition, -) -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) class OpenAIModel: @@ -121,7 +119,7 @@ async def _get_tools_definition(self, tools: list[ToolDefinition]) -> list[dict] result.append(function_call) return result - def _needs_schema_in_prompt(self, output: OutputConfig) -> bool: + def _needs_schema_in_prompt(self, request: ModelRequest) -> bool: """Check whether the schema must be injected into the prompt. Models that only support ``json_object`` mode (e.g. DeepSeek) never @@ -130,21 +128,21 @@ def _needs_schema_in_prompt(self, output: OutputConfig) -> bool: knows what structure to produce. Args: - output: The output configuration. + request: The model request with output_format and output_schema. Returns: True when the schema should be injected into the messages. """ - if output.format != 'json' or not output.schema: + if request.output_format != 'json' or not request.output_schema: return False # DeepSeek models use json_object mode β€” schema never reaches the API. return self._model.startswith('deepseek') - def _get_response_format(self, output: OutputConfig) -> dict | None: + def _get_response_format(self, request: ModelRequest) -> dict | None: """Determines the response format configuration based on the output settings. Args: - output: The output configuration specifying the desired format and optional schema. + request: The model request with output_format and output_schema. Returns: A dictionary representing the response format, which may include: @@ -152,17 +150,19 @@ def _get_response_format(self, output: OutputConfig) -> dict | None: - 'type': 'json_object' if the model supports JSON mode and no schema is provided. - 'type': 'text' as the default fallback. """ - if output.format == 'json': + if request.output_format == 'json': # DeepSeek models: always use 'json_object' (schema is injected # into the prompt by _get_openai_request_config instead). if self._model.startswith('deepseek'): return {'type': 'json_object'} - if output.schema: + if request.output_schema: return { 'type': 'json_schema', 'json_schema': { - 'name': output.schema.get('title', 'Response'), - 'schema': _ensure_strict_json_schema(output.schema, path=(), root=output.schema), + 'name': request.output_schema.get('title', 'Response'), + 'schema': _ensure_strict_json_schema( + request.output_schema, path=(), root=request.output_schema + ), 'strict': True, }, } @@ -173,7 +173,7 @@ def _get_response_format(self, output: OutputConfig) -> dict | None: return {'type': 'text'} - def _clean_json_response(self, response: 'GenerateResponse', request: 'GenerateRequest') -> 'GenerateResponse': + def _clean_json_response(self, response: ModelResponse, request: ModelRequest) -> ModelResponse: """Strip markdown fences from JSON responses for json_object-mode models. Only applies when the model uses ``json_object`` mode (e.g. DeepSeek) @@ -186,12 +186,7 @@ def _clean_json_response(self, response: 'GenerateResponse', request: 'GenerateR Returns: The response with cleaned text parts, or the original response. """ - if ( - not request.output - or request.output.format != 'json' - or not self._model.startswith('deepseek') - or response.message is None - ): + if request.output_format != 'json' or not self._model.startswith('deepseek') or response.message is None: return response cleaned_parts: list[Part] = [] @@ -208,7 +203,7 @@ def _clean_json_response(self, response: 'GenerateResponse', request: 'GenerateR cleaned_parts.append(part) if changed: - return GenerateResponse( + return ModelResponse( request=request, message=Message(role=response.message.role, content=cleaned_parts), finish_reason=response.finish_reason, @@ -244,7 +239,7 @@ def _build_schema_instruction(schema: dict[str, Any]) -> dict[str, str]: ), } - async def _get_openai_request_config(self, request: GenerateRequest) -> dict: + async def _get_openai_request_config(self, request: ModelRequest) -> dict: """Get the OpenAI request configuration. Args: @@ -257,8 +252,8 @@ async def _get_openai_request_config(self, request: GenerateRequest) -> dict: # For models that only support json_object mode, inject the schema # into the messages so the model knows the expected output structure. - if request.output and self._needs_schema_in_prompt(request.output) and request.output.schema: - schema_msg = self._build_schema_instruction(request.output.schema) + if self._needs_schema_in_prompt(request) and request.output_schema: + schema_msg = self._build_schema_instruction(request.output_schema) messages = [schema_msg, *messages] openai_config: dict[str, Any] = { @@ -272,8 +267,8 @@ async def _get_openai_request_config(self, request: GenerateRequest) -> dict: openai_config['tool_choice'] = 'none' elif request.tool_choice: openai_config['tool_choice'] = request.tool_choice - if request.output: - response_format = self._get_response_format(request.output) + if request.output_format: + response_format = self._get_response_format(request) if response_format: # pyrefly: ignore[bad-typed-dict-key] - response_format dict is valid for OpenAI API openai_config['response_format'] = response_format @@ -281,14 +276,14 @@ async def _get_openai_request_config(self, request: GenerateRequest) -> dict: openai_config.update(**request.config.model_dump(exclude_none=True)) return openai_config - async def _generate(self, request: GenerateRequest) -> GenerateResponse: + async def _generate(self, request: ModelRequest) -> ModelResponse: """Processes the request using OpenAI's chat completion API and returns the generated response. Args: - request: The GenerateRequest object containing the input text and configuration. + request: The ModelRequest object containing the input text and configuration. Returns: - A GenerateResponse object containing the generated message. + A ModelResponse object containing the generated message. """ openai_config = await self._get_openai_request_config(request=request) logger.debug('OpenAI generate request', model=self._model, streaming=False) @@ -299,23 +294,23 @@ async def _generate(self, request: GenerateRequest) -> GenerateResponse: finish_reason=str(response.choices[0].finish_reason) if response.choices else None, ) - result = GenerateResponse( + result = ModelResponse( request=request, message=MessageConverter.to_genkit(MessageAdapter(response.choices[0].message)), ) return self._clean_json_response(result, request) async def _generate_stream( - self, request: GenerateRequest, callback: Callable[[GenerateResponseChunk], None] - ) -> GenerateResponse: + self, request: ModelRequest, callback: Callable[[ModelResponseChunk], None] + ) -> ModelResponse: """Streams responses from the OpenAI client and sends chunks to a callback. Args: - request: The GenerateRequest object containing generation parameters. - callback: A function to receive streamed GenerateResponseChunk objects. + request: The ModelRequest object containing generation parameters. + callback: A function to receive streamed ModelResponseChunk objects. Returns: - GenerateResponse: A final message with accumulated content after streaming is complete. + ModelResponse: A final message with accumulated content after streaming is complete. """ openai_config = await self._get_openai_request_config(request=request) openai_config['stream'] = True @@ -332,7 +327,7 @@ async def _generate_stream( message = MessageConverter.to_genkit(MessageAdapter(delta)) accumulated_content.extend(message.content) callback( - GenerateResponseChunk( + ModelResponseChunk( role=Role.MODEL, content=message.content, ) @@ -345,7 +340,7 @@ async def _generate_stream( reasoning_part = Part(root=ReasoningPart(reasoning=reasoning_text)) accumulated_content.append(reasoning_part) callback( - GenerateResponseChunk( + ModelResponseChunk( role=Role.MODEL, content=[reasoning_part], ) @@ -368,7 +363,7 @@ async def _generate_stream( ) for tool_call in delta.tool_calls ] - callback(GenerateResponseChunk(role=Role.MODEL, content=content)) + callback(ModelResponseChunk(role=Role.MODEL, content=content)) if tool_calls: message = MessageConverter.to_genkit( @@ -376,13 +371,13 @@ async def _generate_stream( ) accumulated_content.extend(message.content) - result = GenerateResponse( + result = ModelResponse( request=request, message=Message(role=Role.MODEL, content=accumulated_content), ) return self._clean_json_response(result, request) - async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: """Processes the request using OpenAI's chat completion API. Args: @@ -390,7 +385,7 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> Gen ctx: The context of the action run. Returns: - A GenerateResponse containing the model's response. + A ModelResponse containing the model's response. """ request.config = self.normalize_config(request.config) @@ -406,7 +401,7 @@ def normalize_config(config: object) -> OpenAIConfig: if isinstance(config, OpenAIConfig): return config - if isinstance(config, (GenerationCommonConfig, CoreGenerationCommonConfig)): + if isinstance(config, (ModelConfig, ModelConfig)): return OpenAIConfig( temperature=config.temperature, max_tokens=int(config.max_output_tokens) if config.max_output_tokens is not None else None, @@ -416,11 +411,8 @@ def normalize_config(config: object) -> OpenAIConfig: if isinstance(config, dict): config_dict = cast(dict[str, Any], config) - if config_dict.get('topK'): - del config_dict['topK'] - if config_dict.get('topP'): - config_dict['top_p'] = config_dict['topP'] - del config_dict['topP'] + if config_dict.get('top_k'): + del config_dict['top_k'] return OpenAIConfig(**config_dict) raise ValueError(f'Expected request.config to be a dict or OpenAIConfig, got {type(config).__name__}.') diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py index 1fa9221d51..448aaf73f7 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py @@ -24,11 +24,11 @@ else: from enum import StrEnum -from genkit.plugins.compat_oai.typing import SupportedOutputFormat -from genkit.types import ( +from genkit import ( ModelInfo, Supports, ) +from genkit.plugins.compat_oai.typing import SupportedOutputFormat OPENAI = 'openai' MODEL_GARDEN = 'model-garden' diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/utils.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/utils.py index 4cf937151b..fb3bf327c0 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/utils.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/utils.py @@ -23,10 +23,10 @@ from collections.abc import Callable from typing import Any -from genkit.types import ( - GenerateRequest, +from genkit import ( MediaPart, Message, + ModelRequest, Part, ReasoningPart, Role, @@ -58,7 +58,7 @@ def strip_markdown_fences(text: str) -> str: return text -def _find_text(request: GenerateRequest) -> str | None: +def _find_text(request: ModelRequest) -> str | None: """Find the first text content from the first message, if any. Args: @@ -76,7 +76,7 @@ def _find_text(request: GenerateRequest) -> str | None: ) -def _extract_text(request: GenerateRequest) -> str: +def _extract_text(request: ModelRequest) -> str: """Extract text content from the first message. Args: @@ -160,8 +160,8 @@ def decode_data_uri_bytes(url: str) -> bytes: raise ValueError('Invalid base64 data provided in media URL') from e -def extract_config_dict(request: GenerateRequest) -> dict[str, Any]: - """Extract the config from a GenerateRequest as a mutable dictionary. +def extract_config_dict(request: ModelRequest) -> dict[str, Any]: + """Extract the config from a ModelRequest as a mutable dictionary. Handles both dict configs and Pydantic model configs uniformly. @@ -180,7 +180,7 @@ def extract_config_dict(request: GenerateRequest) -> dict[str, Any]: return {} -def _extract_media(request: GenerateRequest) -> tuple[str, str]: +def _extract_media(request: ModelRequest) -> tuple[str, str]: """Extract media content from the first message. Finds the first part with a MediaPart root and returns its URL and diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/openai_plugin.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/openai_plugin.py index 55997957ca..fe3b3121dd 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/openai_plugin.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/openai_plugin.py @@ -18,19 +18,23 @@ """OpenAI OpenAI API Compatible Plugin for Genkit.""" import enum -from typing import Any, TypeAlias +from typing import Any, Literal, TypeAlias, cast from openai import AsyncOpenAI from openai.types import Model -from genkit.ai import ActionRunContext, Plugin -from genkit.blocks.embedding import EmbedderOptions, EmbedderSupports, embedder_action_metadata -from genkit.blocks.model import model_action_metadata -from genkit.core._loop_local import _loop_local_client -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.schema import to_json_schema -from genkit.core.typing import GenerationCommonConfig +from genkit import Embedding, EmbedRequest, EmbedResponse, ModelInfo, ModelRequest, ModelResponse, Supports +from genkit.embedder import EmbedderOptions, EmbedderSupports, embedder_action_metadata +from genkit.model import ModelConfig, model_action_metadata +from genkit.plugin_api import ( + Action, + ActionKind, + ActionMetadata, + ActionRunContext, + Plugin, + loop_local_client, + to_json_schema, +) from genkit.plugins.compat_oai.models import ( SUPPORTED_EMBEDDING_MODELS, SUPPORTED_IMAGE_MODELS, @@ -46,7 +50,6 @@ ) from genkit.plugins.compat_oai.models.model_info import get_default_openai_model_info from genkit.plugins.compat_oai.typing import OpenAIConfig -from genkit.types import Embedding, EmbedRequest, EmbedResponse, GenerateRequest, GenerateResponse, ModelInfo, Supports def open_ai_name(name: str) -> str: @@ -177,7 +180,7 @@ def _multimodal_action_metadata( """ return model_action_metadata( name=open_ai_name(name), - config_schema=GenerationCommonConfig, + config_schema=ModelConfig, info=_get_multimodal_info_dict(name, model_type, supported_models), ) @@ -206,7 +209,7 @@ def __init__(self, **openai_params: Any) -> None: # noqa: ANN401 other configuration settings required by OpenAI's API. """ self._openai_params = openai_params - self._runtime_client = _loop_local_client(lambda: AsyncOpenAI(**self._openai_params)) + self._runtime_client = loop_local_client(lambda: AsyncOpenAI(**self._openai_params)) self._list_actions_cache: list[ActionMetadata] | None = None async def init(self) -> list[Action]: @@ -313,7 +316,7 @@ def _create_model_action(self, name: str) -> Action: # Create the model handler model_info = self.get_model_info(clean_name) or {} - async def _generate(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def _generate(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: openai_model = OpenAIModelHandler(OpenAIModel(clean_name, self._runtime_client())) return await openai_model.generate(request, ctx) @@ -350,7 +353,7 @@ def _create_multimodal_action( clean_name = name.replace('openai/', '') if name.startswith('openai/') else name info_dict = _get_multimodal_info_dict(clean_name, model_type, supported_models) - async def _generate(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def _generate(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: model_instance = model_class(clean_name, self._runtime_client()) return await model_instance.generate(request, ctx) @@ -393,22 +396,41 @@ async def embed_fn(request: EmbedRequest) -> EmbedResponse: ) texts.append(doc_text) - # Get optional parameters with proper types - dimensions = None - encoding_format = None + # Get optional parameters (omit when None; OpenAI create() uses Omit, not None) + dimensions: int | None = None + encoding_format: Literal['base64', 'float'] | None = None if request.options: if dim_val := request.options.get('dimensions'): dimensions = int(dim_val) - if enc_val := request.options.get('encodingFormat'): - encoding_format = str(enc_val) if enc_val in ('float', 'base64') else None - - # Create embeddings for each document - response = await self._runtime_client().embeddings.create( - model=clean_name, - input=texts, - dimensions=dimensions, # type: ignore[arg-type] - encoding_format=encoding_format, # type: ignore[arg-type] - ) + enc_val = request.options.get('encodingFormat') + if enc_val in ('float', 'base64'): + encoding_format = cast(Literal['base64', 'float'], enc_val) + + # Call with only non-None optional params to satisfy strict typings + if dimensions is not None and encoding_format is not None: + response = await self._runtime_client().embeddings.create( + model=clean_name, + input=texts, + dimensions=dimensions, + encoding_format=encoding_format, + ) + elif dimensions is not None: + response = await self._runtime_client().embeddings.create( + model=clean_name, + input=texts, + dimensions=dimensions, + ) + elif encoding_format is not None: + response = await self._runtime_client().embeddings.create( + model=clean_name, + input=texts, + encoding_format=encoding_format, + ) + else: + response = await self._runtime_client().embeddings.create( + model=clean_name, + input=texts, + ) # Convert OpenAI response to Genkit format embeddings = [Embedding(embedding=item.embedding) for item in response.data] @@ -463,7 +485,7 @@ async def list_actions(self) -> list[ActionMetadata]: actions.append( model_action_metadata( name=open_ai_name(name), - config_schema=GenerationCommonConfig, + config_schema=ModelConfig, info={ 'label': f'OpenAI - {name}', 'supports': Supports( diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/typing.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/typing.py index f66532b498..b25af367a3 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/typing.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/typing.py @@ -38,7 +38,9 @@ else: from enum import StrEnum -from pydantic import BaseModel, ConfigDict, Field +from pydantic import ConfigDict, Field + +from genkit.model import ModelConfig class ReasoningEffort(StrEnum): @@ -112,7 +114,7 @@ class WebSearchContextSize(StrEnum): HIGH = 'high' -class OpenAIConfig(BaseModel): +class OpenAIConfig(ModelConfig): """OpenAI configuration for Genkit. This schema provides full control over OpenAI Chat Completions API parameters. diff --git a/py/plugins/compat-oai/tests/audio_model_test.py b/py/plugins/compat-oai/tests/audio_model_test.py index 9effb9b8ec..031f8c6891 100644 --- a/py/plugins/compat-oai/tests/audio_model_test.py +++ b/py/plugins/compat-oai/tests/audio_model_test.py @@ -23,6 +23,15 @@ import pytest +from genkit import ( + Media, + MediaPart, + Message, + ModelRequest, + Part, + Role, + TextPart, +) from genkit.plugins.compat_oai.models.audio import ( SUPPORTED_STT_MODELS, SUPPORTED_TTS_MODELS, @@ -35,23 +44,14 @@ _to_tts_params, _to_tts_response, ) -from genkit.types import ( - GenerateRequest, - Media, - MediaPart, - Message, - Part, - Role, - TextPart, -) class TestExtractText: - """Tests for extracting text from GenerateRequest.""" + """Tests for extracting text from ModelRequest.""" def test_extracts_text(self) -> None: """Verify text extraction from a simple request.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))]), ], @@ -61,17 +61,17 @@ def test_extracts_text(self) -> None: def test_raises_on_empty(self) -> None: """Verify ValueError when messages list is empty.""" - request = GenerateRequest(messages=[]) + request = ModelRequest(messages=[]) with pytest.raises(ValueError, match='No messages found'): _extract_text(request) class TestExtractMedia: - """Tests for extracting media URLs from GenerateRequest.""" + """Tests for extracting media URLs from ModelRequest.""" def test_extracts_media_url(self) -> None: """Verify media URL and content type extraction.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -94,7 +94,7 @@ def test_extracts_media_url(self) -> None: def test_raises_on_no_media(self) -> None: """Verify ValueError when no media content is found.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='no media'))]), ], @@ -104,11 +104,11 @@ def test_raises_on_no_media(self) -> None: class TestToTTSParams: - """Tests for converting GenerateRequest to TTS params.""" + """Tests for converting ModelRequest to TTS params.""" def test_basic_params(self) -> None: """Verify required TTS params with defaults.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='Say hello'))]), ], @@ -120,7 +120,7 @@ def test_basic_params(self) -> None: def test_custom_voice(self) -> None: """Verify custom voice config is applied.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='test'))]), ], @@ -131,19 +131,19 @@ def test_custom_voice(self) -> None: def test_strips_standard_config(self) -> None: """Verify standard GenAI keys are stripped.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='test'))]), ], - config={'temperature': 0.5, 'topK': 40}, + config={'temperature': 0.5, 'top_k': 40}, ) got = _to_tts_params('tts-1', request) assert 'temperature' not in got - assert 'topK' not in got + assert 'top_k' not in got class TestToTTSResponse: - """Tests for converting speech response to GenerateResponse.""" + """Tests for converting speech response to ModelResponse.""" def test_converts_audio_to_media_part(self) -> None: """Verify audio bytes are encoded as base64 data URI.""" @@ -172,12 +172,12 @@ def test_opus_format(self) -> None: class TestToSTTParams: - """Tests for converting GenerateRequest to STT params.""" + """Tests for converting ModelRequest to STT params.""" def test_basic_params(self) -> None: """Verify required STT params from audio media input.""" audio_data = base64.b64encode(b'fake audio').decode('ascii') - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -201,7 +201,7 @@ def test_basic_params(self) -> None: def test_with_prompt_context(self) -> None: """Verify prompt text is included when present alongside media.""" audio_data = base64.b64encode(b'fake audio').decode('ascii') - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -224,7 +224,7 @@ def test_with_prompt_context(self) -> None: class TestToSTTResponse: - """Tests for converting transcription result to GenerateResponse.""" + """Tests for converting transcription result to ModelResponse.""" def test_transcription_object(self) -> None: """Verify Transcription object is converted to text part.""" @@ -285,7 +285,7 @@ async def test_generate_calls_speech_create(self) -> None: mock_client.audio.speech.create = AsyncMock(return_value=mock_response) model = OpenAITTSModel('tts-1', mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='Say hello'))]), ], @@ -316,7 +316,7 @@ async def test_generate_calls_transcription_create(self) -> None: model = OpenAISTTModel('whisper-1', mock_client) audio_data = base64.b64encode(b'fake audio').decode('ascii') - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, diff --git a/py/plugins/compat-oai/tests/compat_oai_model_test.py b/py/plugins/compat-oai/tests/compat_oai_model_test.py index 1fef2890ad..479c115de0 100644 --- a/py/plugins/compat-oai/tests/compat_oai_model_test.py +++ b/py/plugins/compat-oai/tests/compat_oai_model_test.py @@ -17,31 +17,30 @@ """Tests for OpenAI compatible model implementation.""" -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, PropertyMock import pytest -from genkit.core.action._action import ActionRunContext -from genkit.plugins.compat_oai.models import OpenAIModel -from genkit.plugins.compat_oai.models.utils import strip_markdown_fences -from genkit.plugins.compat_oai.typing import OpenAIConfig -from genkit.types import ( - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationCommonConfig, +from genkit import ( Message, - OutputConfig, + ModelConfig, + ModelRequest, + ModelResponse, + ModelResponseChunk, Part, Role, TextPart, ) +from genkit.plugin_api import ActionRunContext +from genkit.plugins.compat_oai.models import OpenAIModel +from genkit.plugins.compat_oai.models.utils import strip_markdown_fences +from genkit.plugins.compat_oai.typing import OpenAIConfig -def test_get_messages(sample_request: GenerateRequest) -> None: +def test_get_messages(sample_request: ModelRequest) -> None: """Test _get_messages method. - Ensures the method correctly converts GenerateRequest messages into OpenAI-compatible ChatMessage format. + Ensures the method correctly converts ModelRequest messages into OpenAI-compatible ChatMessage format. """ model = OpenAIModel(model='gpt-4', client=MagicMock()) messages = model._get_messages(sample_request.messages) @@ -54,7 +53,7 @@ def test_get_messages(sample_request: GenerateRequest) -> None: @pytest.mark.asyncio -async def test_get_openai_config(sample_request: GenerateRequest) -> None: +async def test_get_openai_config(sample_request: ModelRequest) -> None: """Test _get_openai_request_config method. Ensures the method correctly constructs the OpenAI API configuration dictionary. @@ -69,8 +68,8 @@ async def test_get_openai_config(sample_request: GenerateRequest) -> None: @pytest.mark.asyncio -async def test__generate(sample_request: GenerateRequest) -> None: - """Test generate method calls OpenAI API and returns GenerateResponse.""" +async def test__generate(sample_request: ModelRequest) -> None: + """Test generate method calls OpenAI API and returns ModelResponse.""" mock_message = MagicMock() mock_message.content = 'Hello, user!' mock_message.role = 'model' @@ -87,14 +86,14 @@ async def test__generate(sample_request: GenerateRequest) -> None: response = await model._generate(sample_request) mock_client.chat.completions.create.assert_called_once() - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None assert response.message.role == Role.MODEL assert response.message.content[0].root.text == 'Hello, user!' @pytest.mark.asyncio -async def test__generate_stream(sample_request: GenerateRequest) -> None: +async def test__generate_stream(sample_request: ModelRequest) -> None: """Test generate_stream method ensures it processes streamed responses correctly.""" mock_client = MagicMock() @@ -129,7 +128,7 @@ async def __anext__(self) -> object: model = OpenAIModel(model='gpt-4', client=mock_client) collected_chunks = [] - def callback(chunk: GenerateResponseChunk) -> None: + def callback(chunk: ModelResponseChunk) -> None: collected_chunks.append(chunk.content[0].root.text) await model._generate_stream(sample_request, callback) @@ -145,24 +144,24 @@ def callback(chunk: GenerateResponseChunk) -> None: ], ) @pytest.mark.asyncio -async def test_generate(stream: bool, sample_request: GenerateRequest) -> None: +async def test_generate(stream: bool, sample_request: ModelRequest) -> None: """Tests for generate.""" ctx_mock = MagicMock(spec=ActionRunContext) - ctx_mock.is_streaming = stream + type(ctx_mock).is_streaming = PropertyMock(return_value=stream) - mock_response = GenerateResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='mocked'))])) + mock_response = ModelResponse(message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='mocked'))])) model = OpenAIModel(model='gpt-4', client=MagicMock()) - model._generate_stream = AsyncMock(return_value=mock_response) # type: ignore[method-assign] - model._generate = AsyncMock(return_value=mock_response) # type: ignore[method-assign] - model.normalize_config = MagicMock(return_value={}) # type: ignore[method-assign] + model._generate_stream = AsyncMock(return_value=mock_response) + model._generate = AsyncMock(return_value=mock_response) + model.normalize_config = MagicMock(return_value={}) response = await model.generate(sample_request, ctx_mock) assert response == mock_response if stream: - model._generate_stream.assert_called_once() # type: ignore[union-attr] + model._generate_stream.assert_called_once() else: - model._generate.assert_called_once() # type: ignore[union-attr] + model._generate.assert_called_once() @pytest.mark.parametrize( @@ -171,7 +170,7 @@ async def test_generate(stream: bool, sample_request: GenerateRequest) -> None: (OpenAIConfig(model='test'), OpenAIConfig(model='test')), ({'model': 'test'}, OpenAIConfig(model='test')), ( - GenerationCommonConfig(temperature=0.7), + ModelConfig(temperature=0.7), OpenAIConfig(temperature=0.7), ), ( @@ -207,32 +206,32 @@ class TestNeedsSchemaInPrompt: def test_true_for_deepseek_with_json_and_schema(self) -> None: """Returns True for DeepSeek model with json format and schema.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - output = OutputConfig(format='json', schema=_SAMPLE_SCHEMA) - assert model._needs_schema_in_prompt(output) is True + request = ModelRequest(messages=[], output_format='json', output_schema=_SAMPLE_SCHEMA) + assert model._needs_schema_in_prompt(request) is True def test_false_for_gpt_with_json_and_schema(self) -> None: """Returns False for GPT models even with json format and schema.""" model = OpenAIModel(model='gpt-4o', client=MagicMock()) - output = OutputConfig(format='json', schema=_SAMPLE_SCHEMA) - assert model._needs_schema_in_prompt(output) is False + request = ModelRequest(messages=[], output_format='json', output_schema=_SAMPLE_SCHEMA) + assert model._needs_schema_in_prompt(request) is False def test_false_for_deepseek_without_schema(self) -> None: """Returns False for DeepSeek when no schema is provided.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - output = OutputConfig(format='json') - assert model._needs_schema_in_prompt(output) is False + request = ModelRequest(messages=[], output_format='json') + assert model._needs_schema_in_prompt(request) is False def test_false_for_deepseek_with_text_format(self) -> None: """Returns False for DeepSeek when format is text.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - output = OutputConfig(format='text') - assert model._needs_schema_in_prompt(output) is False + request = ModelRequest(messages=[], output_format='text') + assert model._needs_schema_in_prompt(request) is False def test_false_for_no_format(self) -> None: """Returns False when output has no format set.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - output = OutputConfig() - assert model._needs_schema_in_prompt(output) is False + request = ModelRequest(messages=[]) + assert model._needs_schema_in_prompt(request) is False class TestBuildSchemaInstruction: @@ -264,11 +263,12 @@ class TestSchemaInjectionInConfig: async def test_deepseek_injects_schema_message(self) -> None: """DeepSeek request prepends a schema system message.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a character'))]), ], - output=OutputConfig(format='json', schema=_SAMPLE_SCHEMA), + output_format='json', + output_schema=_SAMPLE_SCHEMA, ) config = await model._get_openai_request_config(request) @@ -284,11 +284,12 @@ async def test_deepseek_injects_schema_message(self) -> None: async def test_gpt_does_not_inject_schema_message(self) -> None: """GPT request does not prepend a schema system message.""" model = OpenAIModel(model='gpt-4o', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='Generate a character'))]), ], - output=OutputConfig(format='json', schema=_SAMPLE_SCHEMA), + output_format='json', + output_schema=_SAMPLE_SCHEMA, ) config = await model._get_openai_request_config(request) @@ -301,11 +302,11 @@ async def test_gpt_does_not_inject_schema_message(self) -> None: async def test_deepseek_without_schema_no_injection(self) -> None: """DeepSeek request without a schema does not inject anything.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))]), ], - output=OutputConfig(format='json'), + output_format='json', ) config = await model._get_openai_request_config(request) @@ -317,12 +318,13 @@ async def test_deepseek_without_schema_no_injection(self) -> None: async def test_deepseek_preserves_existing_system_message(self) -> None: """Schema injection does not clobber an existing system message.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='You are helpful'))]), Message(role=Role.USER, content=[Part(root=TextPart(text='Generate'))]), ], - output=OutputConfig(format='json', schema=_SAMPLE_SCHEMA), + output_format='json', + output_schema=_SAMPLE_SCHEMA, ) config = await model._get_openai_request_config(request) @@ -377,11 +379,12 @@ class TestCleanJsonResponse: def test_cleans_deepseek_json_response(self) -> None: """Strips markdown fences from DeepSeek JSON response.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='json', schema=_SAMPLE_SCHEMA), + output_format='json', + output_schema=_SAMPLE_SCHEMA, ) - response = GenerateResponse( + response = ModelResponse( request=request, message=Message( role=Role.MODEL, @@ -395,12 +398,13 @@ def test_cleans_deepseek_json_response(self) -> None: def test_no_op_for_gpt_model(self) -> None: """Does not modify responses from non-DeepSeek models.""" model = OpenAIModel(model='gpt-4o', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='json', schema=_SAMPLE_SCHEMA), + output_format='json', + output_schema=_SAMPLE_SCHEMA, ) fenced_text = '```json\n{"name": "John", "level": 5}\n```' - response = GenerateResponse( + response = ModelResponse( request=request, message=Message( role=Role.MODEL, @@ -414,12 +418,12 @@ def test_no_op_for_gpt_model(self) -> None: def test_no_op_for_text_output(self) -> None: """Does not modify responses when output format is not json.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='text'), + output_format='text', ) text = '```json\n{"a": 1}\n```' - response = GenerateResponse( + response = ModelResponse( request=request, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=text))]), ) @@ -430,11 +434,11 @@ def test_no_op_for_text_output(self) -> None: def test_no_op_for_no_output(self) -> None: """Does not modify responses when no output config is set.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], ) text = '```json\n{"a": 1}\n```' - response = GenerateResponse( + response = ModelResponse( request=request, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=text))]), ) @@ -445,12 +449,13 @@ def test_no_op_for_no_output(self) -> None: def test_no_op_when_no_fences(self) -> None: """Does not modify clean JSON responses.""" model = OpenAIModel(model='deepseek-chat', client=MagicMock()) - request = GenerateRequest( + request = ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hi'))])], - output=OutputConfig(format='json', schema=_SAMPLE_SCHEMA), + output_format='json', + output_schema=_SAMPLE_SCHEMA, ) text = '{"name": "John", "level": 5}' - response = GenerateResponse( + response = ModelResponse( request=request, message=Message(role=Role.MODEL, content=[Part(root=TextPart(text=text))]), ) diff --git a/py/plugins/compat-oai/tests/compat_oai_plugin_test.py b/py/plugins/compat-oai/tests/compat_oai_plugin_test.py index f44a42bd2e..c401e27514 100644 --- a/py/plugins/compat-oai/tests/compat_oai_plugin_test.py +++ b/py/plugins/compat-oai/tests/compat_oai_plugin_test.py @@ -25,9 +25,7 @@ import pytest from openai.types import Model -from genkit.core._loop_local import _loop_local_client -from genkit.core.action import ActionMetadata -from genkit.core.action.types import ActionKind +from genkit.plugin_api import ActionKind, ActionMetadata, loop_local_client from genkit.plugins.compat_oai.openai_plugin import OpenAI, openai_model @@ -103,7 +101,7 @@ async def test_openai_plugin_list_actions() -> None: async def test_openai_runtime_clients_are_loop_local() -> None: """Runtime OpenAI clients are cached per event loop.""" plugin = OpenAI(api_key='test-key') - plugin._runtime_client = _loop_local_client(lambda: object()) + plugin._runtime_client = loop_local_client(lambda: object()) first = plugin._runtime_client() second = plugin._runtime_client() diff --git a/py/plugins/compat-oai/tests/compat_oai_utils_test.py b/py/plugins/compat-oai/tests/compat_oai_utils_test.py index ec06bfb378..dd252c1026 100644 --- a/py/plugins/compat-oai/tests/compat_oai_utils_test.py +++ b/py/plugins/compat-oai/tests/compat_oai_utils_test.py @@ -21,22 +21,11 @@ import pytest from pydantic import BaseModel -from genkit.plugins.compat_oai.models.utils import ( - DictMessageAdapter, - MessageAdapter, - MessageConverter, - _extract_media, - _extract_text, - _find_text, - decode_data_uri_bytes, - extract_config_dict, - parse_data_uri_content_type, -) -from genkit.types import ( - GenerateRequest, +from genkit import ( Media, MediaPart, Message, + ModelRequest, Part, ReasoningPart, Role, @@ -46,6 +35,17 @@ ToolResponse, ToolResponsePart, ) +from genkit.plugins.compat_oai.models.utils import ( + DictMessageAdapter, + MessageAdapter, + MessageConverter, + _extract_media, + _extract_text, + _find_text, + decode_data_uri_bytes, + extract_config_dict, + parse_data_uri_content_type, +) class TestParseDataUriContentType: @@ -162,9 +162,9 @@ def test_data_uri_with_no_content_type(self) -> None: class TestExtractConfigDict: """Tests for extract_config_dict.""" - def _make_request(self, config: object = None) -> GenerateRequest: - """Create a minimal GenerateRequest with the given config.""" - return GenerateRequest( + def _make_request(self, config: object = None) -> ModelRequest: + """Create a minimal ModelRequest with the given config.""" + return ModelRequest( messages=[ Message( role=Role.USER, @@ -206,7 +206,7 @@ class TestFindText: def test_returns_text_from_first_message(self) -> None: """Find and return text from the first message's text part.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -218,12 +218,12 @@ def test_returns_text_from_first_message(self) -> None: def test_returns_none_for_no_messages(self) -> None: """Return None when there are no messages.""" - request = GenerateRequest(messages=[]) + request = ModelRequest(messages=[]) assert _find_text(request) is None def test_returns_none_for_no_text_parts(self) -> None: """Return None when message has only media parts.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -235,7 +235,7 @@ def test_returns_none_for_no_text_parts(self) -> None: def test_returns_first_text_part(self) -> None: """Return the first text part when multiple exist.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -254,7 +254,7 @@ class TestExtractText: def test_returns_text_when_present(self) -> None: """Extract and return text when present.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -266,13 +266,13 @@ def test_returns_text_when_present(self) -> None: def test_raises_for_no_messages(self) -> None: """Raise ValueError when request has no messages.""" - request = GenerateRequest(messages=[]) + request = ModelRequest(messages=[]) with pytest.raises(ValueError, match='No messages found'): _extract_text(request) def test_raises_for_no_text_content(self) -> None: """Raise ValueError when no text parts exist.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -289,7 +289,7 @@ class TestExtractMedia: def test_extracts_media_url_and_content_type(self) -> None: """Extract URL and content type from a media part.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -312,7 +312,7 @@ def test_extracts_media_url_and_content_type(self) -> None: def test_parses_content_type_from_data_uri_when_missing(self) -> None: """Parse content type from data URI when not explicitly provided.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -334,13 +334,13 @@ def test_parses_content_type_from_data_uri_when_missing(self) -> None: def test_raises_for_no_messages(self) -> None: """Raise ValueError when request has no messages.""" - request = GenerateRequest(messages=[]) + request = ModelRequest(messages=[]) with pytest.raises(ValueError, match='No messages found'): _extract_media(request) def test_raises_for_no_media_parts(self) -> None: """Raise ValueError when message has no media parts.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -353,7 +353,7 @@ def test_raises_for_no_media_parts(self) -> None: def test_skips_text_parts_finds_media(self) -> None: """Find media part even when text parts come first.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -376,7 +376,7 @@ def test_skips_text_parts_finds_media(self) -> None: def test_content_type_from_data_uri_without_base64_qualifier(self) -> None: """Parse content type from data URI that omits ;base64 qualifier.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, diff --git a/py/plugins/compat-oai/tests/conftest.py b/py/plugins/compat-oai/tests/conftest.py index 5f3540982e..2fb0ae29e3 100644 --- a/py/plugins/compat-oai/tests/conftest.py +++ b/py/plugins/compat-oai/tests/conftest.py @@ -19,20 +19,20 @@ import pytest -from genkit.plugins.compat_oai.typing import OpenAIConfig -from genkit.types import ( - GenerateRequest, +from genkit import ( Message, + ModelRequest, Part, Role, TextPart, ) +from genkit.plugins.compat_oai.typing import OpenAIConfig @pytest.fixture -def sample_request() -> GenerateRequest: - """Fixture to create a sample GenerateRequest object.""" - return GenerateRequest( +def sample_request() -> ModelRequest: + """Fixture to create a sample ModelRequest object.""" + return ModelRequest( messages=[ Message( role=Role.SYSTEM, diff --git a/py/plugins/compat-oai/tests/image_model_test.py b/py/plugins/compat-oai/tests/image_model_test.py index 3d9649d7b8..ecd75b357a 100644 --- a/py/plugins/compat-oai/tests/image_model_test.py +++ b/py/plugins/compat-oai/tests/image_model_test.py @@ -22,6 +22,14 @@ import pytest +from genkit import ( + MediaPart, + Message, + ModelRequest, + Part, + Role, + TextPart, +) from genkit.plugins.compat_oai.models.image import ( SUPPORTED_IMAGE_MODELS, OpenAIImageModel, @@ -29,22 +37,14 @@ _to_generate_response, _to_image_generate_params, ) -from genkit.types import ( - GenerateRequest, - MediaPart, - Message, - Part, - Role, - TextPart, -) class TestExtractPromptText: - """Tests for extracting text from GenerateRequest messages.""" + """Tests for extracting text from ModelRequest messages.""" def test_extracts_text_from_first_message(self) -> None: """Verify text extraction from a simple single-message request.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='a sunset'))]), ], @@ -54,13 +54,13 @@ def test_extracts_text_from_first_message(self) -> None: def test_raises_on_empty_messages(self) -> None: """Verify ValueError when messages list is empty.""" - request = GenerateRequest(messages=[]) + request = ModelRequest(messages=[]) with pytest.raises(ValueError, match='No messages found'): _extract_prompt_text(request) def test_raises_on_no_text_content(self) -> None: """Verify ValueError when message has no text parts.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[]), ], @@ -70,11 +70,11 @@ def test_raises_on_no_text_content(self) -> None: class TestToImageGenerateParams: - """Tests for converting GenerateRequest to OpenAI image params.""" + """Tests for converting ModelRequest to OpenAI image params.""" def test_basic_params(self) -> None: """Verify required params are set with correct defaults.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='a cat'))]), ], @@ -86,7 +86,7 @@ def test_basic_params(self) -> None: def test_config_passthrough(self) -> None: """Verify image-specific config options pass through.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='a dog'))]), ], @@ -99,20 +99,20 @@ def test_config_passthrough(self) -> None: def test_strips_standard_genai_config(self) -> None: """Verify standard GenAI keys are stripped from params.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='test'))]), ], - config={'temperature': 0.5, 'topK': 40, 'topP': 0.9}, + config={'temperature': 0.5, 'top_k': 40, 'top_p': 0.9}, ) got = _to_image_generate_params('dall-e-3', request) assert 'temperature' not in got - assert 'topK' not in got - assert 'topP' not in got + assert 'top_k' not in got + assert 'top_p' not in got def test_version_override(self) -> None: """Verify model version override via config.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='test'))]), ], @@ -122,8 +122,8 @@ def test_version_override(self) -> None: assert got['model'] == 'dall-e-3-custom' -class TestToGenerateResponse: - """Tests for converting OpenAI ImagesResponse to GenerateResponse.""" +class TestToModelResponse: + """Tests for converting OpenAI ImagesResponse to ModelResponse.""" def test_empty_data(self) -> None: """Verify empty image data produces empty content.""" @@ -213,7 +213,7 @@ async def test_generate_calls_client(self) -> None: mock_client.images.generate = AsyncMock(return_value=mock_response) model = OpenAIImageModel('dall-e-3', mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.USER, content=[Part(root=TextPart(text='a mountain'))]), ], diff --git a/py/plugins/compat-oai/tests/tool_calling_test.py b/py/plugins/compat-oai/tests/tool_calling_test.py index 86b9b9f2e4..bdbafc951e 100644 --- a/py/plugins/compat-oai/tests/tool_calling_test.py +++ b/py/plugins/compat-oai/tests/tool_calling_test.py @@ -22,12 +22,12 @@ import pytest +from genkit import ModelRequest, ModelResponseChunk, TextPart, ToolRequestPart from genkit.plugins.compat_oai.models import OpenAIModel -from genkit.types import GenerateRequest, GenerateResponseChunk, TextPart, ToolRequestPart @pytest.mark.asyncio -async def test_generate_with_tool_calls_executes_tools(sample_request: GenerateRequest) -> None: +async def test_generate_with_tool_calls_executes_tools(sample_request: ModelRequest) -> None: """Test generate with tool calls executes tools.""" mock_tool_call = MagicMock() mock_tool_call.id = 'tool123' @@ -86,7 +86,7 @@ async def test_generate_with_tool_calls_executes_tools(sample_request: GenerateR @pytest.mark.asyncio -async def test_generate_stream_with_tool_calls(sample_request: GenerateRequest) -> None: +async def test_generate_stream_with_tool_calls(sample_request: ModelRequest) -> None: """Test generate_stream processes tool calls streamed in chunks correctly.""" mock_client = MagicMock() @@ -137,7 +137,7 @@ async def __anext__(self) -> object: model = OpenAIModel(model='gpt-4', client=mock_client) collected_chunks = [] - def callback(chunk: GenerateResponseChunk) -> None: + def callback(chunk: ModelResponseChunk) -> None: collected_chunks.append(chunk.content[0].root) await model._generate_stream(sample_request, callback) diff --git a/py/plugins/dev-local-vectorstore/LICENSE b/py/plugins/dev-local-vectorstore/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/plugins/dev-local-vectorstore/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/plugins/dev-local-vectorstore/README.md b/py/plugins/dev-local-vectorstore/README.md deleted file mode 100644 index 589294914d..0000000000 --- a/py/plugins/dev-local-vectorstore/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Dev Local Vector Store Plugin - -Local file-based vectorstore plugin that provides retriever and indexer helper for Genkit. diff --git a/py/plugins/dev-local-vectorstore/pyproject.toml b/py/plugins/dev-local-vectorstore/pyproject.toml deleted file mode 100644 index 735f076d4d..0000000000 --- a/py/plugins/dev-local-vectorstore/pyproject.toml +++ /dev/null @@ -1,83 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Framework :: AsyncIO", - "Framework :: Pydantic", - "Framework :: Pydantic :: 2", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Typing :: Typed", - "License :: OSI Approved :: Apache Software License", -] -dependencies = [ - "aiofiles>=24.1.0", - "genkit", - "pytest-mock", - "pytest-asyncio", - "strenum>=0.4.15; python_version < '3.11'", -] -description = "Genkit Local Vector Store Plugin" -keywords = [ - "genkit", - "ai", - "llm", - "machine-learning", - "artificial-intelligence", - "generative-ai", - "vector-store", - "embeddings", - "local", -] -license = "Apache-2.0" -name = "genkit-plugin-dev-local-vectorstore" -readme = "README.md" -requires-python = ">=3.10" -version = "0.5.1" - -[project.urls] -"Bug Tracker" = "https://github.com/firebase/genkit/issues" -Changelog = "https://github.com/firebase/genkit/blob/main/py/plugins/dev-local-vectorstore/CHANGELOG.md" -"Documentation" = "https://firebase.google.com/docs/genkit" -"Homepage" = "https://github.com/firebase/genkit" -"Repository" = "https://github.com/firebase/genkit/tree/main/py" - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -only-include = ["src/genkit/plugins/dev_local_vectorstore"] -sources = ["src"] diff --git a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/__init__.py b/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/__init__.py deleted file mode 100644 index 65a7aac817..0000000000 --- a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/__init__.py +++ /dev/null @@ -1,164 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Development local vector store plugin for Genkit. - -This plugin provides a simple, file-based vector store for local development -and testing. It's not intended for production use but is ideal for prototyping -RAG applications without setting up external infrastructure. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Vector Store β”‚ A special database for finding "similar" things. β”‚ - β”‚ β”‚ Like a librarian who knows all related books. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Local Dev Store β”‚ A simple vector store that runs on your computer. β”‚ - β”‚ β”‚ No cloud setup needed - just start coding! β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Embeddings β”‚ Numbers that represent the meaning of text. β”‚ - β”‚ β”‚ "Happy" and "joyful" get similar numbers. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Indexer β”‚ Stores documents with their embeddings. β”‚ - β”‚ β”‚ Like a librarian cataloging new books. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Retriever β”‚ Finds documents matching a query. β”‚ - β”‚ β”‚ Like asking "show me docs about X". β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Cosine Similarity β”‚ Math to compare how similar two embeddings are. β”‚ - β”‚ β”‚ 1.0 = identical, 0 = totally different. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ RAG β”‚ Retrieval-Augmented Generation. Find relevant β”‚ - β”‚ β”‚ docs first, then let AI answer using them. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW LOCAL VECTOR SEARCH WORKS β”‚ - β”‚ β”‚ - β”‚ STEP 1: INDEX YOUR DOCUMENTS β”‚ - β”‚ ───────────────────────────── β”‚ - β”‚ Your Documents: ["How to bake cookies", "Cookie recipes", ...] β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) Each doc gets converted to numbers β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ "bake cookies" β†’ [0.2, -0.5, 0.8, ...] β”‚ - β”‚ β”‚ (any model) β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Stored in local file/memory β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Local Store β”‚ No database needed! β”‚ - β”‚ β”‚ (JSON file) β”‚ Just a file on your computer. β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ - β”‚ STEP 2: SEARCH YOUR DOCUMENTS β”‚ - β”‚ ───────────────────────────── β”‚ - β”‚ Query: "How do I make chocolate chip cookies?" β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Query converted to embedding β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ Query β†’ [0.21, -0.48, 0.79, ...] (similar!) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (4) Find nearest neighbors (cosine similarity) β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Similarity β”‚ "Cookie recipes" scores 0.95 β”‚ - β”‚ β”‚ Search β”‚ "How to bake cookies" scores 0.92 β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (5) Return top matches β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Your App β”‚ Now AI can answer using these docs! β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Architecture Overview:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Dev Local Vector Store Plugin β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Plugin Entry Point (__init__.py) β”‚ - β”‚ └── define_dev_local_vector_store() - Create local vector store β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ plugin_api.py - Plugin API β”‚ - β”‚ β”œβ”€β”€ define_dev_local_vector_store() - Main factory function β”‚ - β”‚ └── Returns indexer and retriever pair β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ local_vector_store_api.py - Vector Store Implementation β”‚ - β”‚ β”œβ”€β”€ In-memory vector storage β”‚ - β”‚ └── Cosine similarity search β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ indexer.py - Document Indexing β”‚ - β”‚ └── Add documents to the local store β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ retriever.py - Document Retrieval β”‚ - β”‚ └── Retrieve similar documents by query β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Data Flow β”‚ - β”‚ β”‚ - β”‚ Documents ──► Embedder ──► Indexer ──► Local Store (file/memory) β”‚ - β”‚ β”‚ - β”‚ Query ──► Embedder ──► Retriever ──► Similarity Search ──► Results β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - ```python - from genkit import Genkit - from genkit.plugins.dev_local_vectorstore import define_dev_local_vector_store - - ai = Genkit(...) - - # Create a local vector store - store = define_dev_local_vector_store( - ai, - name='my_store', - embedder='googleai/gemini-embedding-001', - ) - - # Index documents - await ai.index(indexer=store.indexer, documents=[...]) - - # Retrieve similar documents - results = await ai.retrieve(retriever=store.retriever, query='...') - ``` - -Caveats: - - NOT for production use (no persistence guarantees) - - Data is stored in memory or local files - - No concurrent access support - - Use Firebase, Vertex AI, or other production stores for real applications - -See Also: - - Genkit documentation: https://genkit.dev/ -""" - -from .plugin_api import define_dev_local_vector_store - -__all__ = [ - define_dev_local_vector_store.__name__, -] diff --git a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/indexer.py b/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/indexer.py deleted file mode 100644 index edc073e825..0000000000 --- a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/indexer.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Indexer for dev-local-vectorstore.""" - -import asyncio -import hashlib -import json -from typing import Any - -from genkit.ai import Genkit -from genkit.blocks.document import Document -from genkit.blocks.retriever import IndexerRequest -from genkit.types import Embedding - -from .constant import DbValue -from .local_vector_store_api import LocalVectorStoreAPI - - -class DevLocalVectorStoreIndexer(LocalVectorStoreAPI): - """Indexer for development-level local vector store.""" - - def __init__( - self, - ai: Genkit, - index_name: str, - embedder: str, - embedder_options: dict[str, Any] | None = None, - ) -> None: - """Initialize the DevLocalVectorStoreIndexer. - - Args: - ai: Genkit instance used to embed documents. - index_name: Name of the index. - embedder: The embedder to use for document embeddings. - embedder_options: Optional configuration to pass to the embedder. - """ - super().__init__(index_name=index_name) - self.ai = ai - self.embedder = embedder - self.embedder_options = embedder_options - - async def index(self, request: IndexerRequest) -> None: - """Index documents into the local vector store.""" - docs = request.documents - # pyrefly: ignore[missing-attribute] - inherited from LocalVectorStoreAPI - data = await self._load_filestore() - - embed_resp = await self.ai.embed_many( - embedder=self.embedder, - content=docs, - options=self.embedder_options, - ) - if not embed_resp: - raise ValueError('Embedder returned no embeddings for documents') - - tasks = [] - for doc_data, emb in zip(docs, embed_resp, strict=True): - tasks.append( - self.process_document( - document=Document.from_document_data(document_data=doc_data), - embedding=Embedding(embedding=emb.embedding), - data=data, - ) - ) - - await asyncio.gather(*tasks) - - # pyrefly: ignore[missing-attribute] - _dump_filestore inherited from LocalVectorStoreAPI - await self._dump_filestore(data) - - async def process_document(self, document: Document, embedding: Embedding, data: dict[str, DbValue]) -> None: - """Process a single document and add its embedding to the store.""" - embedding_docs = document.get_embedding_documents([embedding]) - self._add_document(data=data, embedding=embedding, doc=embedding_docs[0]) - - def _add_document( - self, - data: dict[str, DbValue], - embedding: Embedding, - doc: Document, - ) -> None: - # pyrefly: ignore[missing-attribute] - _serialize_data inherited from LocalVectorStoreAPI - data_str = json.dumps(self._serialize_data(data=data), ensure_ascii=False) - # MD5 used for content-based ID generation, not security - idx = hashlib.md5(data_str.encode('utf-8'), usedforsecurity=False).hexdigest() - if idx not in data: - data[idx] = DbValue( - doc=doc, - embedding=embedding, - ) diff --git a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/local_vector_store_api.py b/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/local_vector_store_api.py deleted file mode 100644 index 1d4b18b1e3..0000000000 --- a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/local_vector_store_api.py +++ /dev/null @@ -1,79 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Base API for dev-local-vectorstore. - -Provides async file-backed storage for document embeddings using -``aiofiles`` to avoid blocking the event loop during reads and writes. -""" - -import json -from functools import cached_property - -import aiofiles -import aiofiles.os - -from genkit.codec import dump_json - -from .constant import DbValue - - -class LocalVectorStoreAPI: - """Base class for development local vector store operations.""" - - _LOCAL_FILESTORE_TEMPLATE = '__db_{index_name}.json' - - def __init__( - self, - index_name: str, - ) -> None: - """Initialize the LocalVectorStoreAPI.""" - self.index_name = index_name - - @cached_property - def index_file_name(self) -> str: - """Get the filename of the index file.""" - return self._LOCAL_FILESTORE_TEMPLATE.format(index_name=self.index_name) - - async def _load_filestore(self) -> dict[str, DbValue]: - """Load the filestore asynchronously to avoid blocking the event loop.""" - data: dict[str, object] = {} - if await aiofiles.os.path.exists(self.index_file_name): - async with aiofiles.open(self.index_file_name, encoding='utf-8') as f: - contents = await f.read() - data = json.loads(contents) - return self._deserialize_data(data) - - async def _dump_filestore(self, data: dict[str, DbValue]) -> None: - """Dump the filestore asynchronously to avoid blocking the event loop.""" - serialized_data = self._serialize_data(data) - async with aiofiles.open(self.index_file_name, 'w', encoding='utf-8') as f: - await f.write(dump_json(serialized_data, indent=2)) - - @staticmethod - def _serialize_data(data: dict[str, DbValue]) -> dict[str, object]: - result: dict[str, object] = {} - for k, v in data.items(): - result[k] = DbValue.model_dump(v, exclude_none=True) - return result - - @staticmethod - def _deserialize_data(data: dict[str, object]) -> dict[str, DbValue]: - result: dict[str, DbValue] = {} - for k, v in data.items(): - result[k] = DbValue.model_validate(v) - return result diff --git a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/plugin_api.py b/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/plugin_api.py deleted file mode 100644 index f776d30307..0000000000 --- a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/plugin_api.py +++ /dev/null @@ -1,97 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Local file-based vectorstore helper that provides retriever and indexer for Genkit.""" - -from typing import Any - -from genkit.ai import Genkit -from genkit.blocks.retriever import ( - IndexerOptions, - RetrieverOptions, - indexer_action_metadata, - retriever_action_metadata, -) -from genkit.core.registry import ActionKind -from genkit.core.schema import to_json_schema - -from .indexer import DevLocalVectorStoreIndexer -from .retriever import DevLocalVectorStoreRetriever, RetrieverOptionsSchema - - -def define_dev_local_vector_store( - ai: Genkit, - *, - name: str, - embedder: str, - embedder_options: dict[str, Any] | None = None, -) -> tuple[str, str]: - """Define and register a dev local vector store retriever and indexer. - - NOT INTENDED FOR USE IN PRODUCTION - - Args: - ai: The Genkit instance to register the retriever and indexer with. - name: Name of the retriever and indexer. - embedder: The embedder to use (e.g., 'vertexai/gemini-embedding-001'). - embedder_options: Optional configuration to pass to the embedder. - - Returns: - Tuple of (retriever_name, indexer_name). - """ - # Create and register retriever - retriever = DevLocalVectorStoreRetriever( - ai=ai, - index_name=name, - embedder=embedder, - embedder_options=embedder_options, - ) - - ai.registry.register_action( - kind=ActionKind.RETRIEVER, - name=name, - fn=retriever.retrieve, - metadata=retriever_action_metadata( - name=name, - options=RetrieverOptions( - label=name, - config_schema=to_json_schema(RetrieverOptionsSchema), - ), - ).metadata, - ) - - # Create and register indexer - indexer = DevLocalVectorStoreIndexer( - ai=ai, - index_name=name, - embedder=embedder, - embedder_options=embedder_options, - ) - - ai.registry.register_action( - kind=ActionKind.INDEXER, - name=name, - fn=indexer.index, - metadata=indexer_action_metadata( - name=name, - options=IndexerOptions(label=name), - ).metadata, - ) - - return (name, name) - - -define_dev_local_vector_store_deprecated = define_dev_local_vector_store diff --git a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/py.typed b/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/py.typed deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/retriever.py b/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/retriever.py deleted file mode 100644 index 372ad0c80a..0000000000 --- a/py/plugins/dev-local-vectorstore/src/genkit/plugins/dev_local_vectorstore/retriever.py +++ /dev/null @@ -1,115 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Retriever for dev-local-vectorstore.""" - -from typing import Any - -from pydantic import BaseModel, Field - -from genkit.ai import ActionRunContext, Document, Genkit -from genkit.types import Embedding, RetrieverRequest, RetrieverResponse - -from .local_vector_store_api import LocalVectorStoreAPI - - -class ScoredDocument(BaseModel): - """Document with an associated similarity score.""" - - score: float - document: Document - - -class RetrieverOptionsSchema(BaseModel): - """Schema for retriever options.""" - - limit: int | None = Field(title='Number of documents to retrieve', default=None) - - -class DevLocalVectorStoreRetriever(LocalVectorStoreAPI): - """Retriever for development-level local vector store.""" - - def __init__( - self, - ai: Genkit, - index_name: str, - embedder: str, - embedder_options: dict[str, Any] | None = None, - ) -> None: - """Initialize the DevLocalVectorStoreRetriever. - - Args: - ai: Genkit instance used to embed queries. - index_name: Name of the index. - embedder: The embedder to use for query embeddings. - embedder_options: Optional configuration to pass to the embedder. - """ - super().__init__(index_name=index_name) - self.ai = ai - self.embedder = embedder - self.embedder_options = embedder_options - - async def retrieve(self, request: RetrieverRequest, _: ActionRunContext) -> RetrieverResponse: - """Retrieve documents from the vector store.""" - document = Document.from_document_data(document_data=request.query) - - embed_resp = await self.ai.embed( - embedder=self.embedder, - content=document, - options=self.embedder_options, - ) - if not embed_resp: - raise ValueError('Embedder returned no embeddings for query') - - k = 3 - if isinstance(request.options, dict) and (limit_val := request.options.get('limit')) is not None: - k = int(limit_val) - - docs = await self._get_closest_documents( - k=k, - query_embeddings=Embedding(embedding=embed_resp[0].embedding), - ) - - return RetrieverResponse(documents=[d.document for d in docs]) - - async def _get_closest_documents(self, k: int, query_embeddings: Embedding) -> list[ScoredDocument]: - # pyrefly: ignore[missing-attribute] - _load_filestore inherited from LocalVectorStoreAPI - db = await self._load_filestore() - scored_documents = [] - - for val in db.values(): - this_embedding = val.embedding.embedding - score = self.cosine_similarity(query_embeddings.embedding, this_embedding) - scored_documents.append( - ScoredDocument( - score=score, - document=Document.from_document_data(document_data=val.doc), - ) - ) - - scored_documents = sorted(scored_documents, key=lambda d: d.score, reverse=True) - return scored_documents[:k] - - @classmethod - def cosine_similarity(cls, a: list[float], b: list[float]) -> float: - """Calculate cosine similarity between two vectors.""" - return cls.dot(a, b) / ((cls.dot(a, a) ** 0.5) * (cls.dot(b, b) ** 0.5)) - - @staticmethod - def dot(a: list[float], b: list[float]) -> float: - """Calculate dot product of two vectors.""" - return sum(av * bv for av, bv in zip(a, b, strict=False)) diff --git a/py/plugins/dev-local-vectorstore/tests/.gitkeep b/py/plugins/dev-local-vectorstore/tests/.gitkeep deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/plugins/dev-local-vectorstore/tests/dev_local_vectorstore_plugin_test.py b/py/plugins/dev-local-vectorstore/tests/dev_local_vectorstore_plugin_test.py deleted file mode 100644 index 671d6828ef..0000000000 --- a/py/plugins/dev-local-vectorstore/tests/dev_local_vectorstore_plugin_test.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Dev Local Vector Store plugin.""" - -from genkit.plugins.dev_local_vectorstore import define_dev_local_vector_store - - -def test_define_dev_local_vector_store_callable() -> None: - """Test define_dev_local_vector_store is callable.""" - assert callable(define_dev_local_vector_store) - - -def test_define_dev_local_vector_store_exported() -> None: - """Test define_dev_local_vector_store is exported from package.""" - from genkit.plugins.dev_local_vectorstore import define_dev_local_vector_store as func - - assert func is not None - assert callable(func) diff --git a/py/plugins/dev-local-vectorstore/tests/dlvs_api_test.py b/py/plugins/dev-local-vectorstore/tests/dlvs_api_test.py deleted file mode 100644 index 45f3122383..0000000000 --- a/py/plugins/dev-local-vectorstore/tests/dlvs_api_test.py +++ /dev/null @@ -1,181 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for LocalVectorStoreAPI base class.""" - -import json -import pathlib -import tempfile - -import aiofiles.os -import pytest - -from genkit.plugins.dev_local_vectorstore.constant import DbValue -from genkit.plugins.dev_local_vectorstore.local_vector_store_api import LocalVectorStoreAPI -from genkit.types import DocumentData, Embedding, TextPart - - -def _make_db_value(text: str = 'hello', embedding: list[float] | None = None) -> DbValue: - """Create a DbValue for testing.""" - return DbValue( - # Pydantic's discriminated union accepts TextPart directly at runtime, - # but the static type is list[Part]. Wrapping in Part(root=...) causes - # a ValidationError, so the type: ignore is intentional. - doc=DocumentData(content=[TextPart(text=text)]), # type: ignore[list-item] - embedding=Embedding(embedding=embedding or [0.1, 0.2, 0.3]), - ) - - -class TestIndexFileName: - """Tests for index file naming.""" - - def test_index_file_name_format(self) -> None: - """Index file name follows the expected template.""" - api = LocalVectorStoreAPI(index_name='test-index') - assert api.index_file_name == '__db_test-index.json' - - def test_index_file_name_with_special_chars(self) -> None: - """Index file name handles special characters in index name.""" - api = LocalVectorStoreAPI(index_name='my_index_123') - assert api.index_file_name == '__db_my_index_123.json' - - def test_index_file_name_cached(self) -> None: - """Index file name is cached (cached_property).""" - api = LocalVectorStoreAPI(index_name='cached') - name1 = api.index_file_name - name2 = api.index_file_name - assert name1 is name2 - - -class TestSerializeData: - """Tests for data serialization.""" - - def test_serialize_empty_data(self) -> None: - """Serialize empty dict returns empty dict.""" - result = LocalVectorStoreAPI._serialize_data({}) - assert result == {} - - def test_serialize_single_entry(self) -> None: - """Serialize single DbValue entry.""" - data = {'key1': _make_db_value('test doc')} - result = LocalVectorStoreAPI._serialize_data(data) - assert 'key1' in result - serialized = result['key1'] - assert isinstance(serialized, dict) - - def test_serialize_multiple_entries(self) -> None: - """Serialize multiple DbValue entries preserves all keys.""" - data = { - 'a': _make_db_value('first'), - 'b': _make_db_value('second'), - 'c': _make_db_value('third'), - } - result = LocalVectorStoreAPI._serialize_data(data) - assert set(result.keys()) == {'a', 'b', 'c'} - - def test_serialize_excludes_none(self) -> None: - """Serialization excludes None values.""" - data = {'key1': _make_db_value('test')} - result = LocalVectorStoreAPI._serialize_data(data) - serialized = result['key1'] - assert isinstance(serialized, dict) - for _key, value in _flatten_dict(serialized): - assert value is not None - - -class TestDeserializeData: - """Tests for data deserialization.""" - - def test_deserialize_empty_data(self) -> None: - """Deserialize empty dict returns empty dict.""" - result = LocalVectorStoreAPI._deserialize_data({}) - assert result == {} - - def test_roundtrip_serialize_deserialize(self) -> None: - """Serialize then deserialize produces equivalent data.""" - original = { - 'key1': _make_db_value('hello world', [0.5, -0.3, 0.8]), - 'key2': _make_db_value('goodbye', [0.1, 0.9, -0.2]), - } - serialized = LocalVectorStoreAPI._serialize_data(original) - deserialized = LocalVectorStoreAPI._deserialize_data(serialized) - - assert set(deserialized.keys()) == set(original.keys()) - for key in original: - assert deserialized[key].embedding.embedding == original[key].embedding.embedding - orig_text = original[key].doc.content[0].root.text - deser_text = deserialized[key].doc.content[0].root.text - assert deser_text == orig_text - - -class TestFileStoreOperations: - """Tests for file store load/dump operations.""" - - @pytest.mark.asyncio - async def test_load_nonexistent_file_returns_empty(self) -> None: - """Loading from nonexistent file returns empty dict.""" - api = LocalVectorStoreAPI(index_name='nonexistent_xyz_test') - result = await api._load_filestore() - assert result == {} - - @pytest.mark.asyncio - async def test_dump_and_load_roundtrip(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Dump then load produces equivalent data.""" - with tempfile.TemporaryDirectory() as tmpdir: - index_name = 'roundtrip_test' - - new_template = str(pathlib.Path(tmpdir) / '__db_{index_name}.json') - monkeypatch.setattr(LocalVectorStoreAPI, '_LOCAL_FILESTORE_TEMPLATE', new_template) - api = LocalVectorStoreAPI(index_name=index_name) - - data = { - 'doc1': _make_db_value('stored document', [0.1, 0.2]), - } - await api._dump_filestore(data) - - loaded = await api._load_filestore() - assert 'doc1' in loaded - assert loaded['doc1'].embedding.embedding == [0.1, 0.2] - - @pytest.mark.asyncio - async def test_dump_creates_valid_json(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Dump creates a valid JSON file.""" - with tempfile.TemporaryDirectory() as tmpdir: - index_name = 'json_test' - new_template = str(pathlib.Path(tmpdir) / '__db_{index_name}.json') - monkeypatch.setattr(LocalVectorStoreAPI, '_LOCAL_FILESTORE_TEMPLATE', new_template) - api = LocalVectorStoreAPI(index_name=index_name) - - data = {'key': _make_db_value('test')} - await api._dump_filestore(data) - - assert await aiofiles.os.path.exists(api.index_file_name) - async with aiofiles.open(api.index_file_name, encoding='utf-8') as f: - content = json.loads(await f.read()) - assert isinstance(content, dict) - assert 'key' in content - - -def _flatten_dict(d: dict, prefix: str = '') -> list[tuple[str, object]]: - """Flatten a nested dict for inspection.""" - items: list[tuple[str, object]] = [] - for k, v in d.items(): - new_key = f'{prefix}.{k}' if prefix else k - if isinstance(v, dict): - items.extend(_flatten_dict(v, new_key)) - else: - items.append((new_key, v)) - return items diff --git a/py/plugins/dev-local-vectorstore/tests/dlvs_indexer_test.py b/py/plugins/dev-local-vectorstore/tests/dlvs_indexer_test.py deleted file mode 100644 index 9a299a0aba..0000000000 --- a/py/plugins/dev-local-vectorstore/tests/dlvs_indexer_test.py +++ /dev/null @@ -1,192 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for DevLocalVectorStoreIndexer.""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from genkit.blocks.document import Document -from genkit.plugins.dev_local_vectorstore.constant import DbValue -from genkit.plugins.dev_local_vectorstore.indexer import DevLocalVectorStoreIndexer -from genkit.types import DocumentData, Embedding, TextPart - - -class TestIndexerInit: - """Tests for DevLocalVectorStoreIndexer initialization.""" - - def test_init_stores_parameters(self) -> None: - """Constructor stores ai, index_name, embedder, and options.""" - ai = MagicMock() - indexer = DevLocalVectorStoreIndexer( - ai=ai, - index_name='test-idx', - embedder='test-embedder', - embedder_options={'dim': 128}, - ) - assert indexer.ai is ai - assert indexer.index_name == 'test-idx' - assert indexer.embedder == 'test-embedder' - assert indexer.embedder_options == {'dim': 128} - - def test_init_default_options(self) -> None: - """Default embedder_options is None.""" - ai = MagicMock() - indexer = DevLocalVectorStoreIndexer( - ai=ai, - index_name='idx', - embedder='emb', - ) - assert indexer.embedder_options is None - - -class TestAddDocument: - """Tests for _add_document method.""" - - def test_add_document_creates_entry(self) -> None: - """Adding a document creates an entry in the data dict.""" - ai = MagicMock() - indexer = DevLocalVectorStoreIndexer(ai=ai, index_name='test', embedder='emb') - data: dict[str, DbValue] = {} - doc = Document.from_text('hello') - embedding = Embedding(embedding=[0.1, 0.2, 0.3]) - - indexer._add_document(data=data, embedding=embedding, doc=doc) - - assert len(data) == 1 - - def test_add_document_uses_md5_key(self) -> None: - """Document key is MD5 hash of serialized data.""" - ai = MagicMock() - indexer = DevLocalVectorStoreIndexer(ai=ai, index_name='test', embedder='emb') - data: dict[str, DbValue] = {} - doc = Document.from_text('test') - embedding = Embedding(embedding=[0.5]) - - indexer._add_document(data=data, embedding=embedding, doc=doc) - - # The key should be a hex MD5 hash - key = next(iter(data.keys())) - assert len(key) == 32 # MD5 hex digest length - assert all(c in '0123456789abcdef' for c in key) - - def test_add_document_with_existing_key_no_overwrite(self) -> None: - """Adding a document whose key already exists doesn't overwrite.""" - ai = MagicMock() - indexer = DevLocalVectorStoreIndexer(ai=ai, index_name='test', embedder='emb') - - # First, add a document to populate a key - data: dict[str, DbValue] = {} - doc = Document.from_text('original') - embedding = Embedding(embedding=[0.1]) - indexer._add_document(data=data, embedding=embedding, doc=doc) - - assert len(data) == 1 - key = next(iter(data.keys())) - data[key] - - # Create a different DbValue but use the same key by pre-inserting - new_doc = Document.from_text('replacement') - new_embedding = Embedding(embedding=[0.99]) - DbValue( - # Pydantic's discriminated union accepts TextPart directly at runtime, - # but the static type is list[Part]. Wrapping in Part(root=...) causes - # a ValidationError, so the type: ignore is intentional. - doc=DocumentData(content=[TextPart(text='replacement')]), # type: ignore[list-item] - embedding=new_embedding, - ) - - # Manually set the same key β€” _add_document should skip it - # (this is the actual dedup behavior: if idx already in data, skip) - data_copy = dict(data) - indexer._add_document(data=data_copy, embedding=new_embedding, doc=new_doc) - - # Should have 2 entries since different content produces different hash - assert len(data_copy) == 2 - - def test_add_different_documents(self) -> None: - """Adding different documents creates separate entries.""" - ai = MagicMock() - indexer = DevLocalVectorStoreIndexer(ai=ai, index_name='test', embedder='emb') - data: dict[str, DbValue] = {} - - doc1 = Document.from_text('first') - doc2 = Document.from_text('second') - emb1 = Embedding(embedding=[0.1]) - emb2 = Embedding(embedding=[0.9]) - - indexer._add_document(data=data, embedding=emb1, doc=doc1) - indexer._add_document(data=data, embedding=emb2, doc=doc2) - - assert len(data) == 2 - - -class TestIndexMethod: - """Tests for the async index method.""" - - @pytest.mark.asyncio - async def test_index_calls_embed_many(self) -> None: - """Index method calls ai.embed_many with correct parameters.""" - ai = MagicMock() - ai.embed_many = AsyncMock( - return_value=[ - Embedding(embedding=[0.1, 0.2]), - Embedding(embedding=[0.3, 0.4]), - ] - ) - - indexer = DevLocalVectorStoreIndexer(ai=ai, index_name='test', embedder='my-embedder') - - from genkit.blocks.retriever import IndexerRequest - - docs = [ - DocumentData(content=[TextPart(text='doc1')]), # type: ignore[list-item] - DocumentData(content=[TextPart(text='doc2')]), # type: ignore[list-item] - ] - request = IndexerRequest(documents=docs) - - with patch.object(indexer, '_load_filestore', return_value={}): - with patch('aiofiles.open', new_callable=MagicMock) as mock_aiofiles: - mock_file = AsyncMock() - mock_aiofiles.return_value.__aenter__ = AsyncMock(return_value=mock_file) - mock_aiofiles.return_value.__aexit__ = AsyncMock(return_value=False) - - await indexer.index(request) - - ai.embed_many.assert_called_once_with( - embedder='my-embedder', - content=docs, - options=None, - ) - - @pytest.mark.asyncio - async def test_index_raises_on_empty_embeddings(self) -> None: - """Index raises ValueError when embedder returns empty response.""" - ai = MagicMock() - ai.embed_many = AsyncMock(return_value=[]) - - indexer = DevLocalVectorStoreIndexer(ai=ai, index_name='test', embedder='emb') - - from genkit.blocks.retriever import IndexerRequest - - request = IndexerRequest( - documents=[DocumentData(content=[TextPart(text='doc')])] # type: ignore[list-item] - ) - - with patch.object(indexer, '_load_filestore', return_value={}): - with pytest.raises(ValueError, match='no embeddings'): - await indexer.index(request) diff --git a/py/plugins/dev-local-vectorstore/tests/dlvs_retriever_test.py b/py/plugins/dev-local-vectorstore/tests/dlvs_retriever_test.py deleted file mode 100644 index 4414d31910..0000000000 --- a/py/plugins/dev-local-vectorstore/tests/dlvs_retriever_test.py +++ /dev/null @@ -1,131 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for DevLocalVectorStoreRetriever.""" - -from genkit.plugins.dev_local_vectorstore.retriever import ( - DevLocalVectorStoreRetriever, - RetrieverOptionsSchema, - ScoredDocument, -) - - -class TestCosineSimilarity: - """Tests for cosine similarity calculation.""" - - def test_identical_vectors(self) -> None: - """Identical vectors have cosine similarity of 1.0.""" - a = [1.0, 2.0, 3.0] - result = DevLocalVectorStoreRetriever.cosine_similarity(a, a) - assert abs(result - 1.0) < 1e-9 - - def test_orthogonal_vectors(self) -> None: - """Orthogonal vectors have cosine similarity of 0.0.""" - a = [1.0, 0.0] - b = [0.0, 1.0] - result = DevLocalVectorStoreRetriever.cosine_similarity(a, b) - assert abs(result) < 1e-9 - - def test_opposite_vectors(self) -> None: - """Opposite vectors have cosine similarity of -1.0.""" - a = [1.0, 0.0] - b = [-1.0, 0.0] - result = DevLocalVectorStoreRetriever.cosine_similarity(a, b) - assert abs(result - (-1.0)) < 1e-9 - - def test_similar_vectors(self) -> None: - """Similar vectors have high cosine similarity.""" - a = [1.0, 2.0, 3.0] - b = [1.1, 2.1, 3.1] - result = DevLocalVectorStoreRetriever.cosine_similarity(a, b) - assert result > 0.99 - - def test_dissimilar_vectors(self) -> None: - """Dissimilar vectors have lower cosine similarity.""" - a = [1.0, 0.0, 0.0] - b = [0.0, 0.0, 1.0] - result = DevLocalVectorStoreRetriever.cosine_similarity(a, b) - assert abs(result) < 0.01 - - -class TestDotProduct: - """Tests for dot product calculation.""" - - def test_basic_dot_product(self) -> None: - """Dot product of [1,2,3] and [4,5,6] = 32.""" - result = DevLocalVectorStoreRetriever.dot([1.0, 2.0, 3.0], [4.0, 5.0, 6.0]) - assert abs(result - 32.0) < 1e-9 - - def test_zero_vector(self) -> None: - """Dot product with zero vector is 0.""" - result = DevLocalVectorStoreRetriever.dot([1.0, 2.0], [0.0, 0.0]) - assert abs(result) < 1e-9 - - def test_single_dimension(self) -> None: - """Dot product of single-dimension vectors.""" - result = DevLocalVectorStoreRetriever.dot([3.0], [4.0]) - assert abs(result - 12.0) < 1e-9 - - def test_negative_values(self) -> None: - """Dot product with negative values.""" - result = DevLocalVectorStoreRetriever.dot([1.0, -2.0], [-3.0, 4.0]) - assert abs(result - (-11.0)) < 1e-9 - - -class TestScoredDocument: - """Tests for ScoredDocument model.""" - - def test_create_scored_document(self) -> None: - """ScoredDocument can be created with score and document.""" - from genkit.ai import Document - - doc = Document.from_text('test') - scored = ScoredDocument(score=0.95, document=doc) - assert scored.score == 0.95 - assert scored.document is doc - - def test_scored_document_ordering(self) -> None: - """ScoredDocuments can be sorted by score.""" - from genkit.ai import Document - - docs = [ - ScoredDocument(score=0.3, document=Document.from_text('low')), - ScoredDocument(score=0.9, document=Document.from_text('high')), - ScoredDocument(score=0.6, document=Document.from_text('mid')), - ] - sorted_docs = sorted(docs, key=lambda d: d.score, reverse=True) - assert sorted_docs[0].score == 0.9 - assert sorted_docs[1].score == 0.6 - assert sorted_docs[2].score == 0.3 - - -class TestRetrieverOptionsSchema: - """Tests for RetrieverOptionsSchema.""" - - def test_default_limit_is_none(self) -> None: - """Default limit is None.""" - opts = RetrieverOptionsSchema() - assert opts.limit is None - - def test_custom_limit(self) -> None: - """Custom limit value is preserved.""" - opts = RetrieverOptionsSchema(limit=5) - assert opts.limit == 5 - - def test_schema_serialization(self) -> None: - """Schema can be serialized to JSON.""" - schema = RetrieverOptionsSchema.model_json_schema() - assert 'limit' in schema.get('properties', {}) diff --git a/py/plugins/fastapi/src/genkit/plugins/fastapi/handler.py b/py/plugins/fastapi/src/genkit/plugins/fastapi/handler.py index 145f8b657d..e6a66285c5 100644 --- a/py/plugins/fastapi/src/genkit/plugins/fastapi/handler.py +++ b/py/plugins/fastapi/src/genkit/plugins/fastapi/handler.py @@ -17,15 +17,21 @@ """Genkit FastAPI handler for serving flows as HTTP endpoints.""" import asyncio +import json from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any +from pydantic import BaseModel + from fastapi import Request, Response from fastapi.responses import StreamingResponse -from genkit.ai import FlowWrapper, Genkit -from genkit.codec import dump_dict, dump_json -from genkit.core.context import ContextProvider, RequestData -from genkit.core.error import GenkitError, get_callable_json +from genkit import Action, Genkit, GenkitError +from genkit.plugin_api import ContextProvider, RequestData, get_callable_json + + +def _to_dict(obj: Any) -> Any: # noqa: ANN401 + """Convert object to dict if it's a Pydantic model, otherwise return as-is.""" + return obj.model_dump() if isinstance(obj, BaseModel) else obj class _FastAPIRequestData(RequestData): @@ -41,7 +47,7 @@ def __init__(self, request: Request, body: dict[str, Any] | None) -> None: def genkit_fastapi_handler( ai: Genkit, context_provider: ContextProvider | None = None, -) -> Callable[[Callable[[], FlowWrapper]], Callable[[Request], Awaitable[Response | dict[str, Any]]]]: +) -> Callable[[Callable[[], Action]], Callable[[Request], Awaitable[Response | dict[str, Any]]]]: """Decorator for serving Genkit flows via FastAPI. Example: @@ -71,11 +77,11 @@ async def chat(): context_provider: Optional function to extract context from the request. Returns: - A decorator that wraps a function returning a FlowWrapper. + A decorator that wraps a function returning a Action. """ def decorator( - fn: Callable[[], FlowWrapper], + fn: Callable[[], Action], ) -> Callable[[Request], Awaitable[Response | dict[str, Any]]]: async def handler(request: Request) -> Response | dict[str, Any]: result = fn() @@ -83,7 +89,7 @@ async def handler(request: Request) -> Response | dict[str, Any]: if asyncio.iscoroutine(result): result = await result flow = result - if not isinstance(flow, FlowWrapper): + if not isinstance(flow, Action): raise GenkitError( status='INVALID_ARGUMENT', message='genkit_fastapi_handler must wrap a function that returns a @flow', @@ -93,11 +99,11 @@ async def handler(request: Request) -> Response | dict[str, Any]: if 'data' not in body: err = GenkitError( status='INVALID_ARGUMENT', - message='Flow request must be wrapped in {"data": ...} object', + message='Action request must be wrapped in {"data": ...} object', ) return Response( status_code=400, - content=dump_json(get_callable_json(err)), + content=json.dumps(get_callable_json(err), separators=(',', ':')), media_type='application/json', ) @@ -119,26 +125,26 @@ async def handler(request: Request) -> Response | dict[str, Any]: async def event_stream() -> AsyncIterator[str]: try: - stream_iter, response_future = flow._action.stream(body.get('data'), context=action_context) - async for chunk in stream_iter: - yield f'data: {dump_json({"message": dump_dict(chunk)})}\n\n' + stream_response = flow.stream(body.get('data'), context=action_context) + async for chunk in stream_response.stream: + yield f'data: {json.dumps({"message": _to_dict(chunk)}, separators=(",", ":"))}\n\n' - result = await response_future - yield f'data: {dump_json({"result": dump_dict(result.response)})}\n\n' + result = await stream_response.response + yield f'data: {json.dumps({"result": _to_dict(result)}, separators=(",", ":"))}\n\n' except Exception as e: ex = e.cause if isinstance(e, GenkitError) else e - yield f'error: {dump_json({"error": dump_dict(get_callable_json(ex))})}' + yield f'error: {json.dumps({"error": _to_dict(get_callable_json(ex))}, separators=(",", ":"))}' return StreamingResponse(event_stream(), media_type='text/event-stream') else: try: - response = await flow._action.arun_raw(body.get('data'), context=action_context) - return {'result': dump_dict(response.response)} + response = await flow.run(body.get('data'), context=action_context) + return {'result': _to_dict(response.response)} except Exception as e: ex = e.cause if isinstance(e, GenkitError) else e return Response( status_code=500, - content=dump_json(get_callable_json(ex)), + content=json.dumps(get_callable_json(ex), separators=(',', ':')), media_type='application/json', ) diff --git a/py/plugins/firebase/LICENSE b/py/plugins/firebase/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/plugins/firebase/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/plugins/firebase/README.md b/py/plugins/firebase/README.md deleted file mode 100644 index 856daf4004..0000000000 --- a/py/plugins/firebase/README.md +++ /dev/null @@ -1,27 +0,0 @@ -# Genkit Firebase plugin - -This Genkit plugin provides a set of tools and utilities for working with -Firebase. - -## Telemetry - -The Firebase plugin provides easy integration with Google Cloud Observability (Cloud Trace and Cloud Monitoring). - -To enable telemetry: - -```python -from genkit.plugins.firebase import add_firebase_telemetry - -# Enable telemetry (defaults to production-only export) -add_firebase_telemetry() -``` - -### Configuration - -`add_firebase_telemetry` supports the following options: - -- `project_id`: Firebase project ID (optional, auto-detected). -- `force_dev_export`: Set to `True` to export telemetry in dev environment (defaults to `False`). -- `log_input_and_output`: Set to `True` to log model inputs and outputs (defaults to `False` / redacted). -- `disable_metrics`: Set to `True` to disable metrics export. -- `disable_traces`: Set to `True` to disable trace export. diff --git a/py/plugins/firebase/pyproject.toml b/py/plugins/firebase/pyproject.toml deleted file mode 100644 index 9fdfd37d70..0000000000 --- a/py/plugins/firebase/pyproject.toml +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Framework :: AsyncIO", - "Framework :: Pydantic", - "Framework :: Pydantic :: 2", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Typing :: Typed", - "License :: OSI Approved :: Apache Software License", -] -dependencies = [ - "genkit", - "genkit-plugin-google-cloud", - "google-cloud-firestore", - "strenum>=0.4.15; python_version < '3.11'", -] -description = "Genkit Firebase Plugin" -keywords = [ - "genkit", - "ai", - "llm", - "machine-learning", - "artificial-intelligence", - "generative-ai", - "firebase", - "google", - "firestore", - "telemetry", -] -license = "Apache-2.0" -name = "genkit-plugin-firebase" -readme = "README.md" -requires-python = ">=3.10" -version = "0.5.1" - -[project.optional-dependencies] -# Telemetry export uses the Google Cloud telemetry exporters (Cloud Trace/Monitoring). -telemetry = ["genkit-plugin-google-cloud"] - -[project.urls] -"Bug Tracker" = "https://github.com/firebase/genkit/issues" -Changelog = "https://github.com/firebase/genkit/blob/main/py/plugins/firebase/CHANGELOG.md" -"Documentation" = "https://firebase.google.com/docs/genkit" -"Homepage" = "https://github.com/firebase/genkit" -"Repository" = "https://github.com/firebase/genkit/tree/main/py" - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -only-include = ["src/genkit/plugins/firebase"] -sources = ["src"] diff --git a/py/plugins/firebase/src/genkit/plugins/firebase/__init__.py b/py/plugins/firebase/src/genkit/plugins/firebase/__init__.py deleted file mode 100644 index 1cd1a84508..0000000000 --- a/py/plugins/firebase/src/genkit/plugins/firebase/__init__.py +++ /dev/null @@ -1,267 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Firebase plugin for Genkit. - -This plugin provides Firebase integrations for Genkit, including Firestore -vector stores for RAG and Firebase telemetry export to Google Cloud. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Firebase β”‚ Google's app development platform. Like a β”‚ - β”‚ β”‚ toolbox for building apps with database, auth. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Firestore β”‚ A NoSQL database that syncs in real-time. β”‚ - β”‚ β”‚ Store data as flexible documents. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Vector Store β”‚ A database that can find "similar" items. β”‚ - β”‚ β”‚ Like Google but for YOUR documents. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ RAG β”‚ Retrieval-Augmented Generation. AI looks up β”‚ - β”‚ β”‚ your docs before answering. Fewer hallucinations! β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Embeddings β”‚ Convert text to numbers that capture meaning. β”‚ - β”‚ β”‚ "Cat" and "kitten" become similar numbers. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Indexer β”‚ Stores documents with their embeddings. β”‚ - β”‚ β”‚ Like adding books to a library catalog. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Retriever β”‚ Finds documents matching a query. β”‚ - β”‚ β”‚ Like a librarian finding relevant books. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow (RAG with Firestore):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW FIRESTORE VECTOR SEARCH WORKS β”‚ - β”‚ β”‚ - β”‚ STEP 1: INDEXING (Store your documents) β”‚ - β”‚ ──────────────────────────────────────── β”‚ - β”‚ Your Documents β”‚ - β”‚ ["How to reset password", "Billing FAQ", ...] β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) Convert text to embeddings β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ "Reset password" β†’ [0.12, -0.34, ...] β”‚ - β”‚ β”‚ (Gemini, etc.) β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Store in Firestore with vectors β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Firestore β”‚ Document + embedding stored together β”‚ - β”‚ β”‚ (Vector Index) β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ - β”‚ STEP 2: RETRIEVAL (Find relevant documents) β”‚ - β”‚ ───────────────────────────────────────────── β”‚ - β”‚ User Query: "How do I change my password?" β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Convert query to embedding β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ Query β†’ [0.11, -0.33, ...] (similar!) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (4) Find nearest neighbors β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Firestore β”‚ "Reset password" doc is 95% match! β”‚ - β”‚ β”‚ Vector Search β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (5) Return matching documents β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Your App β”‚ AI uses these docs to answer accurately β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Architecture Overview:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Firebase Plugin β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Plugin Entry Point (__init__.py) β”‚ - β”‚ β”œβ”€β”€ define_firestore_vector_store() - Create Firestore vector store β”‚ - β”‚ └── add_firebase_telemetry() - Enable Cloud observability β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ firestore.py - Firestore Vector Store β”‚ - β”‚ β”œβ”€β”€ define_firestore_vector_store() - Main factory function β”‚ - β”‚ β”œβ”€β”€ Firestore indexer implementation β”‚ - β”‚ └── Firestore retriever implementation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ retriever.py - Retriever Implementation β”‚ - β”‚ └── FirestoreRetriever (vector similarity search) β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Firestore Vector Store Flow β”‚ - β”‚ β”‚ - β”‚ Documents ──► Embedder ──► Firestore (with vector index) β”‚ - β”‚ β”‚ - β”‚ Query ──► Embedder ──► Firestore Vector Search ──► Results β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Overview: - The Firebase plugin enables: - - Firestore as a vector store for document retrieval (RAG) - - Telemetry export to Google Cloud Trace and Monitoring - -Key Components: - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Component β”‚ Purpose β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ define_firestore_vector_storeβ”‚ Create a Firestore-backed vector store β”‚ - β”‚ add_firebase_telemetry() β”‚ Enable Cloud Trace/Monitoring export β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Example: - Using Firestore vector store: - - ```python - from genkit import Genkit - from genkit.plugins.firebase import define_firestore_vector_store - - ai = Genkit(...) - - # Define a Firestore vector store - store = define_firestore_vector_store( - ai, - name='my_store', - collection='documents', - embedder='vertexai/text-embedding-005', - ) - - # Index documents - await ai.index(indexer=store.indexer, documents=[...]) - - # Retrieve documents - docs = await ai.retrieve(retriever=store.retriever, query='...') - ``` - - Enabling telemetry: - - ```python - from genkit.plugins.firebase import add_firebase_telemetry - - # Export traces to Cloud Trace (disabled in dev mode by default) - add_firebase_telemetry() - ``` - -Caveats: - - Requires firebase-admin SDK and Google Cloud credentials - - Telemetry is disabled by default in development mode (GENKIT_ENV=dev) - -See Also: - - Firestore: https://firebase.google.com/docs/firestore - - Genkit documentation: https://genkit.dev/ -""" - -from typing import Any - -from opentelemetry.sdk.trace.sampling import Sampler - -from .constant import FirebaseTelemetryConfig -from .firestore import define_firestore_vector_store - - -def package_name() -> str: - """Get the package name for the Firebase plugin. - - Returns: - The fully qualified package name as a string. - """ - return 'genkit.plugins.firebase' - - -def add_firebase_telemetry( - config: FirebaseTelemetryConfig | None = None, - *, - project_id: str | None = None, - credentials: dict[str, Any] | None = None, - sampler: Sampler | None = None, - log_input_and_output: bool = False, - force_dev_export: bool = False, - disable_metrics: bool = False, - disable_traces: bool = False, - metric_export_interval_ms: int | None = None, - metric_export_timeout_ms: int | None = None, -) -> None: - """Add Firebase telemetry export to Google Cloud Observability. - - Exports traces to Cloud Trace, metrics to Cloud Monitoring, and logs to - Cloud Logging. In development (GENKIT_ENV=dev), telemetry is disabled by - default unless force_dev_export is True. - - Args: - config: FirebaseTelemetryConfig object. If provided, kwargs are ignored. - project_id: Firebase project ID. Auto-detected from environment if None. - credentials: Service account credentials dictionary. - sampler: OpenTelemetry trace sampler. - log_input_and_output: If True, logs feature inputs/outputs. WARNING: May log PII. - force_dev_export: If True, exports in dev mode. - disable_metrics: If True, disables metrics export. - disable_traces: If True, disables trace export. - metric_export_interval_ms: Metrics export interval in ms. Minimum: 1000ms. - metric_export_timeout_ms: Metrics export timeout in ms. - - Example:: - - # Using kwargs - add_firebase_telemetry(project_id='my-project', log_input_and_output=True) - - # Using config object - config = FirebaseTelemetryConfig(project_id='my-project') - add_firebase_telemetry(config) - """ - try: - # Imported lazily so Firestore-only users don't need telemetry deps. - from .telemetry import add_firebase_telemetry as _add_firebase_telemetry - except ImportError as e: - raise ImportError( - 'Firebase telemetry requires the Google Cloud telemetry exporter. ' - 'Install it with: pip install "genkit-plugin-firebase[telemetry]"' - ) from e - - if config is not None: - _add_firebase_telemetry(config) - else: - _add_firebase_telemetry( - FirebaseTelemetryConfig( - project_id=project_id, - credentials=credentials, - sampler=sampler, - log_input_and_output=log_input_and_output, - force_dev_export=force_dev_export, - disable_metrics=disable_metrics, - disable_traces=disable_traces, - metric_export_interval_ms=metric_export_interval_ms, - metric_export_timeout_ms=metric_export_timeout_ms, - ) - ) - - -__all__ = [ - 'add_firebase_telemetry', - 'define_firestore_vector_store', - 'FirebaseTelemetryConfig', - 'package_name', -] diff --git a/py/plugins/firebase/src/genkit/plugins/firebase/constant.py b/py/plugins/firebase/src/genkit/plugins/firebase/constant.py deleted file mode 100644 index aa9929e06a..0000000000 --- a/py/plugins/firebase/src/genkit/plugins/firebase/constant.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Firebase constants and configuration models.""" - -from collections.abc import Callable -from typing import Annotated, Any - -from google.cloud.firestore_v1 import DocumentSnapshot -from opentelemetry.sdk.trace.sampling import Sampler -from pydantic import BaseModel, Field - -MetadataTransformFn = Callable[[DocumentSnapshot], dict[str, Any]] - - -class FirebaseTelemetryConfig(BaseModel): - """Configuration for Firebase telemetry export to Google Cloud Observability. - - Args: - project_id: Firebase project ID. Auto-detected from environment if None. - credentials: Service account credentials dictionary. - sampler: OpenTelemetry trace sampler. - log_input_and_output: If True, logs feature inputs/outputs. WARNING: May log PII. - force_dev_export: If True, exports telemetry in dev mode (GENKIT_ENV=dev). - disable_metrics: If True, disables metrics export. - disable_traces: If True, disables trace export. - metric_export_interval_ms: Metrics export interval in ms. Minimum: 1000ms. - metric_export_timeout_ms: Metrics export timeout in ms. - """ - - project_id: str | None = None - credentials: dict[str, Any] | None = None - sampler: Sampler | None = None - log_input_and_output: bool = False - force_dev_export: bool = False - disable_metrics: bool = False - disable_traces: bool = False - metric_export_interval_ms: Annotated[int, Field(ge=1000)] | None = None - metric_export_timeout_ms: int | None = None - model_config = {'arbitrary_types_allowed': True} diff --git a/py/plugins/firebase/src/genkit/plugins/firebase/firestore.py b/py/plugins/firebase/src/genkit/plugins/firebase/firestore.py deleted file mode 100644 index 10e96f44c8..0000000000 --- a/py/plugins/firebase/src/genkit/plugins/firebase/firestore.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Firestore vector store operations for Genkit.""" - -from collections.abc import Callable - -from google.cloud import firestore -from google.cloud.firestore_v1 import DocumentSnapshot -from google.cloud.firestore_v1.base_vector_query import DistanceMeasure - -from genkit.ai import Genkit -from genkit.blocks.retriever import RetrieverOptions, retriever_action_metadata -from genkit.core.action.types import ActionKind -from genkit.core.typing import DocumentPart -from genkit.plugins.firebase.retriever import FirestoreRetriever - -from .constant import MetadataTransformFn - - -def firestore_action_name(name: str) -> str: - """Create a firestore action name. - - Args: - name: Base name for the action - - Returns: - str: Firestore action name. - - """ - return f'firestore/{name}' - - -def define_firestore_vector_store( - ai: Genkit, - *, - name: str, - embedder: str, - embedder_options: dict[str, object] | None = None, - collection: str, - vector_field: str, - content_field: str | Callable[[DocumentSnapshot], list['DocumentPart']], - firestore_client: firestore.Client, - distance_measure: DistanceMeasure = DistanceMeasure.COSINE, - metadata_fields: list[str] | MetadataTransformFn | None = None, -) -> str: - """Define and register a Firestore vector store retriever. - - Args: - ai: The Genkit instance to register the retriever with. - name: Name of the retriever. - embedder: The embedder to use (e.g., 'vertexai/gemini-embedding-001'). - embedder_options: Optional configuration to pass to the embedder. - collection: The name of the Firestore collection to query. - vector_field: The name of the field containing the vector embeddings. - content_field: The name of the field containing the document content, you wish to return. - firestore_client: The Firestore database instance from which to query. - distance_measure: The distance measure to use when comparing vectors. Defaults to 'COSINE'. - metadata_fields: Optional list of metadata fields to include. - - Returns: - The registered retriever name. - """ - retriever = FirestoreRetriever( - ai=ai, - name=name, - embedder=embedder, - embedder_options=embedder_options, - firestore_client=firestore_client, - collection=collection, - vector_field=vector_field, - content_field=content_field, - distance_measure=distance_measure, - metadata_fields=metadata_fields, - ) - - action_name = firestore_action_name(name) - - ai.registry.register_action( - kind=ActionKind.RETRIEVER, - name=action_name, - fn=retriever.retrieve, - metadata=retriever_action_metadata( - name=action_name, - options=RetrieverOptions(label=name), - ).metadata, - ) - - return action_name diff --git a/py/plugins/firebase/src/genkit/plugins/firebase/py.typed b/py/plugins/firebase/src/genkit/plugins/firebase/py.typed deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/plugins/firebase/src/genkit/plugins/firebase/retriever.py b/py/plugins/firebase/src/genkit/plugins/firebase/retriever.py deleted file mode 100644 index cdc9d0688b..0000000000 --- a/py/plugins/firebase/src/genkit/plugins/firebase/retriever.py +++ /dev/null @@ -1,196 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Firestore retriever implementation for Genkit.""" - -from collections.abc import Callable - -from google.cloud import firestore -from google.cloud.firestore_v1 import DocumentSnapshot -from google.cloud.firestore_v1.base_vector_query import DistanceMeasure -from google.cloud.firestore_v1.vector import Vector - -from genkit.ai import Genkit -from genkit.core.typing import DocumentPart, TextPart -from genkit.types import ActionRunContext, Document, GenkitError, RetrieverRequest, RetrieverResponse - -from .constant import MetadataTransformFn - - -class FirestoreRetriever: - """Retrieves documents from Google Cloud Firestore using vector similarity search. - - Attributes: - ai: Genkit instance used to embed queries. - name: Name of the retriever. - embedder: The embedder to use for query embeddings. - embedder_options: Optional configuration to pass to the embedder. - firestore_client: The initialized Firestore client from the configuration. - collection: The name of the Firestore collection to query. - vector_field: The name of the field containing the vector embeddings. - content_field: The name of the field containing the document content. - distance_measure: The distance measure to use when comparing vectors. - metadata_fields: Optional list of metadata fields to include. - """ - - def __init__( - self, - ai: Genkit, - name: str, - embedder: str, - embedder_options: dict[str, object] | None, - firestore_client: firestore.Client, - collection: str, - vector_field: str, - content_field: str | Callable[[DocumentSnapshot], list[DocumentPart]], - distance_measure: DistanceMeasure = DistanceMeasure.COSINE, - metadata_fields: list[str] | MetadataTransformFn | None = None, - ) -> None: - """Initialize the FirestoreRetriever. - - Args: - ai: Genkit instance used to embed queries. - name: Name of the retriever. - embedder: The embedder to use for query embeddings. - embedder_options: Optional configuration to pass to the embedder. - firestore_client: The Firestore database instance from which to query. - collection: The name of the Firestore collection to query. - vector_field: The name of the field containing the vector embeddings. - content_field: The name of the field containing the document content. - distance_measure: The distance measure to use when comparing vectors. - metadata_fields: Optional list of metadata fields to include. - """ - self.ai = ai - self.name = name - self.embedder = embedder - self.embedder_options = embedder_options - self.firestore_client = firestore_client - self.collection = collection - self.vector_field = vector_field - self.content_field = content_field - self.distance_measure = distance_measure - self.metadata_fields = metadata_fields - self._validate_config() - - def _validate_config(self) -> None: - """Validate the FirestoreRetriever configuration. - - Raises: - ValueError: If the configuration is invalid. - """ - if not self.collection: - raise ValueError('Firestore Retriever config must include firestore collection name.') - if not self.vector_field: - raise ValueError('Firestore Retriever config must include vector field name.') - if not self.embedder: - raise ValueError('Firestore Retriever config must include embedder name.') - if not self.firestore_client: - raise ValueError('Firestore Retriever config must include firestore client.') - - def _to_content(self, doc_snapshot: DocumentSnapshot) -> list[DocumentPart]: - """Convert a Firestore document snapshot to a list of content dictionaries. - - Args: - doc_snapshot: A Firestore DocumentSnapshot object. - - Returns: - A list of dictionaries containing the content of the document. - """ - content_field = self.content_field - if callable(content_field): - return content_field(doc_snapshot) - else: - content = doc_snapshot.get(content_field) - return [DocumentPart(root=TextPart(text=str(content)))] if content else [] - - def _to_metadata(self, doc_snapshot: DocumentSnapshot) -> dict[str, object]: - """Convert a Firestore document snapshot to a list of metadata dictionaries. - - Args: - doc_snapshot: A Firestore DocumentSnapshot object. - - Returns: - A list of dictionaries containing the metadata of the document. - """ - metadata: dict[str, object] = {} - metadata_fields = self.metadata_fields - if metadata_fields: - if callable(metadata_fields): - # pyrefly: ignore[bad-assignment] - MetadataTransformFn returns dict[str, Any] - metadata = metadata_fields(doc_snapshot) - else: - doc_dict = doc_snapshot.to_dict() or {} - for field in metadata_fields: - if field in doc_dict: - metadata[field] = doc_dict[field] - else: - metadata = doc_snapshot.to_dict() or {} - vector_field = self.vector_field - content_field = self.content_field - if vector_field in metadata: - del metadata[vector_field] - if isinstance(content_field, str) and content_field in metadata: - del metadata[content_field] - return metadata - - def _to_document(self, doc_snapshot: DocumentSnapshot) -> Document: - """Convert a Firestore document snapshot to a Genkit Document object. - - Args: - doc_snapshot: A Firestore DocumentSnapshot object. - - Returns: - A Genkit Document object. - """ - return Document(content=self._to_content(doc_snapshot), metadata=self._to_metadata(doc_snapshot)) - - async def retrieve(self, request: RetrieverRequest, _: ActionRunContext) -> RetrieverResponse: - """Retrieves documents from Firestore using native vector similarity search. - - Args: - request: A RetrieverRequest Object - - Returns: - A RetrieverResponse Object containing retrieved documents - """ - query = Document.from_document_data(document_data=request.query) - query_embedding_result = await self.ai.embed( - embedder=self.embedder, - content=query, - options=self.embedder_options, - ) - - if not query_embedding_result: - raise GenkitError(message='Embedder returned no embeddings') - - query_embedding = query_embedding_result[0].embedding - query_vector = Vector(query_embedding) - collection = self.firestore_client.collection(self.collection) - - limit = 10 - if isinstance(request.options, dict) and (limit_val := request.options.get('limit')) is not None: - limit = int(limit_val) - - vector_query = collection.find_nearest( - vector_field=self.vector_field, - query_vector=query_vector, - distance_measure=self.distance_measure, - limit=limit, - ) - query_snapshot = vector_query.get() - documents = [self._to_document(doc) for doc in query_snapshot] - return RetrieverResponse(documents=documents) diff --git a/py/plugins/firebase/src/genkit/plugins/firebase/telemetry.py b/py/plugins/firebase/src/genkit/plugins/firebase/telemetry.py deleted file mode 100644 index a2bd26220f..0000000000 --- a/py/plugins/firebase/src/genkit/plugins/firebase/telemetry.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Firebase telemetry integration.""" - -from genkit.core.logging import get_logger -from genkit.plugins.google_cloud.telemetry.config import GcpTelemetry - -from .constant import FirebaseTelemetryConfig - -logger = get_logger(__name__) - - -def add_firebase_telemetry(config: FirebaseTelemetryConfig) -> None: - """Add Firebase telemetry export to Google Cloud Observability. - - Exports traces to Cloud Trace, metrics to Cloud Monitoring, and logs to - Cloud Logging. In development (GENKIT_ENV=dev), telemetry is disabled by - default unless force_dev_export is set to True. - - Args: - config: FirebaseTelemetryConfig object with telemetry configuration. - """ - manager = GcpTelemetry( - project_id=config.project_id, - credentials=config.credentials, - sampler=config.sampler, - log_input_and_output=config.log_input_and_output, - force_dev_export=config.force_dev_export, - disable_metrics=config.disable_metrics, - disable_traces=config.disable_traces, - metric_export_interval_ms=config.metric_export_interval_ms, - metric_export_timeout_ms=config.metric_export_timeout_ms, - ) - - if not manager.project_id: - logger.warning( - 'Firebase project ID not found. Set FIREBASE_PROJECT_ID, GOOGLE_CLOUD_PROJECT, ' - 'or GCLOUD_PROJECT environment variable, or pass project_id parameter.' - ) - - manager.initialize() diff --git a/py/plugins/firebase/src/genkit/plugins/firebase/tests/firebase_telemetry_test.py b/py/plugins/firebase/src/genkit/plugins/firebase/tests/firebase_telemetry_test.py deleted file mode 100644 index 93b04ba194..0000000000 --- a/py/plugins/firebase/src/genkit/plugins/firebase/tests/firebase_telemetry_test.py +++ /dev/null @@ -1,139 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Firebase telemetry functionality.""" - -from unittest.mock import MagicMock, patch - -from opentelemetry.sdk.trace import ReadableSpan - -from genkit.plugins.firebase import add_firebase_telemetry -from genkit.plugins.google_cloud.telemetry.metrics import record_generate_metrics - - -def _create_model_span( - model_name: str = 'gemini-pro', - path: str = '/{myflow,t:flow}', - output: str = '{"usage": {"inputTokens": 100, "outputTokens": 50}}', - is_ok: bool = True, - start_time: int = 1000000000, - end_time: int = 1500000000, -) -> MagicMock: - """Helper function to create a model action span for testing. - - Args: - model_name: The model name for genkit:name attribute - path: The genkit:path value - output: The genkit:output JSON string - is_ok: Whether the span status is ok - start_time: Span start time in nanoseconds - end_time: Span end time in nanoseconds - - Returns: - A mocked ReadableSpan with model action attributes - """ - mock_span = MagicMock(spec=ReadableSpan) - mock_span.attributes = { - 'genkit:type': 'action', - 'genkit:metadata:subtype': 'model', - 'genkit:name': model_name, - 'genkit:path': path, - 'genkit:output': output, - } - mock_span.status.is_ok = is_ok - mock_span.start_time = start_time - mock_span.end_time = end_time - return mock_span - - -@patch('genkit.plugins.firebase.telemetry.GcpTelemetry') -def test_firebase_telemetry_initializes_gcp_telemetry(mock_gcp_telemetry_cls: MagicMock) -> None: - """Test that Firebase telemetry initializes GcpTelemetry with correct defaults.""" - mock_manager = MagicMock() - mock_gcp_telemetry_cls.return_value = mock_manager - - add_firebase_telemetry() - - mock_gcp_telemetry_cls.assert_called_once() - kwargs = mock_gcp_telemetry_cls.call_args.kwargs - assert kwargs['force_dev_export'] is False - mock_manager.initialize.assert_called_once() - - -def test_firebase_telemetry_passes_configuration() -> None: - """Test that configuration options are passed to GcpTelemetry.""" - with patch('genkit.plugins.firebase.telemetry.GcpTelemetry') as mock_gcp_telemetry_cls: - add_firebase_telemetry( - project_id='test-project', - log_input_and_output=True, - force_dev_export=True, - disable_metrics=True, - ) - - kwargs = mock_gcp_telemetry_cls.call_args.kwargs - assert kwargs['project_id'] == 'test-project' - assert kwargs['log_input_and_output'] is True - assert kwargs['force_dev_export'] is True - assert kwargs['disable_metrics'] is True - - -@patch('genkit.plugins.google_cloud.telemetry.metrics._output_tokens') -@patch('genkit.plugins.google_cloud.telemetry.metrics._input_tokens') -@patch('genkit.plugins.google_cloud.telemetry.metrics._latency') -@patch('genkit.plugins.google_cloud.telemetry.metrics._failures') -@patch('genkit.plugins.google_cloud.telemetry.metrics._requests') -def test_record_generate_metrics_with_model_action( - mock_requests: MagicMock, - mock_failures: MagicMock, - mock_latency: MagicMock, - mock_input_tokens: MagicMock, - mock_output_tokens: MagicMock, -) -> None: - """Test that metrics are recorded for model action spans with usage data.""" - # Setup mocks - mock_request_counter = MagicMock() - mock_latency_histogram = MagicMock() - mock_input_counter = MagicMock() - mock_output_counter = MagicMock() - - mock_requests.return_value = mock_request_counter - mock_failures.return_value = MagicMock() - mock_latency.return_value = mock_latency_histogram - mock_input_tokens.return_value = mock_input_counter - mock_output_tokens.return_value = mock_output_counter - - # Create test span using helper - mock_span = _create_model_span( - model_name='gemini-pro', - path='/{myflow,t:flow}', - output='{"usage": {"inputTokens": 100, "outputTokens": 50}}', - ) - - # Execute - record_generate_metrics(mock_span) - - # Verify dimensions - expected_dimensions = {'model': 'gemini-pro', 'source': 'myflow', 'error': 'none'} - - # Verify requests counter - mock_request_counter.add.assert_called_once_with(1, expected_dimensions) - - # Verify latency (500ms = 1.5s - 1.0s) - mock_latency_histogram.record.assert_called_once_with(500.0, expected_dimensions) - - # Verify token counts - mock_input_counter.add.assert_called_once_with(100, expected_dimensions) - mock_output_counter.add.assert_called_once_with(50, expected_dimensions) diff --git a/py/plugins/firebase/src/genkit/py.typed b/py/plugins/firebase/src/genkit/py.typed deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/py/plugins/firebase/tests/firebase_plugin_test.py b/py/plugins/firebase/tests/firebase_plugin_test.py deleted file mode 100644 index afc9f44303..0000000000 --- a/py/plugins/firebase/tests/firebase_plugin_test.py +++ /dev/null @@ -1,99 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Firebase plugin.""" - -import sys -from unittest.mock import MagicMock, patch - -import pytest - -from genkit.plugins.firebase import ( - FirebaseTelemetryConfig, - add_firebase_telemetry, - define_firestore_vector_store, - package_name, -) - - -def test_package_name() -> None: - """Test package_name returns correct value.""" - assert package_name() == 'genkit.plugins.firebase' - - -@patch('genkit.plugins.firebase.telemetry.GcpTelemetry') -def test_add_firebase_telemetry_calls_gcp_telemetry(mock_gcp_telemetry_cls: MagicMock) -> None: - """Test add_firebase_telemetry delegates to GCP telemetry.""" - mock_manager = MagicMock() - mock_gcp_telemetry_cls.return_value = mock_manager - - add_firebase_telemetry() - - mock_gcp_telemetry_cls.assert_called_once() - mock_manager.initialize.assert_called_once() - - -@patch('genkit.plugins.firebase.telemetry.GcpTelemetry') -def test_add_firebase_telemetry_with_config(mock_gcp_telemetry_cls: MagicMock) -> None: - """Test add_firebase_telemetry accepts config object.""" - mock_manager = MagicMock() - mock_gcp_telemetry_cls.return_value = mock_manager - - config = FirebaseTelemetryConfig( - project_id='test-project', - log_input_and_output=True, - force_dev_export=True, - ) - add_firebase_telemetry(config) - - mock_gcp_telemetry_cls.assert_called_once_with( - project_id='test-project', - credentials=None, - sampler=None, - log_input_and_output=True, - force_dev_export=True, - disable_metrics=False, - disable_traces=False, - metric_export_interval_ms=None, - metric_export_timeout_ms=None, - ) - mock_manager.initialize.assert_called_once() - - -def test_define_firestore_vector_store_exported() -> None: - """Test define_firestore_vector_store is exported.""" - # Just verify the function is importable and callable - assert callable(define_firestore_vector_store) - - -def test_add_firebase_telemetry_raises_on_missing_deps() -> None: - """Test that an informative ImportError is raised if telemetry deps are missing.""" - # Temporarily remove the module from sys.modules to simulate it not being installed. - with patch.dict(sys.modules, {'genkit.plugins.firebase.telemetry': None}): - with pytest.raises(ImportError, match='Firebase telemetry requires the Google Cloud telemetry exporter'): - add_firebase_telemetry() - - -def test_firebase_telemetry_config_validation() -> None: - """Test Pydantic validation on FirebaseTelemetryConfig.""" - # Valid config - config = FirebaseTelemetryConfig(project_id='test-project') - assert config.project_id == 'test-project' - assert config.log_input_and_output is False - - # Invalid metric interval (< 1000ms) should raise ValidationError - with pytest.raises(ValueError): - FirebaseTelemetryConfig(metric_export_interval_ms=500) diff --git a/py/plugins/firebase/tests/firebase_retriever_test.py b/py/plugins/firebase/tests/firebase_retriever_test.py deleted file mode 100644 index b04fe53c17..0000000000 --- a/py/plugins/firebase/tests/firebase_retriever_test.py +++ /dev/null @@ -1,210 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for FirestoreRetriever.""" - -from unittest.mock import MagicMock - -import pytest - -from genkit.core.typing import DocumentPart, TextPart -from genkit.plugins.firebase.retriever import FirestoreRetriever - - -def _make_retriever(**overrides: object) -> FirestoreRetriever: - """Create a FirestoreRetriever with sensible defaults for testing.""" - defaults: dict[str, object] = { - 'ai': MagicMock(), - 'name': 'test-retriever', - 'embedder': 'test-embedder', - 'embedder_options': None, - 'firestore_client': MagicMock(), - 'collection': 'test-collection', - 'vector_field': 'embedding', - 'content_field': 'content', - } - defaults.update(overrides) - return FirestoreRetriever(**defaults) # type: ignore[arg-type] - - -class TestFirestoreRetrieverInit: - """Tests for FirestoreRetriever initialization.""" - - def test_init_stores_all_params(self) -> None: - """Constructor stores all configuration parameters.""" - r = _make_retriever() - assert r.name == 'test-retriever' - assert r.embedder == 'test-embedder' - assert r.collection == 'test-collection' - assert r.vector_field == 'embedding' - assert r.content_field == 'content' - assert r.embedder_options is None - - -class TestFirestoreRetrieverValidation: - """Tests for config validation.""" - - def test_empty_collection_raises(self) -> None: - """Empty collection name raises ValueError.""" - with pytest.raises(ValueError, match='collection'): - _make_retriever(collection='') - - def test_empty_vector_field_raises(self) -> None: - """Empty vector field name raises ValueError.""" - with pytest.raises(ValueError, match='vector field'): - _make_retriever(vector_field='') - - def test_empty_embedder_raises(self) -> None: - """Empty embedder name raises ValueError.""" - with pytest.raises(ValueError, match='embedder'): - _make_retriever(embedder='') - - def test_none_firestore_client_raises(self) -> None: - """None firestore client raises ValueError.""" - with pytest.raises(ValueError, match='firestore client'): - _make_retriever(firestore_client=None) - - -class TestFirestoreRetrieverToContent: - """Tests for _to_content conversion.""" - - def test_string_content_field(self) -> None: - """String content_field reads from doc snapshot field.""" - r = _make_retriever(content_field='body') - doc_snapshot = MagicMock() - doc_snapshot.get.return_value = 'Hello world' - - parts = r._to_content(doc_snapshot) - - assert len(parts) == 1 - doc_snapshot.get.assert_called_once_with('body') - - def test_string_content_field_empty(self) -> None: - """Empty content returns empty list.""" - r = _make_retriever(content_field='body') - doc_snapshot = MagicMock() - doc_snapshot.get.return_value = None - - parts = r._to_content(doc_snapshot) - assert parts == [] - - def test_callable_content_field(self) -> None: - """Callable content_field is invoked with doc snapshot.""" - custom_content = [DocumentPart(root=TextPart(text='custom'))] - content_fn = MagicMock(return_value=custom_content) - r = _make_retriever(content_field=content_fn) - doc_snapshot = MagicMock() - - parts = r._to_content(doc_snapshot) - - content_fn.assert_called_once_with(doc_snapshot) - assert parts == custom_content - - -class TestFirestoreRetrieverToMetadata: - """Tests for _to_metadata conversion.""" - - def test_no_metadata_fields_returns_dict_minus_vector_content(self) -> None: - """Without metadata_fields config, returns all fields except vector and content.""" - r = _make_retriever( - vector_field='vec', - content_field='body', - metadata_fields=None, - ) - doc_snapshot = MagicMock() - doc_snapshot.to_dict.return_value = { - 'vec': [0.1, 0.2], - 'body': 'text', - 'author': 'Alice', - 'date': '2025-01-01', - } - - metadata = r._to_metadata(doc_snapshot) - - assert 'vec' not in metadata - assert 'body' not in metadata - assert metadata['author'] == 'Alice' - assert metadata['date'] == '2025-01-01' - - def test_metadata_fields_list_filters(self) -> None: - """List of metadata_fields only includes specified fields.""" - r = _make_retriever(metadata_fields=['author']) - doc_snapshot = MagicMock() - doc_snapshot.to_dict.return_value = { - 'author': 'Bob', - 'date': '2025-01-01', - 'hidden': 'secret', - } - - metadata = r._to_metadata(doc_snapshot) - - assert metadata == {'author': 'Bob'} - assert 'date' not in metadata - assert 'hidden' not in metadata - - def test_metadata_fields_list_missing_field_skipped(self) -> None: - """Missing fields in metadata_fields list are silently skipped.""" - r = _make_retriever(metadata_fields=['author', 'nonexistent']) - doc_snapshot = MagicMock() - doc_snapshot.to_dict.return_value = {'author': 'Carol'} - - metadata = r._to_metadata(doc_snapshot) - - assert metadata == {'author': 'Carol'} - - def test_callable_metadata_fields(self) -> None: - """Callable metadata_fields is invoked with doc snapshot.""" - meta_fn = MagicMock(return_value={'custom_key': 'custom_val'}) - r = _make_retriever(metadata_fields=meta_fn) - doc_snapshot = MagicMock() - - metadata = r._to_metadata(doc_snapshot) - - meta_fn.assert_called_once_with(doc_snapshot) - assert metadata == {'custom_key': 'custom_val'} - - -class TestFirestoreRetrieverToDocument: - """Tests for _to_document conversion.""" - - def test_to_document_combines_content_and_metadata(self) -> None: - """_to_document creates a Document with content and metadata.""" - r = _make_retriever(content_field='body', metadata_fields=['author']) - doc_snapshot = MagicMock() - doc_snapshot.get.return_value = 'Hello' - doc_snapshot.to_dict.return_value = {'body': 'Hello', 'author': 'Dave'} - - doc = r._to_document(doc_snapshot) - - assert len(doc.content) == 1 - assert doc.metadata is not None - assert doc.metadata['author'] == 'Dave' - - -class TestFirestoreActionName: - """Tests for firestore_action_name helper.""" - - def test_action_name_format(self) -> None: - """Action name is prefixed with 'firestore/'.""" - from genkit.plugins.firebase.firestore import firestore_action_name - - assert firestore_action_name('my-store') == 'firestore/my-store' - - def test_action_name_preserves_special_chars(self) -> None: - """Special characters in name are preserved.""" - from genkit.plugins.firebase.firestore import firestore_action_name - - assert firestore_action_name('my_store-v2') == 'firestore/my_store-v2' diff --git a/py/plugins/flask/pyproject.toml b/py/plugins/flask/pyproject.toml index 454d5dc9f6..4715271cbd 100644 --- a/py/plugins/flask/pyproject.toml +++ b/py/plugins/flask/pyproject.toml @@ -46,7 +46,7 @@ dependencies = [ "genkit", "genkit-plugin-google-genai", "pydantic>=2.10.5", - "flask", + "flask>=3.1.3", ] description = "Genkit Firebase Plugin" keywords = [ diff --git a/py/plugins/flask/src/genkit/plugins/flask/handler.py b/py/plugins/flask/src/genkit/plugins/flask/handler.py index a417e750fa..b143c69018 100644 --- a/py/plugins/flask/src/genkit/plugins/flask/handler.py +++ b/py/plugins/flask/src/genkit/plugins/flask/handler.py @@ -17,15 +17,57 @@ """Genkit Flask plugin.""" import asyncio -from collections.abc import AsyncIterator, Awaitable, Callable, Iterable -from typing import Any, TypeAlias +import json +from asyncio import AbstractEventLoop +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable +from typing import Any, TypeAlias, TypeVar + +from pydantic import BaseModel from flask import Response, request -from genkit.ai import FlowWrapper, Genkit -from genkit.aio.loop import create_loop, iter_over_async -from genkit.codec import dump_dict, dump_json -from genkit.core.context import ContextProvider, RequestData -from genkit.core.error import GenkitError, get_callable_json +from genkit import Genkit, GenkitError +from genkit._core._action import Action +from genkit.plugin_api import ( + ContextProvider, + RequestData, + get_callable_json, +) + + +def _to_dict(obj: Any) -> Any: # noqa: ANN401 + """Convert object to dict if it's a Pydantic model, otherwise return as-is.""" + return obj.model_dump() if isinstance(obj, BaseModel) else obj + + +T = TypeVar('T') + + +def _create_loop() -> AbstractEventLoop: + """Creates a new asyncio event loop or returns the current one.""" + try: + return asyncio.get_event_loop() + except Exception: + return asyncio.new_event_loop() + + +def _iter_over_async(ait: AsyncIterable[T], loop: AbstractEventLoop) -> Iterable[T]: + """Synchronously iterates over an AsyncIterable using a specified event loop.""" + ait_iter = ait.__aiter__() + + async def get_next() -> tuple[bool, T | None]: + try: + obj = await ait_iter.__anext__() + return False, obj + except StopAsyncIteration: + return True, None + + while True: + done, obj = loop.run_until_complete(get_next()) + if done: + break + assert obj is not None + yield obj + # Type alias for Flask-compatible route handler return type FlaskRouteReturn: TypeAlias = Response | dict[str, object] | Iterable[Any] @@ -47,7 +89,7 @@ def __init__(self) -> None: def genkit_flask_handler( ai: Genkit, context_provider: ContextProvider | None = None, -) -> Callable[[FlowWrapper], Callable[..., Awaitable[FlaskRouteReturn]]]: +) -> Callable[[Action], Callable[..., Awaitable[FlaskRouteReturn]]]: """A decorator for serving Genkit flows via a flask sever. ```python @@ -67,10 +109,10 @@ async def say_hi(name: str, ctx): ``` """ - loop = create_loop() + loop = _create_loop() - def decorator(flow: FlowWrapper) -> Callable[..., Awaitable[FlaskRouteReturn]]: - if not isinstance(flow, FlowWrapper): + def decorator(flow: Action) -> Callable[..., Awaitable[FlaskRouteReturn]]: + if not isinstance(flow, Action): raise GenkitError(status='INVALID_ARGUMENT', message='must apply @genkit_flask_handler on a @flow') async def handler() -> FlaskRouteReturn: @@ -93,29 +135,29 @@ async def handler() -> FlaskRouteReturn: async def async_gen() -> AsyncIterator[str]: try: - stream, response = flow._action.stream(input_data.get('data'), context=action_context) - async for chunk in stream: - yield f'data: {dump_json({"message": dump_dict(chunk)})}\n\n' + stream_response = flow.stream(input_data.get('data'), context=action_context) + async for chunk in stream_response.stream: + yield f'data: {json.dumps({"message": _to_dict(chunk)}, separators=(",", ":"))}\n\n' - result = await response - yield f'data: {dump_json({"result": dump_dict(result.response)})}\n\n' + result = await stream_response.response + yield f'data: {json.dumps({"result": _to_dict(result)}, separators=(",", ":"))}\n\n' except Exception as e: ex = e if isinstance(ex, GenkitError): ex = ex.cause - yield f'error: {dump_json({"error": dump_dict(get_callable_json(ex))})}' + yield f'error: {json.dumps({"error": _to_dict(get_callable_json(ex))}, separators=(",", ":"))}' - iter = iter_over_async(async_gen(), loop) + iter = _iter_over_async(async_gen(), loop) return iter else: try: - response = await flow._action.arun_raw(input_data.get('data'), context=action_context) - return {'result': dump_dict(response.response)} + response = await flow.run(input_data.get('data'), context=action_context) + return {'result': _to_dict(response.response)} except Exception as e: ex = e if isinstance(ex, GenkitError): ex = ex.cause - return Response(status=500, response=dump_json(get_callable_json(ex))) + return Response(status=500, response=json.dumps(get_callable_json(ex), separators=(',', ':'))) return handler diff --git a/py/plugins/flask/tests/flask_exports_test.py b/py/plugins/flask/tests/flask_exports_test.py index 2a871a0d12..aee9cb5d43 100644 --- a/py/plugins/flask/tests/flask_exports_test.py +++ b/py/plugins/flask/tests/flask_exports_test.py @@ -16,7 +16,7 @@ """Tests for Flask plugin module exports and integration types.""" -from genkit.core.context import RequestData +from genkit.plugins.flask.handler import RequestData class TestFlaskModuleExports: diff --git a/py/plugins/flask/tests/flask_handler_test.py b/py/plugins/flask/tests/flask_handler_test.py index 89329d2b33..b68610d15c 100644 --- a/py/plugins/flask/tests/flask_handler_test.py +++ b/py/plugins/flask/tests/flask_handler_test.py @@ -18,7 +18,7 @@ import pytest -from genkit.core.error import GenkitError +from genkit._core._error import GenkitError from genkit.plugins.flask.handler import genkit_flask_handler @@ -26,7 +26,7 @@ class TestGenkitFlaskHandlerValidation: """Tests that genkit_flask_handler rejects non-flow inputs.""" def test_rejects_plain_function(self) -> None: - """The decorator must reject arguments that are not FlowWrapper.""" + """The decorator must reject arguments that are not Flow.""" class FakeGenkit: pass diff --git a/py/plugins/flask/tests/flask_test.py b/py/plugins/flask/tests/flask_test.py index 99f0fd841b..afadbd5599 100644 --- a/py/plugins/flask/tests/flask_test.py +++ b/py/plugins/flask/tests/flask_test.py @@ -21,8 +21,8 @@ from flask import Flask, Request -from genkit.ai import ActionRunContext, Genkit -from genkit.core.context import RequestData +from genkit import ActionRunContext, Genkit +from genkit.plugin_api import RequestData from genkit.plugins.flask import genkit_flask_handler diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/action.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/action.py index a88684c72c..efb6835c75 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/action.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/action.py @@ -37,11 +37,12 @@ import structlog from opentelemetry.sdk.trace import ReadableSpan +from genkit.plugin_api import to_display_path + from .gcp_logger import gcp_logger from .utils import ( create_common_log_attributes, extract_outer_feature_name_from_path, - to_display_path, truncate, truncate_path, ) diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/config.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/config.py index 5b7181b4ca..f8921505a5 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/config.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/config.py @@ -36,9 +36,7 @@ from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import get_current_span, span as trace_span -from genkit.core.environment import is_dev_environment -from genkit.core.logging import get_logger -from genkit.core.tracing import add_custom_exporter +from genkit.plugin_api import add_custom_exporter, is_dev_environment from .constants import ( DEFAULT_METRIC_EXPORT_INTERVAL_MS, @@ -50,7 +48,7 @@ from .metrics_exporter import GenkitMetricExporter from .trace_exporter import GcpAdjustingTraceExporter, GenkitGCPExporter -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) def resolve_project_id( diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/engagement.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/engagement.py index c56e5f3054..b6367bb03a 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/engagement.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/engagement.py @@ -40,7 +40,7 @@ from opentelemetry import metrics from opentelemetry.sdk.trace import ReadableSpan -from genkit.core import GENKIT_VERSION +from genkit.plugin_api import GENKIT_VERSION from .gcp_logger import gcp_logger from .utils import create_common_log_attributes, truncate diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/feature.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/feature.py index 0f51329ee6..f3a15328e9 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/feature.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/feature.py @@ -37,13 +37,12 @@ from opentelemetry import metrics from opentelemetry.sdk.trace import ReadableSpan -from genkit.core import GENKIT_VERSION +from genkit.plugin_api import GENKIT_VERSION, to_display_path from .gcp_logger import gcp_logger from .utils import ( create_common_log_attributes, extract_error_name, - to_display_path, truncate, truncate_path, ) diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/generate.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/generate.py index 2429fa25ac..f48f17342e 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/generate.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/generate.py @@ -93,14 +93,13 @@ from opentelemetry import metrics from opentelemetry.sdk.trace import ReadableSpan -from genkit.core import GENKIT_VERSION +from genkit.plugin_api import GENKIT_VERSION, to_display_path from .gcp_logger import gcp_logger from .utils import ( create_common_log_attributes, extract_error_name, extract_outer_feature_name_from_path, - to_display_path, truncate, truncate_path, ) @@ -425,10 +424,10 @@ def _record_generate_action_config_logs( metadata['threadName'] = thread_name config = input_data.get('config', {}) - if config.get('maxOutputTokens'): - metadata['maxOutputTokens'] = config['maxOutputTokens'] - if config.get('stopSequences'): - metadata['stopSequences'] = config['stopSequences'] + if config.get('max_output_tokens'): + metadata['maxOutputTokens'] = config['max_output_tokens'] + if config.get('stop_sequences'): + metadata['stopSequences'] = config['stop_sequences'] gcp_logger.log_structured(f'Config[{path}, {model}]', metadata) diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/path.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/path.py index b92a620610..54a00043c2 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/path.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/path.py @@ -37,7 +37,7 @@ from opentelemetry import metrics from opentelemetry.sdk.trace import ReadableSpan -from genkit.core import GENKIT_VERSION +from genkit.plugin_api import GENKIT_VERSION, to_display_path from .gcp_logger import gcp_logger from .utils import ( @@ -46,7 +46,6 @@ extract_error_name, extract_error_stack, extract_outer_feature_name_from_path, - to_display_path, truncate_path, ) diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/trace_exporter.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/trace_exporter.py index 81b5a202bb..8c4ae2c5ec 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/trace_exporter.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/trace_exporter.py @@ -29,7 +29,7 @@ from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult -from genkit.core.trace.adjusting_exporter import AdjustingTraceExporter, RedactedSpan +from genkit.plugin_api import AdjustingTraceExporter, RedactedSpan from .action import action_telemetry from .constants import ( diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/tracing.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/tracing.py index 617a2067f5..4086fb7ee6 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/tracing.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/tracing.py @@ -89,7 +89,7 @@ 1. **GcpAdjustingTraceExporter**: Extends AdjustingTraceExporter to add GCP-specific telemetry recording before spans are adjusted and exported. - 2. **AdjustingTraceExporter** (from genkit.core.trace): Base class that + 2. **AdjustingTraceExporter** (from genkit._core._trace): Base class that handles PII redaction, error marking, and label normalization. 3. **GenkitGCPExporter**: Extends CloudTraceSpanExporter with retry logic diff --git a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/utils.py b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/utils.py index 4e208c583b..b0e5104cd5 100644 --- a/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/utils.py +++ b/py/plugins/google-cloud/src/genkit/plugins/google_cloud/telemetry/utils.py @@ -22,7 +22,6 @@ Functions: - truncate(): Limit string length for log content - truncate_path(): Limit Genkit path string length - - to_display_path(): Convert internal path to display format - extract_outer_feature_name_from_path(): Get root feature from path - create_common_log_attributes(): Build log attributes dict - extract_error_*(): Error info extraction helpers @@ -188,30 +187,3 @@ def create_common_log_attributes(span: ReadableSpan, project_id: str | None = No 'logging.googleapis.com/trace': f'projects/{project_id}/traces/{format(span_context.trace_id, "032x")}', 'logging.googleapis.com/trace_sampled': '1' if is_sampled else '0', } - - -def to_display_path(qualified_path: str) -> str: - """Convert a qualified Genkit path to a display path. - - Simplifies paths like '/{myFlow,t:flow}/{step,t:flowStep}' to 'myFlow/step'. - - Args: - qualified_path: The full Genkit path. - - Returns: - A simplified display path. - """ - if not qualified_path: - return '' - - # Extract names from path segments like '{name,t:type}' - parts = [] - for segment in qualified_path.split('/'): - if segment.startswith('{'): - match = re.match(r'\{([^,}]+)', segment) - if match: - parts.append(match.group(1)) - elif segment: - parts.append(segment) - - return '/'.join(parts) diff --git a/py/plugins/google-cloud/tests/gcp_telemetry_utils_test.py b/py/plugins/google-cloud/tests/gcp_telemetry_utils_test.py index 81c8359eb0..42e08b6b4b 100644 --- a/py/plugins/google-cloud/tests/gcp_telemetry_utils_test.py +++ b/py/plugins/google-cloud/tests/gcp_telemetry_utils_test.py @@ -24,6 +24,7 @@ from opentelemetry.trace import TraceFlags +from genkit.plugin_api import to_display_path from genkit.plugins.google_cloud.telemetry.utils import ( MAX_LOG_CONTENT_CHARS, MAX_PATH_CHARS, @@ -33,7 +34,6 @@ extract_error_stack, extract_outer_feature_name_from_path, extract_outer_flow_name_from_path, - to_display_path, truncate, truncate_path, ) @@ -297,29 +297,27 @@ def test_none_context_returns_empty(self) -> None: # to_display_path() # --------------------------------------------------------------------------- class TestToDisplayPath: - """Tests for ToDisplayPath.""" + """Tests for to_display_path (now in genkit.plugin_api).""" def test_simple_flow_path(self) -> None: """Simple flow path.""" assert to_display_path('/{myFlow,t:flow}') == 'myFlow' def test_nested_path(self) -> None: - """Nested path.""" + """Nested path uses ' > ' separator (matching JS).""" result = to_display_path('/{myFlow,t:flow}/{step,t:flowStep}') - assert result == 'myFlow/step' + assert result == 'myFlow > step' def test_three_level_path(self) -> None: """Three level path.""" result = to_display_path('/{myFlow,t:flow}/{step,t:flowStep}/{googleai/gemini-pro,t:action,s:model}') - # The function extracts the name part before the first comma in braces, - # but nested slashes create extra segments. - assert 'myFlow' in result - assert 'step' in result + assert result == 'myFlow > step > googleai/gemini-pro' def test_empty_string_returns_empty(self) -> None: """Empty string returns empty.""" assert to_display_path('') == '' - def test_plain_segments_preserved(self) -> None: - """Plain segments preserved.""" - assert to_display_path('foo/bar') == 'foo/bar' + def test_plain_segments_not_matched(self) -> None: + """Plain segments without type annotations are not extracted.""" + # The regex only matches {name,t:type} patterns + assert to_display_path('foo/bar') == '' diff --git a/py/plugins/google-cloud/tests/tracing_test.py b/py/plugins/google-cloud/tests/tracing_test.py index 1bc80984c5..efc4d8ca03 100644 --- a/py/plugins/google-cloud/tests/tracing_test.py +++ b/py/plugins/google-cloud/tests/tracing_test.py @@ -31,14 +31,17 @@ from unittest import mock from unittest.mock import MagicMock, patch -from genkit.core.environment import EnvVar, GenkitEnvironment +# Environment variable and value constants (matching genkit._core._environment) +_GENKIT_ENV = 'GENKIT_ENV' +_ENV_DEV = 'dev' +_ENV_PROD = 'prod' def test_add_gcp_telemetry_wraps_with_gcp_adjusting_exporter() -> None: """Test that add_gcp_telemetry wraps the exporter with GcpAdjustingTraceExporter.""" # Set production environment and clear project-related env vars to ensure project_id is None with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}, clear=False), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}, clear=False), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter') as mock_gcp_exporter, patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter') as mock_adjusting, patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter') as mock_add_exporter, @@ -81,7 +84,7 @@ def test_add_gcp_telemetry_wraps_with_gcp_adjusting_exporter() -> None: def test_add_gcp_telemetry_with_log_input_and_output_enabled() -> None: """Test that log_input_and_output=True disables PII redaction (JS parity).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter'), patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter') as mock_adjusting, patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter'), @@ -104,7 +107,7 @@ def test_add_gcp_telemetry_with_log_input_and_output_enabled() -> None: def test_add_gcp_telemetry_with_project_id() -> None: """Test that project_id is passed to GcpAdjustingTraceExporter (JS/Go parity).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter'), patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter') as mock_adjusting, patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter'), @@ -127,7 +130,7 @@ def test_add_gcp_telemetry_with_project_id() -> None: def test_add_gcp_telemetry_skips_in_dev_without_force() -> None: """Test that telemetry is skipped in dev environment without force_dev_export (JS/Go parity).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_DEV}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter') as mock_gcp_exporter, patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter') as mock_add_exporter, ): @@ -144,7 +147,7 @@ def test_add_gcp_telemetry_skips_in_dev_without_force() -> None: def test_add_gcp_telemetry_exports_in_dev_with_force() -> None: """Test that telemetry is exported in dev environment with force_dev_export=True (JS/Go parity).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_DEV}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter') as mock_gcp_exporter, patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter'), patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter') as mock_add_exporter, @@ -167,7 +170,7 @@ def test_add_gcp_telemetry_exports_in_dev_with_force() -> None: def test_add_gcp_telemetry_disable_traces() -> None: """Test that disable_traces=True skips trace export (JS/Go parity).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter') as mock_gcp_exporter, patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter') as mock_add_exporter, patch('genkit.plugins.google_cloud.telemetry.config.GoogleCloudResourceDetector'), @@ -189,7 +192,7 @@ def test_add_gcp_telemetry_disable_traces() -> None: def test_add_gcp_telemetry_disable_metrics() -> None: """Test that disable_metrics=True skips metrics export (JS/Go parity).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter'), patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter'), patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter'), @@ -214,7 +217,7 @@ def test_add_gcp_telemetry_disable_metrics() -> None: def test_add_gcp_telemetry_custom_metric_interval() -> None: """Test that metric_export_interval_ms is passed correctly (JS/Go parity).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter'), patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter'), patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter'), @@ -239,7 +242,7 @@ def test_add_gcp_telemetry_custom_metric_interval() -> None: def test_add_gcp_telemetry_enforces_minimum_interval() -> None: """Test that metric_export_interval_ms enforces minimum 5000ms (GCP requirement).""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter'), patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter'), patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter'), @@ -316,7 +319,7 @@ def test_resolve_project_id_from_credentials() -> None: def test_legacy_force_export_parameter() -> None: """Test that legacy force_export parameter still works but shows warning.""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.DEV}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_DEV}), patch('genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter') as mock_gcp_exporter, patch('genkit.plugins.google_cloud.telemetry.config.GcpAdjustingTraceExporter'), patch('genkit.plugins.google_cloud.telemetry.config.add_custom_exporter'), @@ -344,7 +347,7 @@ def test_legacy_force_export_parameter() -> None: def test_add_gcp_telemetry_is_fail_safe() -> None: """Test that add_gcp_telemetry does not crash if initialization fails.""" with ( - mock.patch.dict(os.environ, {EnvVar.GENKIT_ENV: GenkitEnvironment.PROD}), + mock.patch.dict(os.environ, {_GENKIT_ENV: _ENV_PROD}), patch( 'genkit.plugins.google_cloud.telemetry.config.GenkitGCPExporter', side_effect=Exception('Auth failed'), diff --git a/py/plugins/google-genai/README.md b/py/plugins/google-genai/README.md index 24448fd76a..5d817b0246 100644 --- a/py/plugins/google-genai/README.md +++ b/py/plugins/google-genai/README.md @@ -80,7 +80,7 @@ Built-in evaluators for assessing model output quality. Evaluators are automatic ```python from genkit import Genkit -from genkit.core.typing import BaseDataPoint +from genkit._core.typing import BaseDataPoint from genkit.plugins.google_genai import VertexAI ai = Genkit(plugins=[VertexAI(project='my-project')]) diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/evaluators/evaluation.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/evaluators/evaluation.py index 868a0895db..1760856414 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/evaluators/evaluation.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/evaluators/evaluation.py @@ -70,15 +70,12 @@ from google.auth.transport.requests import Request from pydantic import BaseModel, ConfigDict -from genkit.ai import GENKIT_CLIENT_HEADER -from genkit.blocks.evaluator import EvalFnResponse -from genkit.core.action import Action -from genkit.core.error import GenkitError -from genkit.core.http_client import get_cached_client -from genkit.core.typing import BaseDataPoint, Details, Score +from genkit import GenkitError +from genkit.evaluator import BaseDataPoint, Details, EvalFnResponse, Score +from genkit.plugin_api import GENKIT_CLIENT_HEADER, Action, get_cached_client if TYPE_CHECKING: - from genkit.ai._registry import GenkitRegistry + from genkit import Genkit as GenkitRegistry class VertexAIEvaluationMetricType(StrEnum): diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/google.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/google.py index 72bf0073c6..b89e331afc 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/google.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/google.py @@ -102,24 +102,17 @@ from google.genai.types import HttpOptions, HttpOptionsDict import genkit.plugins.google_genai.constants as const -from genkit.ai import GENKIT_CLIENT_HEADER, Plugin -from genkit.blocks.background_model import BackgroundAction -from genkit.blocks.document import Document -from genkit.blocks.embedding import EmbedderOptions, EmbedderSupports, embedder_action_metadata -from genkit.blocks.model import model_action_metadata -from genkit.blocks.reranker import reranker_action_metadata -from genkit.core._loop_local import _loop_local_client -from genkit.core.action import Action, ActionMetadata -from genkit.core.error import GenkitError -from genkit.core.registry import ActionKind -from genkit.core.schema import to_json_schema -from genkit.core.typing import ( - EvalFnResponse, - EvalRequest, - RankedDocumentData, - RankedDocumentMetadata, - RerankerRequest, - RerankerResponse, +from genkit.embedder import EmbedderOptions, EmbedderSupports, embedder_action_metadata +from genkit.evaluator import EvalFnResponse, EvalRequest +from genkit.model import BackgroundAction, model_action_metadata +from genkit.plugin_api import ( + GENKIT_CLIENT_HEADER, + Action, + ActionKind, + ActionMetadata, + Plugin, + loop_local_client, + to_json_schema, ) from genkit.plugins.google_genai.evaluators import ( VertexAIEvaluationMetricType, @@ -147,15 +140,6 @@ is_veo_model, veo_model_info, ) -from genkit.plugins.google_genai.rerankers.reranker import ( - KNOWN_MODELS as RERANKER_MODELS, - RerankRequest, - VertexRerankerClientOptions, - VertexRerankerConfig, - _from_rerank_response, - _to_reranker_doc, - reranker_rank, -) class GenaiModels: @@ -416,7 +400,7 @@ def __init__( 'http_options': _inject_attribution_headers(http_options, base_url, api_version), } # Single loop-local client accessor used everywhere in plugin runtime paths. - self._runtime_client = _loop_local_client(lambda: genai.client.Client(**self._client_kwargs)) + self._runtime_client = loop_local_client(lambda: genai.client.Client(**self._client_kwargs)) async def init(self) -> list[Action]: """Initialize the plugin. @@ -521,7 +505,7 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: return self._resolve_embedder(name) return None - def _resolve_veo_model(self, name: str) -> 'BackgroundAction': + def _resolve_veo_model(self, name: str) -> BackgroundAction: """Create a BackgroundAction for a Veo video generation model. Args: @@ -774,7 +758,7 @@ def __init__( 'http_options': _inject_attribution_headers(http_options, base_url, api_version), } # Single loop-local client accessor used everywhere in plugin runtime paths. - self._runtime_client = _loop_local_client(lambda: genai.client.Client(**self._client_kwargs)) + self._runtime_client = loop_local_client(lambda: genai.client.Client(**self._client_kwargs)) async def init(self) -> list[Action]: """Initialize the plugin. @@ -797,20 +781,16 @@ async def init(self) -> list[Action]: for name in genai_models.embedders: actions.append(self._resolve_embedder(vertexai_name(name))) - # Register Vertex AI rerankers - for name in RERANKER_MODELS: - actions.append(self._resolve_reranker(vertexai_name(name))) - # Register Vertex AI evaluators # Deferred import to avoid circular dependency - from genkit.ai._registry import GenkitRegistry + from genkit import Genkit if not self._project: raise ValueError( 'VertexAI plugin requires a project ID to use evaluators. ' 'Set the project parameter or GOOGLE_CLOUD_PROJECT environment variable.' ) - registry = GenkitRegistry() + registry = Genkit() actions.extend( create_vertex_evaluators( registry, @@ -856,8 +836,6 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: return self._resolve_model(name) elif action_type == ActionKind.EMBEDDER: return self._resolve_embedder(name) - elif action_type == ActionKind.RERANKER: - return self._resolve_reranker(name) elif action_type == ActionKind.EVALUATOR: return self._resolve_evaluator(name) return None @@ -879,9 +857,9 @@ def _resolve_evaluator(self, name: str) -> Action | None: except ValueError: return None - from genkit.ai._registry import GenkitRegistry + from genkit import Genkit - registry = GenkitRegistry() + registry = Genkit() if not self._project: raise ValueError( 'VertexAI plugin requires a project ID to use evaluators. ' @@ -952,84 +930,6 @@ def _resolve_embedder(self, name: str) -> Action: """ return _create_embedder_action(name, self._runtime_client, VERTEXAI_PLUGIN_NAME) - def _resolve_reranker(self, name: str) -> Action: - """Create an Action object for a Vertex AI reranker. - - Args: - name: The namespaced name of the reranker. - - Returns: - Action object for the reranker. - """ - # Extract local name (remove plugin prefix) - clean_name = name.replace(VERTEXAI_PLUGIN_NAME + '/', '') if name.startswith(VERTEXAI_PLUGIN_NAME) else name - - # Validate project is configured (required for reranker API) - if not self._project: - raise ValueError( - 'VertexAI plugin requires a project ID to use rerankers. ' - 'Set the project parameter or GOOGLE_CLOUD_PROJECT environment variable.' - ) - - # Use project and location stored on the plugin instance during init. - # This avoids accessing private attributes of the client library. - client_options = VertexRerankerClientOptions( - project_id=self._project, - location=self._location, - ) - - async def wrapper( - request: RerankerRequest, - _ctx: Any, # noqa: ANN401 - ) -> RerankerResponse: - """Wrapper that takes RerankerRequest and returns RerankerResponse. - - This matches the signature expected by the Action class (max 2 args). - """ - query_doc = Document.from_document_data(request.query) - documents = [Document.from_document_data(d) for d in request.documents] - options = request.options - - config = VertexRerankerConfig.model_validate(options or {}) - - # Use location from config if provided, otherwise use client default - effective_options = VertexRerankerClientOptions( - project_id=client_options.project_id, - location=config.location or client_options.location, - ) - - query_text = query_doc.text() - if not query_text: - raise GenkitError(message='Reranker query cannot be empty.') - - rerank_request = RerankRequest( - model=clean_name, - query=query_text, - records=[_to_reranker_doc(doc, idx) for idx, doc in enumerate(documents)], - top_n=config.top_n, - ignore_record_details_in_response=config.ignore_record_details_in_response, - ) - - response = await reranker_rank(clean_name, rerank_request, effective_options) - ranked_docs = _from_rerank_response(response, documents) - - # Convert to RerankerResponse format - ranked_docs are RankedDocument instances - response_docs: list[RankedDocumentData] = [] - for doc in ranked_docs: - metadata = RankedDocumentMetadata(score=doc.score if doc.score is not None else 0.0) - response_docs.append(RankedDocumentData(content=doc.content, metadata=metadata)) - - return RerankerResponse(documents=response_docs) - - metadata = reranker_action_metadata(name) - - return Action( - kind=ActionKind.RERANKER, - name=name, - fn=wrapper, - metadata=metadata.metadata, - ) - async def list_actions(self) -> list[ActionMetadata]: """Generate a list of available actions or models. diff --git a/py/packages/genkit/src/genkit/lang/deprecations.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/_deprecations.py similarity index 71% rename from py/packages/genkit/src/genkit/lang/deprecations.py rename to py/plugins/google-genai/src/genkit/plugins/google_genai/models/_deprecations.py index 05fff39660..72ceaa7d28 100644 --- a/py/packages/genkit/src/genkit/lang/deprecations.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/_deprecations.py @@ -14,33 +14,17 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Helpers for managing deprecations. +"""Helpers for managing deprecations in enum members. -This module provides a metaclass for creating deprecated enums. It allows -developers to mark enum members as deprecated and provides warnings when those -members are accessed. - -The metaclass works by overriding the __getattribute__ method of the enum class. -When a deprecated member is accessed, a deprecation warning is issued. - -## Example - - ```python - Deprecations = deprecated_enum_metafactory({ - 'OLD_THING': DeprecationInfo(recommendation='NEW_THING', status=DeprecationStatus.DEPRECATED), - }) - - - class TestEnum(StrEnum, metaclass=Deprecations): ... - ``` +This module provides utilities for handling deprecated enum values, +allowing plugin authors to gracefully deprecate enum members while +maintaining backward compatibility. """ import enum import warnings from dataclasses import dataclass -from genkit.core._compat import override - class DeprecationStatus(enum.Enum): """Defines the deprecation status of an enum member.""" @@ -71,7 +55,6 @@ def deprecated_enum_metafactory( """ class DeprecatedEnumMeta(enum.EnumMeta): - @override def __getattribute__(cls, name: str) -> object: """Get an attribute of the enum class. @@ -82,8 +65,6 @@ def __getattribute__(cls, name: str) -> object: Returns: The attribute value. """ - # This __getattribute__ is called when accessing attributes - # directly on the Enum class itself (e.g., MyEnum.MEMBER). if name in deprecated_map: info = deprecated_map[name] if info.status in ( @@ -96,9 +77,6 @@ def __getattribute__(cls, name: str) -> object: if info.recommendation is not None else f'{cls.__name__}.{name} is {status_str}' ) - # Start with stacklevel=4; adjust if needed based on test - # results (factory adds a frame, metaclass __getattribute__ - # adds a frame) warnings.warn( message, DeprecationWarning, diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/context_caching/utils.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/context_caching/utils.py index 4afdedef9e..fd630aece2 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/context_caching/utils.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/context_caching/utils.py @@ -21,8 +21,7 @@ import structlog -from genkit.core.error import GenkitError -from genkit.core.typing import GenerateRequest +from genkit import GenkitError, ModelRequest from genkit.plugins.google_genai.models.context_caching.constants import ( CONTEXT_CACHE_SUPPORTED_MODELS, INVALID_ARGUMENT_MESSAGES, @@ -31,11 +30,11 @@ logger = structlog.getLogger(__name__) -def generate_cache_key(request: GenerateRequest) -> str: +def generate_cache_key(request: ModelRequest) -> str: """Generates context cache key by hashing the given request instance. Args: - request: `GenerateRequest` instance to hash + request: `ModelRequest` instance to hash Returns: Generated cache key string @@ -43,11 +42,11 @@ def generate_cache_key(request: GenerateRequest) -> str: return hashlib.sha256(json.dumps(request.model_dump(), sort_keys=True).encode()).hexdigest() -def validate_context_cache_request(request: GenerateRequest, model_name: str) -> bool: +def validate_context_cache_request(request: ModelRequest, model_name: str) -> bool: """Verifies that the context cache request could be processed for the request. Args: - request: `GenerateRequest` instance to check + request: `ModelRequest` instance to check model_name: Name of the generation model to check Returns: diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/embedder.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/embedder.py index b5addb7847..1c601a5d60 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/embedder.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/embedder.py @@ -27,8 +27,8 @@ from google import genai from google.genai import types as genai_types +from genkit import Embedding, EmbedRequest, EmbedResponse from genkit.plugins.google_genai.models.utils import PartConverter -from genkit.types import Embedding, EmbedRequest, EmbedResponse class VertexEmbeddingModels(StrEnum): diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/gemini.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/gemini.py index 39970c7bf3..1fbda9417e 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/gemini.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/gemini.py @@ -142,38 +142,25 @@ else: from enum import StrEnum +import json from functools import cached_property -from typing import Annotated, Any, cast +from typing import Annotated, Any, Any as JsonAny, cast from google import genai from google.genai import types as genai_types from google.genai.errors import ClientError from pydantic import BaseModel, ConfigDict, Field, WithJsonSchema -from genkit.ai import ( - ActionRunContext, -) -from genkit.blocks.model import get_basic_usage_stats -from genkit.codec import dump_dict, dump_json -from genkit.core.error import GenkitError, StatusName -from genkit.core.tracing import tracer -from genkit.core.typing import ( - Candidate, - FinishReason, -) -from genkit.lang.deprecations import ( - deprecated_enum_metafactory, -) -from genkit.plugins.google_genai.models.utils import PartConverter -from genkit.types import ( +from genkit import ( Constrained, - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationCommonConfig, - GenerationUsage, + GenkitError, Message, + ModelConfig, ModelInfo, + ModelRequest, + ModelResponse, + ModelResponseChunk, + ModelUsage, Part, Role, Stage, @@ -181,6 +168,24 @@ TextPart, ToolDefinition, ) +from genkit._core._typing import GenerationCommonConfig +from genkit.model import Candidate, FinishReason, get_basic_usage_stats +from genkit.plugin_api import ( + ActionRunContext, + StatusName, + tracer, +) + + +def _to_dict(obj: JsonAny) -> JsonAny: # noqa: ANN401 + """Convert object to dict if it's a Pydantic model, otherwise return as-is.""" + return obj.model_dump() if isinstance(obj, BaseModel) else obj + + +from genkit.plugins.google_genai.models._deprecations import ( # noqa: E402 + deprecated_enum_metafactory, +) +from genkit.plugins.google_genai.models.utils import PartConverter # noqa: E402 class HarmCategory(StrEnum): @@ -299,7 +304,7 @@ class VoiceConfigSchema(BaseModel): prebuilt_voice_config: PrebuiltVoiceConfig | None = Field(None, alias='prebuiltVoiceConfig') -class GeminiConfigSchema(GenerationCommonConfig): +class GeminiConfigSchema(ModelConfig): """Gemini Config Schema.""" model_config = ConfigDict(extra='allow', populate_by_name=True) @@ -419,7 +424,7 @@ class GeminiConfigSchema(GenerationCommonConfig): None, description='Return grounding metadata from links included in the query', alias='urlContext' ) - # inherited from GenerationCommonConfig: + # inherited from ModelConfig: # version, temperature, max_output_tokens, top_k, top_p, stop_sequences temperature: Annotated[ @@ -767,7 +772,7 @@ class GemmaConfigSchema(GeminiConfigSchema): Deprecations = deprecated_enum_metafactory({}) -class VertexAIGeminiVersion(StrEnum, metaclass=Deprecations): +class VertexAIGeminiVersion(StrEnum, metaclass=Deprecations): # pyrefly: ignore[invalid-inheritance] """VertexAIGemini models. Model Support: @@ -827,7 +832,7 @@ class VertexAIGeminiVersion(StrEnum, metaclass=Deprecations): GEMMA_3N_E4B_IT = 'gemma-3n-e4b-it' -class GoogleAIGeminiVersion(StrEnum, metaclass=Deprecations): +class GoogleAIGeminiVersion(StrEnum, metaclass=Deprecations): # pyrefly: ignore[invalid-inheritance] """GoogleAI Gemini models. Model Support: @@ -1047,7 +1052,7 @@ def __init__( self._version = version self._client = client - def _get_tools(self, request: GenerateRequest) -> list[genai_types.Tool]: + def _get_tools(self, request: ModelRequest) -> list[genai_types.Tool]: """Generates VertexAI Gemini compatible tool definitions. Args: @@ -1165,7 +1170,7 @@ def _convert_schema_property( return schema async def _retrieve_cached_content( - self, request: GenerateRequest, model_name: str, cache_config: dict, contents: list[genai_types.Content] + self, request: ModelRequest, model_name: str, cache_config: dict, contents: list[genai_types.Content] ) -> genai_types.CachedContent: """Retrieves cached content from the Google API if exists. @@ -1211,7 +1216,7 @@ async def _retrieve_cached_content( ) return cache - async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: """Handle a generation request. Args: @@ -1259,9 +1264,9 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext) -> Gen if request.config: if isinstance(request.config, dict): - api_version = request.config.get('api_version') or request.config.get('apiVersion') - api_key_override = request.config.get('api_key') or request.config.get('apiKey') - base_url_override = request.config.get('base_url') or request.config.get('baseUrl') + api_version = request.config.get('api_version') + api_key_override = request.config.get('api_key') + base_url_override = request.config.get('base_url') else: api_version = getattr(request.config, 'api_version', None) api_key_override = getattr(request.config, 'api_key', None) @@ -1342,7 +1347,7 @@ async def _generate( request_cfg: genai_types.GenerateContentConfig | None, model_name: str, client: genai.Client | None = None, - ) -> GenerateResponse: + ) -> ModelResponse: """Call google-genai generate. Args: @@ -1357,13 +1362,13 @@ async def _generate( with tracer.start_as_current_span('generate_content') as span: span.set_attribute( 'genkit:input', - dump_json( + json.dumps( { - 'config': dump_dict(request_cfg), - 'contents': [dump_dict(c) for c in request_contents], + 'config': _to_dict(request_cfg), + 'contents': [_to_dict(c) for c in request_contents], 'model': model_name, }, - fallback=lambda _: '[!! failed to serialize !!]', + default=lambda _: '[!! failed to serialize !!]', ), ) client = client or self._client @@ -1402,7 +1407,7 @@ async def _generate( status='INTERNAL', message=f'Unexpected error during generation: {type(e).__name__}: {str(e)}', ) from e - span.set_attribute('genkit:output', dump_json(response)) + span.set_attribute('genkit:output', json.dumps(_to_dict(response), default=str)) content = await self._contents_from_response(response) @@ -1447,14 +1452,14 @@ async def _generate( ) ) - return GenerateResponse( + return ModelResponse( message=Message( content=content, role=Role.MODEL, ), finish_reason=finish_reason, candidates=candidates, - usage=GenerationUsage( + usage=ModelUsage( input_tokens=float(response.usage_metadata.prompt_token_count or 0) if response.usage_metadata else None, @@ -1472,7 +1477,7 @@ async def _streaming_generate( ctx: ActionRunContext, model_name: str, client: genai.Client | None = None, - ) -> GenerateResponse: + ) -> ModelResponse: """Call google-genai generate for streaming. Args: @@ -1488,11 +1493,14 @@ async def _streaming_generate( with tracer.start_as_current_span('generate_content_stream') as span: span.set_attribute( 'genkit:input', - dump_json({ - 'config': dump_dict(request_cfg), - 'contents': [dump_dict(c) for c in request_contents], - 'model': model_name, - }), + json.dumps( + { + 'config': _to_dict(request_cfg), + 'contents': [_to_dict(c) for c in request_contents], + 'model': model_name, + }, + default=str, + ), ) client = client or self._client try: @@ -1525,13 +1533,13 @@ async def _streaming_generate( if content: # Only process if we have content accumulated_content.extend(content) ctx.send_chunk( - chunk=GenerateResponseChunk( + chunk=ModelResponseChunk( content=content, role=Role.MODEL, ) ) - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=accumulated_content, @@ -1558,7 +1566,7 @@ def metadata(self) -> dict: } async def _build_messages( - self, request: GenerateRequest, model_name: str + self, request: ModelRequest, model_name: str ) -> tuple[list[genai_types.Content], genai_types.CachedContent | None]: """Build google-genai request contents from Genkit request. @@ -1618,8 +1626,8 @@ async def _contents_from_response(self, response: genai_types.GenerateContentRes # Ensure we always return a list, even if empty return content if content else [] - async def _genkit_to_googleai_cfg(self, request: GenerateRequest) -> genai_types.GenerateContentConfig | None: - """Converts a Genkit GenerateRequest to a Gemini GenerateContentConfig. + async def _genkit_to_googleai_cfg(self, request: ModelRequest) -> genai_types.GenerateContentConfig | None: + """Converts a Genkit ModelRequest to a Gemini GenerateContentConfig. The conversion follows a linear pipeline: 1. Extract system instructions from messages @@ -1664,18 +1672,20 @@ async def _genkit_to_googleai_cfg(self, request: GenerateRequest) -> genai_types # Tools from top-level field and config-level fields tools.extend(self._get_tools(request)) - if cfg is not None or tools or system_instruction or request.output: + has_output = bool(request.output_format or request.output_schema) + + if cfg is not None or tools or system_instruction or request.output_format: if cfg is None: cfg = genai_types.GenerateContentConfig() - if request.output: + if has_output: response_mime_type = ( - 'application/json' if request.output.format == 'json' and not request.tools else None + 'application/json' if request.output_format == 'json' and not request.tools else None ) cfg.response_mime_type = response_mime_type - if request.output.schema and request.output.constrained: - cfg.response_schema = self._convert_schema_property(request.output.schema) + if request.output_schema and request.output_constrained: + cfg.response_schema = self._convert_schema_property(request.output_schema) if tools: cfg.tools = cast(genai_types.ToolListUnion, tools) @@ -1696,13 +1706,13 @@ async def _genkit_to_googleai_cfg(self, request: GenerateRequest) -> genai_types def _normalize_config_to_dict( self, - config: GeminiConfigSchema | GenerationCommonConfig | dict, + config: GeminiConfigSchema | ModelConfig | dict, ) -> dict[str, Any] | None: """Normalize any config type into a plain dict for uniform processing. Handles three input shapes: - GeminiConfigSchema (and subclasses like TTS/Image): model_dump - - GenerationCommonConfig: model_dump + - ModelConfig: model_dump - dict: route to the appropriate schema first, then model_dump Returns: @@ -1711,7 +1721,7 @@ def _normalize_config_to_dict( """ if isinstance(config, GeminiConfigSchema): schema = config - elif isinstance(config, GenerationCommonConfig): + elif isinstance(config, (ModelConfig, GenerationCommonConfig)): schema = config elif isinstance(config, dict): if 'image_config' in config: @@ -1723,7 +1733,7 @@ def _normalize_config_to_dict( else: return None - dumped = schema.model_dump(exclude_none=True) + dumped = schema.model_dump(exclude_none=True, by_alias=False) return dumped or None def _extract_tools_from_config( @@ -1784,7 +1794,7 @@ def _clean_unsupported_keys(self, config: dict[str, Any]) -> None: if key in config and key not in genai_types.GenerateContentConfig.model_fields: del config[key] - def _create_usage_stats(self, request: GenerateRequest, response: GenerateResponse) -> GenerationUsage: + def _create_usage_stats(self, request: ModelRequest, response: ModelResponse) -> ModelUsage: """Create usage statistics. Args: @@ -1795,7 +1805,7 @@ def _create_usage_stats(self, request: GenerateRequest, response: GenerateRespon usage statistics """ if not response.message: - usage = GenerationUsage() + usage = ModelUsage() usage.input_tokens = 0 usage.output_tokens = 0 usage.total_tokens = 0 diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/imagen.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/imagen.py index 6e6cd0095c..4bfcfb4c44 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/imagen.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/imagen.py @@ -24,27 +24,32 @@ else: from enum import StrEnum +import json from functools import cached_property +from typing import Any from google import genai from google.genai import types as genai_types from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError -from genkit.ai import ActionRunContext -from genkit.codec import dump_dict, dump_json -from genkit.core.tracing import tracer -from genkit.types import ( - GenerateRequest, - GenerateResponse, +from genkit import ( Media, MediaPart, Message, ModelInfo, + ModelRequest, + ModelResponse, Part, Role, Supports, TextPart, ) +from genkit.plugin_api import ActionRunContext, tracer + + +def _to_dict(obj: Any) -> Any: # noqa: ANN401 + """Convert object to dict if it's a Pydantic model, otherwise return as-is.""" + return obj.model_dump() if isinstance(obj, BaseModel) else obj class ImagenVersion(StrEnum): @@ -136,7 +141,7 @@ def __init__(self, version: str | ImagenVersion, client: genai.Client) -> None: self._version = version self._client = client - def _build_prompt(self, request: GenerateRequest) -> str: + def _build_prompt(self, request: ModelRequest) -> str: """Build prompt request from Genkit request. Args: @@ -154,7 +159,7 @@ def _build_prompt(self, request: GenerateRequest) -> str: raise ValueError('Non-text messages are not supported') return ' '.join(prompt) - async def generate(self, request: GenerateRequest, _: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, _: ActionRunContext) -> ModelResponse: """Handle a generation request. Args: @@ -172,25 +177,25 @@ async def generate(self, request: GenerateRequest, _: ActionRunContext) -> Gener with tracer.start_as_current_span('generate_images') as span: span.set_attribute( 'genkit:input', - dump_json({ - 'config': dump_dict(config), + json.dumps({ + 'config': _to_dict(config), 'contents': prompt, 'model': self._version, }), ) response = await self._client.aio.models.generate_images(model=self._version, prompt=prompt, config=config) - span.set_attribute('genkit:output', dump_json(response)) + span.set_attribute('genkit:output', json.dumps(_to_dict(response), default=str)) content = self._contents_from_response(response) - return GenerateResponse( + return ModelResponse( message=Message( content=content, role=Role.MODEL, ) ) - def _get_config(self, request: GenerateRequest) -> genai_types.GenerateImagesConfigOrDict | None: + def _get_config(self, request: ModelRequest) -> genai_types.GenerateImagesConfigOrDict | None: cfg = None if request.config: diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/lyria.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/lyria.py index 4b11e4c1ab..c1023476b0 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/lyria.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/lyria.py @@ -73,10 +73,7 @@ from pydantic import BaseModel, Field -from genkit.core.typing import ( - ModelInfo, - Supports, -) +from genkit import ModelInfo, Supports class LyriaVersion(StrEnum): @@ -152,7 +149,7 @@ def _extract_text(messages: list[Any]) -> str: """Extract text prompt from messages. Args: - messages: The message list from a GenerateRequest. + messages: The message list from a ModelRequest. Returns: The text prompt string. diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/utils.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/utils.py index 50ce68ad2c..76c636b377 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/utils.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/utils.py @@ -52,12 +52,12 @@ from google import genai -from genkit.core.http_client import get_cached_client -from genkit.core.typing import DocumentPart, Metadata -from genkit.types import ( +from genkit import ( CustomPart, + DocumentPart, Media, MediaPart, + Metadata, Part, ReasoningPart, TextPart, @@ -66,6 +66,7 @@ ToolResponse, ToolResponsePart, ) +from genkit.plugin_api import get_cached_client class PartConverter: diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/veo.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/veo.py index 31ff24573d..4f206b8e09 100644 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/models/veo.py +++ b/py/plugins/google-genai/src/genkit/plugins/google_genai/models/veo.py @@ -79,24 +79,20 @@ from google.genai import types as genai_types from pydantic import BaseModel, ConfigDict, Field -from genkit.core.action import ActionRunContext -from genkit.core.tracing import tracer -from genkit.core.typing import ( - Error, - GenerateRequest, - GenerateResponse, - ModelInfo, - Operation, - Supports, -) -from genkit.types import ( +from genkit import ( Media, MediaPart, Message, + ModelInfo, + ModelRequest, + ModelResponse, Part, Role, + Supports, TextPart, ) +from genkit.model import Error, Operation +from genkit.plugin_api import ActionRunContext, tracer class VeoVersion(StrEnum): @@ -169,8 +165,8 @@ def veo_model_info(version: str) -> ModelInfo: ) -def _extract_text(request: GenerateRequest) -> str: - """Extract text prompt from a GenerateRequest. +def _extract_text(request: ModelRequest) -> str: + """Extract text prompt from a ModelRequest. Args: request: The generation request. @@ -287,7 +283,7 @@ def __init__(self, version: str, client: genai.Client) -> None: self._version = version self._client = client - def _build_prompt(self, request: GenerateRequest) -> str: + def _build_prompt(self, request: ModelRequest) -> str: """Build prompt request from Genkit request.""" prompt = [] for message in request.messages: @@ -300,7 +296,7 @@ def _build_prompt(self, request: GenerateRequest) -> str: pass return ' '.join(prompt) - async def generate(self, request: GenerateRequest, _: ActionRunContext) -> GenerateResponse: + async def generate(self, request: ModelRequest, _: ActionRunContext) -> ModelResponse: """Handle a generation request (synchronous/blocking mode for Vertex AI). Args: @@ -333,14 +329,14 @@ async def generate(self, request: GenerateRequest, _: ActionRunContext) -> Gener content = self._contents_from_response(cast(genai_types.GenerateVideosResponse, response)) - return GenerateResponse( + return ModelResponse( message=Message( content=content, role=Role.MODEL, ) ) - async def start(self, request: GenerateRequest, ctx: ActionRunContext) -> Operation: + async def start(self, request: ModelRequest, ctx: ActionRunContext) -> Operation: """Start a video generation operation (background model pattern for GoogleAI). Args: @@ -400,12 +396,10 @@ async def check(self, operation: Operation) -> Operation: return _from_veo_operation(op_dict) - def _get_config(self, request: GenerateRequest) -> genai_types.GenerateVideosConfigOrDict | None: - cfg = None - if request.config: - # Simple cast/validate - cfg = request.config - return cfg + def _get_config(self, request: ModelRequest) -> genai_types.GenerateVideosConfigOrDict | None: + if not request.config: + return None + return cast(genai_types.GenerateVideosConfigOrDict, request.config) def _contents_from_response(self, response: genai_types.GenerateVideosResponse) -> list[Part]: content = [] diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/rerankers/__init__.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/rerankers/__init__.py deleted file mode 100644 index dd14a8f058..0000000000 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/rerankers/__init__.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Vertex AI Rerankers for the Genkit framework. - -This module provides reranking functionality using the Vertex AI Discovery Engine -Ranking API. Rerankers improve RAG (Retrieval-Augmented Generation) quality by -re-scoring documents based on their relevance to a query. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Reranker β”‚ A "second opinion" scorer that re-orders your β”‚ - β”‚ β”‚ search results by relevance. Like asking an expert β”‚ - β”‚ β”‚ to sort your library books by importance. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Semantic Ranker β”‚ Uses AI to understand meaning, not just keywords. β”‚ - β”‚ β”‚ Knows "car" and "automobile" mean the same thing. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ top_n β”‚ How many top results to return after reranking. β”‚ - β”‚ β”‚ "Give me the 5 most relevant documents." β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Score β”‚ A number (0-1) showing how relevant a document is.β”‚ - β”‚ β”‚ Higher = more relevant to your query. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ RAG WITH RERANKING β”‚ - β”‚ β”‚ - β”‚ User Query: "How do neural networks learn?" β”‚ - β”‚ β”‚ β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Retriever β”‚ ◄── Fast initial search, returns ~100 docs β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ [doc1, doc2, doc3, ... doc100] β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Reranker β”‚ ◄── AI-powered relevance scoring β”‚ - β”‚ β”‚ (Vertex) β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ [doc47: 0.95, doc3: 0.87, doc12: 0.82, ...] β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Model β”‚ ◄── Uses top-k most relevant docs β”‚ - β”‚ β”‚ (Gemini) β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β–Ό β”‚ - β”‚ High-quality answer with accurate citations β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Overview: - Vertex AI offers semantic rerankers that use machine learning to score - documents based on their semantic similarity to a query. This is typically - used after initial retrieval to improve the quality of the top-k results. - - Reranking is a two-stage retrieval pattern: - 1. **Fast retrieval**: Get many candidates quickly (e.g., 100 docs) - 2. **Quality reranking**: Score candidates by relevance, keep top-k - -Available Models: - +--------------------------------+-----------------------------------------+ - | Model | Description | - +--------------------------------+-----------------------------------------+ - | semantic-ranker-default@latest | Latest default semantic ranker | - | semantic-ranker-default-004 | Semantic ranker version 004 | - | semantic-ranker-fast-004 | Fast variant (lower latency, less acc.) | - | semantic-ranker-default-003 | Semantic ranker version 003 | - | semantic-ranker-default-002 | Semantic ranker version 002 | - +--------------------------------+-----------------------------------------+ - -Example: - Basic reranking: - - >>> from genkit import Genkit - >>> from genkit.plugins.google_genai import VertexAI - >>> - >>> ai = Genkit(plugins=[VertexAI(project='my-project')]) - >>> - >>> # Rerank documents after retrieval - >>> ranked_docs = await ai.rerank( - ... reranker='vertexai/semantic-ranker-default@latest', - ... query='What is machine learning?', - ... documents=retrieved_docs, - ... options={'top_n': 5}, - ... ) - - Full RAG pipeline with reranking: - - >>> # 1. Retrieve initial candidates - >>> candidates = await ai.retrieve( - ... retriever='my-retriever', - ... query='How do neural networks learn?', - ... options={'limit': 50}, - ... ) - >>> - >>> # 2. Rerank for quality - >>> ranked = await ai.rerank( - ... reranker='vertexai/semantic-ranker-default@latest', - ... query='How do neural networks learn?', - ... documents=candidates, - ... options={'top_n': 5}, - ... ) - >>> - >>> # 3. Generate with top results - >>> response = await ai.generate( - ... model='vertexai/gemini-2.0-flash', - ... prompt='Explain how neural networks learn.', - ... docs=ranked, - ... ) - -Caveats: - - Requires Google Cloud project with Discovery Engine API enabled - - Reranking adds latency - use for quality-critical applications - - Models may silently fall back to default if name is not recognized - -See Also: - - Vertex AI Ranking API: https://cloud.google.com/generative-ai-app-builder/docs/ranking - - RAG best practices: https://genkit.dev/docs/rag -""" - -from genkit.plugins.google_genai.rerankers.reranker import ( - DEFAULT_MODEL_NAME, - KNOWN_MODELS, - RerankRequest, - RerankResponse, - VertexRerankerClientOptions, - VertexRerankerConfig, - _from_rerank_response, - _to_reranker_doc, - is_reranker_model_name, - reranker_rank, -) - -__all__ = [ - 'DEFAULT_MODEL_NAME', - 'KNOWN_MODELS', - 'RerankRequest', - 'RerankResponse', - 'VertexRerankerClientOptions', - 'VertexRerankerConfig', - '_from_rerank_response', - '_to_reranker_doc', - 'is_reranker_model_name', - 'reranker_rank', -] diff --git a/py/plugins/google-genai/src/genkit/plugins/google_genai/rerankers/reranker.py b/py/plugins/google-genai/src/genkit/plugins/google_genai/rerankers/reranker.py deleted file mode 100644 index 196cd4c5b5..0000000000 --- a/py/plugins/google-genai/src/genkit/plugins/google_genai/rerankers/reranker.py +++ /dev/null @@ -1,361 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Vertex AI Reranker implementation. - -This module implements the Vertex AI Discovery Engine Ranking API for reranking -documents based on their semantic relevance to a query. - -Architecture:: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Vertex AI Reranker Module β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Constants & Configuration β”‚ - β”‚ β”œβ”€β”€ DEFAULT_LOCATION (global) β”‚ - β”‚ β”œβ”€β”€ DEFAULT_MODEL_NAME (semantic-ranker-default@latest) β”‚ - β”‚ └── KNOWN_MODELS (supported model registry) β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Request/Response Types (Pydantic) β”‚ - β”‚ β”œβ”€β”€ VertexRerankerConfig - User-facing configuration β”‚ - β”‚ β”œβ”€β”€ VertexRerankerClientOptions - Internal client config β”‚ - β”‚ β”œβ”€β”€ RerankRequest, RerankRequestRecord - API request types β”‚ - β”‚ └── RerankResponse, RerankResponseRecord - API response types β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ API Client β”‚ - β”‚ β”œβ”€β”€ reranker_rank() - Async API call to Discovery Engine β”‚ - β”‚ └── get_vertex_rerank_url() - URL builder for ranking endpoint β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Conversion Functions β”‚ - β”‚ β”œβ”€β”€ _to_reranker_doc() - Document β†’ RerankRequestRecord β”‚ - β”‚ └── _from_rerank_response() - Response β†’ RankedDocument list β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Implementation Notes: - - Uses Google Cloud Application Default Credentials (ADC) for auth - - Calls the Discovery Engine rankingConfigs:rank endpoint - - Supports configurable location and top_n parameters - - Returns RankedDocument instances with scores - -Note: - The actual reranker action registration is handled by the VertexAI plugin - in google.py via the _resolve_reranker method, which uses the conversion - functions and API client defined here. -""" - -from __future__ import annotations - -import asyncio -import json -from typing import Any, ClassVar - -from google.auth import default as google_auth_default -from google.auth.transport.requests import Request -from pydantic import BaseModel, ConfigDict, Field - -from genkit.blocks.document import Document -from genkit.blocks.model import text_from_content -from genkit.blocks.reranker import RankedDocument -from genkit.core.error import GenkitError -from genkit.core.http_client import get_cached_client -from genkit.core.typing import DocumentData - -# Default location for Vertex AI Ranking API (global is recommended per docs) -DEFAULT_LOCATION = 'global' - -# Default reranker model name -DEFAULT_MODEL_NAME = 'semantic-ranker-default@latest' - -# Known reranker models -KNOWN_MODELS: dict[str, str] = { - 'semantic-ranker-default@latest': 'semantic-ranker-default@latest', - 'semantic-ranker-default-004': 'semantic-ranker-default-004', - 'semantic-ranker-fast-004': 'semantic-ranker-fast-004', - 'semantic-ranker-default-003': 'semantic-ranker-default-003', - 'semantic-ranker-default-002': 'semantic-ranker-default-002', -} - - -def is_reranker_model_name(value: str | None) -> bool: - """Check if a value is a valid reranker model name. - - Args: - value: The value to check. - - Returns: - True if the value is a valid reranker model name. - """ - return value is not None and value.startswith('semantic-ranker-') - - -class VertexRerankerConfig(BaseModel): - """Configuration options for Vertex AI reranker. - - Attributes: - top_n: Number of top documents to return. If not specified, all documents - are returned with their scores. - ignore_record_details_in_response: If True, the response will only contain - record ID and score. Defaults to False. - location: Google Cloud location (e.g., "us-central1"). If not specified, - uses the default location from plugin options. - """ - - model_config: ClassVar[ConfigDict] = ConfigDict( - extra='allow', - populate_by_name=True, - ) - - top_n: int | None = Field(default=None, alias='topN') - ignore_record_details_in_response: bool | None = Field( - default=None, - alias='ignoreRecordDetailsInResponse', - ) - location: str | None = None - - -class RerankRequestRecord(BaseModel): - """A record to be reranked. - - Attributes: - id: Unique identifier for the record. - title: Optional title of the record. - content: The content of the record to be ranked. - """ - - id: str - title: str | None = None - content: str - - -class RerankRequest(BaseModel): - """Request body for the rerank API. - - Attributes: - model: The reranker model to use. - query: The query to rank documents against. - records: The records to be ranked. - top_n: Number of top documents to return. - ignore_record_details_in_response: If True, only return ID and score. - """ - - model_config: ClassVar[ConfigDict] = ConfigDict( - extra='allow', - populate_by_name=True, - ) - - model: str - query: str - records: list[RerankRequestRecord] - top_n: int | None = Field(default=None, alias='topN') - ignore_record_details_in_response: bool | None = Field( - default=None, - alias='ignoreRecordDetailsInResponse', - ) - - -class RerankResponseRecord(BaseModel): - """A record in the rerank response. - - Attributes: - id: The record ID. - score: The relevance score (0-1). - content: The record content (if not ignored). - title: The record title (if present). - """ - - id: str - score: float - content: str | None = None - title: str | None = None - - -class RerankResponse(BaseModel): - """Response from the rerank API. - - Attributes: - records: The ranked records with scores. - """ - - records: list[RerankResponseRecord] - - -class VertexRerankerClientOptions(BaseModel): - """Client options for the Vertex AI reranker. - - Attributes: - project_id: Google Cloud project ID. - location: Google Cloud location (e.g., "us-central1"). - """ - - project_id: str - location: str = DEFAULT_LOCATION - - -async def reranker_rank( - model: str, - request: RerankRequest, - client_options: VertexRerankerClientOptions, -) -> RerankResponse: - """Call the Vertex AI Ranking API. - - Args: - model: The reranker model name. - request: The rerank request. - client_options: Client options including project and location. - - Returns: - The rerank response with scored records. - - Raises: - GenkitError: If the API call fails. - """ - url = get_vertex_rerank_url(client_options) - - # Get authentication token - # Use asyncio.to_thread to avoid blocking the event loop during token refresh - credentials, _ = google_auth_default() - await asyncio.to_thread(credentials.refresh, Request()) - token = credentials.token - - if not token: - raise GenkitError( - message='Unable to authenticate your request. ' - 'Please ensure you have valid Google Cloud credentials configured.', - status='UNAUTHENTICATED', - ) - - headers = { - 'Authorization': f'Bearer {token}', - 'x-goog-user-project': client_options.project_id, - 'Content-Type': 'application/json', - } - - # Prepare request body - only include non-None values - request_body: dict[str, Any] = { - 'model': request.model, - 'query': request.query, - 'records': [r.model_dump(exclude_none=True) for r in request.records], - } - if request.top_n is not None: - request_body['topN'] = request.top_n - if request.ignore_record_details_in_response is not None: - request_body['ignoreRecordDetailsInResponse'] = request.ignore_record_details_in_response - - # Use cached client for better connection reuse. - # Note: Auth headers are passed per-request since tokens may expire. - client = get_cached_client( - cache_key='vertex-ai-reranker', - timeout=60.0, - ) - - try: - response = await client.post( - url, - headers=headers, - json=request_body, - ) - - if response.status_code != 200: - error_message = response.text - try: - error_json = response.json() - if 'error' in error_json and 'message' in error_json['error']: - error_message = error_json['error']['message'] - except json.JSONDecodeError: # noqa: S110 - # JSON parsing failed, use raw text - pass - - raise GenkitError( - message=f'Error calling Vertex AI Reranker API: [{response.status_code}] {error_message}', - status='INTERNAL', - ) - - return RerankResponse.model_validate(response.json()) - - except Exception as e: - if isinstance(e, GenkitError): - raise - raise GenkitError( - message=f'Failed to call Vertex AI Reranker API: {e}', - status='UNAVAILABLE', - ) from e - - -def get_vertex_rerank_url(client_options: VertexRerankerClientOptions) -> str: - """Get the URL for the Vertex AI Ranking API. - - Args: - client_options: Client options including project and location. - - Returns: - The API endpoint URL. - """ - return ( - f'https://discoveryengine.googleapis.com/v1/projects/{client_options.project_id}' - f'/locations/{client_options.location}/rankingConfigs/default_ranking_config:rank' - ) - - -def _to_reranker_doc(doc: Document | DocumentData, idx: int) -> RerankRequestRecord: - """Convert a document to a rerank request record. - - Args: - doc: The document to convert. - idx: The index of the document (used as ID). - - Returns: - A rerank request record. - """ - if isinstance(doc, Document): - text = doc.text() - else: - # DocumentData - use text_from_content helper - text = text_from_content(doc.content) - - return RerankRequestRecord( - id=str(idx), - content=text, - ) - - -def _from_rerank_response( - response: RerankResponse, - documents: list[Document], -) -> list[RankedDocument]: - """Convert rerank response to ranked documents. - - Args: - response: The rerank response. - documents: The original documents. - - Returns: - RankedDocument instances with scores, sorted by relevance. - """ - ranked_docs: list[RankedDocument] = [] - for record in response.records: - idx = int(record.id) - original_doc = documents[idx] - - # Create RankedDocument with the score from the API response - ranked_docs.append( - RankedDocument( - content=original_doc.content, - metadata=original_doc.metadata, - score=record.score, - ) - ) - - return ranked_docs diff --git a/py/plugins/google-genai/test/google_plugin_test.py b/py/plugins/google-genai/test/google_plugin_test.py index bc3c55e1b5..ff0728f077 100644 --- a/py/plugins/google-genai/test/google_plugin_test.py +++ b/py/plugins/google-genai/test/google_plugin_test.py @@ -28,8 +28,17 @@ from google.auth.credentials import Credentials from google.genai.types import HttpOptions -from genkit.ai import GENKIT_CLIENT_HEADER, Genkit -from genkit.core.registry import ActionKind +from genkit import ( + ActionKind, + Genkit, + Message, + ModelInfo, + ModelRequest, + Part, + Role, + TextPart, +) +from genkit.plugin_api import GENKIT_CLIENT_HEADER from genkit.plugins.google_genai import GoogleAI, VertexAI from genkit.plugins.google_genai.google import _inject_attribution_headers, googleai_name, vertexai_name from genkit.plugins.google_genai.models.gemini import ( @@ -42,14 +51,6 @@ DEFAULT_IMAGE_SUPPORT, SUPPORTED_MODELS as IMAGE_SUPPORTED_MODELS, ) -from genkit.types import ( - GenerateRequest, - Message, - ModelInfo, - Part, - Role, - TextPart, -) async def _get_runtime_client(plugin: GoogleAI | VertexAI) -> object: @@ -961,7 +962,7 @@ async def test_system_prompt_handling() -> None: mock_client = MagicMock(spec=genai.Client) model = GeminiModel(version='gemini-1.5-flash', client=mock_client) - request = GenerateRequest( + request = ModelRequest( messages=[ Message(role=Role.SYSTEM, content=[Part(root=TextPart(text='You are a helpful assistant'))]), Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))]), diff --git a/py/plugins/google-genai/test/models/googlegenai_embedder_test.py b/py/plugins/google-genai/test/models/googlegenai_embedder_test.py index 835bf8d185..37f4d11f64 100644 --- a/py/plugins/google-genai/test/models/googlegenai_embedder_test.py +++ b/py/plugins/google-genai/test/models/googlegenai_embedder_test.py @@ -20,15 +20,15 @@ from google import genai from pytest_mock import MockerFixture -from genkit.ai import Document +from genkit import ( + Document, + EmbedRequest, + EmbedResponse, +) from genkit.plugins.google_genai.models.embedder import ( Embedder, GeminiEmbeddingModels, ) -from genkit.types import ( - EmbedRequest, - EmbedResponse, -) @pytest.mark.asyncio diff --git a/py/plugins/google-genai/test/models/googlegenai_gemini_test.py b/py/plugins/google-genai/test/models/googlegenai_gemini_test.py index e282435700..41e93ac0aa 100644 --- a/py/plugins/google-genai/test/models/googlegenai_gemini_test.py +++ b/py/plugins/google-genai/test/models/googlegenai_gemini_test.py @@ -32,8 +32,19 @@ from pydantic import BaseModel, Field from pytest_mock import MockerFixture -from genkit.ai import ActionRunContext -from genkit.core.schema import to_json_schema +from genkit import ( + ActionRunContext, + MediaPart, + Message, + ModelInfo, + ModelRequest, + ModelResponse, + Part, + Role, + TextPart, + ToolDefinition, +) +from genkit.plugin_api import to_json_schema from genkit.plugins.google_genai.models.gemini import ( DEFAULT_SUPPORTS_MODEL, GeminiModel, @@ -43,17 +54,6 @@ is_image_model, is_tts_model, ) -from genkit.types import ( - GenerateRequest, - GenerateResponse, - MediaPart, - Message, - ModelInfo, - Part, - Role, - TextPart, - ToolDefinition, -) ALL_VERSIONS = list(GoogleAIGeminiVersion) + list(VertexAIGeminiVersion) IMAGE_GENERATION_VERSIONS = [GoogleAIGeminiVersion.GEMINI_2_0_FLASH_EXP] @@ -66,7 +66,7 @@ async def test_generate_text_response(mocker: MockerFixture, version: str) -> No response_text = 'request answer' request_text = 'response question' - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -102,7 +102,7 @@ async def test_generate_text_response(mocker: MockerFixture, version: str) -> No config=expected_config, ) ]) - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None assert response.message.content[0].root.text == response_text @@ -114,7 +114,7 @@ async def test_generate_stream_text_response(mocker: MockerFixture, version: str response_text = 'request answer' request_text = 'response question' - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -133,7 +133,7 @@ async def test_generate_stream_text_response(mocker: MockerFixture, version: str on_chunk_mock = mocker.MagicMock() gemini = GeminiModel(version, googleai_client_mock) - ctx = ActionRunContext(on_chunk=on_chunk_mock) + ctx = ActionRunContext(streaming_callback=on_chunk_mock) response = await gemini.generate(request, ctx) # Determine expected config based on model type @@ -151,7 +151,7 @@ async def test_generate_stream_text_response(mocker: MockerFixture, version: str config=expected_config, ) ]) - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None assert response.message.content == [] @@ -165,7 +165,7 @@ async def test_generate_media_response(mocker: MockerFixture, version: str) -> N response_mimetype = 'image/png' modalities = ['Text', 'Image'] - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -201,7 +201,7 @@ async def test_generate_media_response(mocker: MockerFixture, version: str) -> N config=genai.types.GenerateContentConfig(response_modalities=modalities), ) ]) - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None content = response.message.content[0] @@ -299,7 +299,7 @@ async def test_generate_with_system_instructions(mocker: MockerFixture) -> None: system_instruction = 'system instruction text' version = GoogleAIGeminiVersion.GEMINI_2_0_FLASH - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -335,7 +335,7 @@ async def test_generate_with_system_instructions(mocker: MockerFixture) -> None: config=genai.types.GenerateContentConfig(system_instruction=expected_system_instruction), ) ]) - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None assert response.message.content[0].root.text == response_text @@ -431,7 +431,7 @@ def test_gemini_model__get_tools( ), ] - request = GenerateRequest( + request = ModelRequest( tools=request_tools, messages=[ Message( @@ -756,7 +756,7 @@ class MockPage(AsyncMock): gemini_model_instance._client = mock_client - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, diff --git a/py/plugins/google-genai/test/models/googlegenai_imagen_test.py b/py/plugins/google-genai/test/models/googlegenai_imagen_test.py index 63f3070127..0172edf3c3 100644 --- a/py/plugins/google-genai/test/models/googlegenai_imagen_test.py +++ b/py/plugins/google-genai/test/models/googlegenai_imagen_test.py @@ -23,17 +23,17 @@ from google import genai from pytest_mock import MockerFixture -from genkit.ai import ActionRunContext -from genkit.plugins.google_genai.models.imagen import ImagenModel, ImagenVersion -from genkit.types import ( - GenerateRequest, - GenerateResponse, +from genkit import ( + ActionRunContext, MediaPart, Message, + ModelRequest, + ModelResponse, Part, Role, TextPart, ) +from genkit.plugins.google_genai.models.imagen import ImagenModel, ImagenVersion @pytest.mark.asyncio @@ -44,7 +44,7 @@ async def test_generate_media_response(mocker: MockerFixture, version: ImagenVer response_byte_string = b'\x89PNG\r\n\x1a\n' response_mimetype = 'image/png' - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -74,7 +74,7 @@ async def test_generate_media_response(mocker: MockerFixture, version: ImagenVer googleai_client_mock.assert_has_calls([ mocker.call.aio.models.generate_images(model=version, prompt=request_text, config=None) ]) - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None content = response.message.content[0] assert isinstance(content.root, MediaPart) diff --git a/py/plugins/google-genai/tests/evaluators_test.py b/py/plugins/google-genai/tests/evaluators_test.py index 823531df31..6718816d33 100644 --- a/py/plugins/google-genai/tests/evaluators_test.py +++ b/py/plugins/google-genai/tests/evaluators_test.py @@ -154,7 +154,7 @@ async def test_evaluator_factory_evaluate_instances_structure() -> None: @pytest.mark.asyncio async def test_evaluator_factory_evaluate_instances_error_handling() -> None: """Test that evaluate_instances raises GenkitError on API failure.""" - from genkit.core.error import GenkitError + from genkit._core._error import GenkitError factory = EvaluatorFactory( project_id='test-project', diff --git a/py/plugins/google-genai/tests/google_genai_plugin_test.py b/py/plugins/google-genai/tests/google_genai_plugin_test.py index 339998e6f3..9cdb68526f 100644 --- a/py/plugins/google-genai/tests/google_genai_plugin_test.py +++ b/py/plugins/google-genai/tests/google_genai_plugin_test.py @@ -24,7 +24,7 @@ import pytest -from genkit.core.registry import ActionKind +from genkit import ActionKind from genkit.plugins.google_genai import ( EmbeddingTaskType, GeminiConfigSchema, diff --git a/py/plugins/google-genai/tests/part_converter_test.py b/py/plugins/google-genai/tests/part_converter_test.py index 0a46b2e841..18ac208cea 100644 --- a/py/plugins/google-genai/tests/part_converter_test.py +++ b/py/plugins/google-genai/tests/part_converter_test.py @@ -25,8 +25,8 @@ import pytest from google import genai +from genkit import Media, MediaPart, Part from genkit.plugins.google_genai.models.utils import PartConverter -from genkit.types import Media, MediaPart, Part class TestIsGeminiNativeUrl: diff --git a/py/plugins/google-genai/tests/rerankers_test.py b/py/plugins/google-genai/tests/rerankers_test.py deleted file mode 100644 index 7cfd0cd8f8..0000000000 --- a/py/plugins/google-genai/tests/rerankers_test.py +++ /dev/null @@ -1,346 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Vertex AI Rerankers.""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from genkit.blocks.document import Document -from genkit.core.typing import TextPart -from genkit.plugins.google_genai.rerankers import ( - DEFAULT_MODEL_NAME, - KNOWN_MODELS, - VertexRerankerConfig, - is_reranker_model_name, -) -from genkit.plugins.google_genai.rerankers.reranker import ( - RerankRequest, - RerankRequestRecord, - RerankResponse, - RerankResponseRecord, - VertexRerankerClientOptions, - _from_rerank_response, - _to_reranker_doc, - get_vertex_rerank_url, -) - - -def test_default_model_name() -> None: - """Test that DEFAULT_MODEL_NAME is set correctly.""" - assert DEFAULT_MODEL_NAME == 'semantic-ranker-default@latest' - - -def test_known_models_contains_expected_models() -> None: - """Test that KNOWN_MODELS contains expected reranker models.""" - assert 'semantic-ranker-default@latest' in KNOWN_MODELS - assert 'semantic-ranker-default-004' in KNOWN_MODELS - assert 'semantic-ranker-fast-004' in KNOWN_MODELS - assert 'semantic-ranker-default-003' in KNOWN_MODELS - assert 'semantic-ranker-default-002' in KNOWN_MODELS - - -def test_is_reranker_model_name_valid() -> None: - """Test is_reranker_model_name returns True for valid names.""" - assert is_reranker_model_name('semantic-ranker-default@latest') is True - assert is_reranker_model_name('semantic-ranker-fast-004') is True - - -def test_is_reranker_model_name_invalid() -> None: - """Test is_reranker_model_name returns False for invalid names.""" - assert is_reranker_model_name('gemini-2.0-flash') is False - assert is_reranker_model_name('gemini-embedding-001') is False - assert is_reranker_model_name(None) is False - assert is_reranker_model_name('') is False - - -def test_vertex_reranker_config() -> None: - """Test VertexRerankerConfig model.""" - config = VertexRerankerConfig(top_n=5) - assert config.top_n == 5 - - -def test_vertex_reranker_config_defaults() -> None: - """Test VertexRerankerConfig default values.""" - config = VertexRerankerConfig() - assert config.top_n is None - assert config.location is None - assert config.ignore_record_details_in_response is None - - -def test_vertex_reranker_config_with_aliases() -> None: - """Test VertexRerankerConfig works with aliases.""" - # Use Python field names (populate_by_name=True allows both) - config = VertexRerankerConfig(top_n=10, ignore_record_details_in_response=True) - assert config.top_n == 10 - assert config.ignore_record_details_in_response is True - - -def test_vertex_reranker_client_options() -> None: - """Test VertexRerankerClientOptions model.""" - options = VertexRerankerClientOptions( - project_id='my-project', - location='us-central1', - ) - assert options.project_id == 'my-project' - assert options.location == 'us-central1' - - -def test_vertex_reranker_client_options_default_location() -> None: - """Test VertexRerankerClientOptions uses default location.""" - options = VertexRerankerClientOptions(project_id='test-project') - assert options.project_id == 'test-project' - assert options.location == 'global' - - -def test_rerank_request_record() -> None: - """Test RerankRequestRecord model.""" - record = RerankRequestRecord( - id='doc-1', - title='Test Document', - content='This is the document content.', - ) - assert record.id == 'doc-1' - assert record.title == 'Test Document' - assert record.content == 'This is the document content.' - - -def test_rerank_request_record_no_title() -> None: - """Test RerankRequestRecord without optional title.""" - record = RerankRequestRecord(id='1', content='Content only') - assert record.id == '1' - assert record.title is None - assert record.content == 'Content only' - - -def test_rerank_request() -> None: - """Test RerankRequest model.""" - records = [ - RerankRequestRecord(id='1', content='Doc 1'), - RerankRequestRecord(id='2', content='Doc 2'), - ] - request = RerankRequest( - query='What is machine learning?', - records=records, - model='semantic-ranker-default@latest', - top_n=5, - ) - assert request.query == 'What is machine learning?' - assert len(request.records) == 2 - assert request.model == 'semantic-ranker-default@latest' - assert request.top_n == 5 - - -def test_rerank_response_record() -> None: - """Test RerankResponseRecord model.""" - record = RerankResponseRecord( - id='doc-1', - score=0.95, - content='Document content', - ) - assert record.id == 'doc-1' - assert record.score == 0.95 - assert record.content == 'Document content' - - -def test_rerank_response_record_minimal() -> None: - """Test RerankResponseRecord with only required fields.""" - record = RerankResponseRecord(id='1', score=0.5) - assert record.id == '1' - assert record.score == 0.5 - assert record.content is None - assert record.title is None - - -def test_rerank_response() -> None: - """Test RerankResponse model.""" - records = [ - RerankResponseRecord(id='1', score=0.9), - RerankResponseRecord(id='2', score=0.7), - ] - response = RerankResponse(records=records) - assert len(response.records) == 2 - assert response.records[0].score == 0.9 - - -def test_get_vertex_rerank_url() -> None: - """Test get_vertex_rerank_url builds correct URL.""" - options = VertexRerankerClientOptions( - project_id='my-project', - location='us-central1', - ) - - url = get_vertex_rerank_url(options) - - assert 'my-project' in url - assert 'us-central1' in url - assert 'discoveryengine.googleapis.com' in url - assert ':rank' in url - - -def test_get_vertex_rerank_url_different_location() -> None: - """Test get_vertex_rerank_url with different location.""" - options = VertexRerankerClientOptions( - project_id='test-project', - location='europe-west1', - ) - - url = get_vertex_rerank_url(options) - - assert 'test-project' in url - assert 'europe-west1' in url - - -def test_to_reranker_doc_from_document() -> None: - """Test _to_reranker_doc converts Document to RerankRequestRecord.""" - from genkit.core.typing import DocumentPart - - doc = Document(content=[DocumentPart(root=TextPart(text='This is document content.'))]) - - record = _to_reranker_doc(doc, 0) - - assert record.content == 'This is document content.' - assert record.id == '0' - - -def test_to_reranker_doc_different_index() -> None: - """Test _to_reranker_doc uses provided index.""" - from genkit.core.typing import DocumentPart - - doc = Document(content=[DocumentPart(root=TextPart(text='Content'))]) - - record = _to_reranker_doc(doc, 5) - - assert record.id == '5' - - -def test_from_rerank_response_basic() -> None: - """Test _from_rerank_response converts response to scored documents.""" - from genkit.core.typing import DocumentPart - - original_docs = [ - Document(content=[DocumentPart(root=TextPart(text='Doc 0'))]), - Document(content=[DocumentPart(root=TextPart(text='Doc 1'))]), - Document(content=[DocumentPart(root=TextPart(text='Doc 2'))]), - ] - - response = RerankResponse( - records=[ - RerankResponseRecord(id='1', score=0.9), - RerankResponseRecord(id='0', score=0.7), - RerankResponseRecord(id='2', score=0.5), - ] - ) - - result = _from_rerank_response(response, original_docs) - - assert len(result) == 3 - for doc in result: - assert doc.metadata is not None - assert 'score' in doc.metadata - - -def test_from_rerank_response_preserves_content() -> None: - """Test _from_rerank_response preserves document content.""" - from genkit.core.typing import DocumentPart - - original_docs = [ - Document(content=[DocumentPart(root=TextPart(text='Original content'))]), - ] - - response = RerankResponse(records=[RerankResponseRecord(id='0', score=0.85)]) - - result = _from_rerank_response(response, original_docs) - - assert len(result) == 1 - assert result[0].text() == 'Original content' - assert result[0].metadata is not None - assert result[0].metadata.get('score') == 0.85 - - -def test_from_rerank_response_preserves_original_metadata() -> None: - """Test _from_rerank_response preserves original document metadata.""" - from genkit.core.typing import DocumentPart - - original_docs = [ - Document( - content=[DocumentPart(root=TextPart(text='Content'))], - metadata={'custom_field': 'value'}, - ), - ] - - response = RerankResponse(records=[RerankResponseRecord(id='0', score=0.85)]) - - result = _from_rerank_response(response, original_docs) - - assert len(result) == 1 - assert result[0].metadata is not None - assert result[0].metadata.get('custom_field') == 'value' - assert result[0].metadata.get('score') == 0.85 - - -def test_from_rerank_response_empty() -> None: - """Test _from_rerank_response handles empty response.""" - response = RerankResponse(records=[]) - - result = _from_rerank_response(response, []) - - assert result == [] - - -@pytest.mark.asyncio -async def test_reranker_api_call_structure() -> None: - """Test that reranker API call is structured correctly.""" - from genkit.plugins.google_genai.rerankers.reranker import reranker_rank - - mock_credentials = MagicMock() - mock_credentials.token = 'mock-token' - mock_credentials.expired = False - - mock_response_data = { - 'records': [ - {'id': '0', 'score': 0.9}, - {'id': '1', 'score': 0.7}, - ] - } - - with patch('genkit.plugins.google_genai.rerankers.reranker.google_auth_default') as mock_auth: - mock_auth.return_value = (mock_credentials, 'test-project') - - # Mock get_cached_client to return a mock client - mock_client = AsyncMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = mock_response_data - mock_client.post = AsyncMock(return_value=mock_response) - mock_client.is_closed = False - - with patch('genkit.plugins.google_genai.rerankers.reranker.get_cached_client', return_value=mock_client): - request = RerankRequest( - model='semantic-ranker-default@latest', - query='test query', - records=[ - RerankRequestRecord(id='0', content='Doc 1'), - RerankRequestRecord(id='1', content='Doc 2'), - ], - ) - options = VertexRerankerClientOptions(project_id='test-project') - - result = await reranker_rank('semantic-ranker-default@latest', request, options) - - assert isinstance(result, RerankResponse) - assert len(result.records) == 2 diff --git a/py/plugins/ollama/src/genkit/plugins/ollama/converters.py b/py/plugins/ollama/src/genkit/plugins/ollama/converters.py index ea8329448e..75795e3467 100644 --- a/py/plugins/ollama/src/genkit/plugins/ollama/converters.py +++ b/py/plugins/ollama/src/genkit/plugins/ollama/converters.py @@ -28,10 +28,10 @@ from typing import Any, Literal, cast -from genkit.types import ( - GenerationCommonConfig, - GenerationUsage, +from genkit import ( Message, + ModelConfig, + ModelUsage, Part, Role, TextPart, @@ -94,11 +94,11 @@ def build_prompt(messages: list[Message]) -> str: def build_request_options_dict( - config: GenerationCommonConfig | dict[str, object] | None, + config: ModelConfig | dict[str, object] | None, ) -> dict[str, Any]: """Build options dict from config for the Ollama API. - Maps Genkit ``GenerationCommonConfig`` fields to Ollama option names. + Maps Genkit ``ModelConfig`` fields to Ollama option names. Args: config: Request configuration. @@ -109,7 +109,7 @@ def build_request_options_dict( if config is None: return {} - if isinstance(config, GenerationCommonConfig): + if isinstance(config, ModelConfig): result: dict[str, Any] = {} if config.top_k is not None: result['top_k'] = config.top_k @@ -165,10 +165,10 @@ def build_response_parts( def get_usage_info( - basic_usage: GenerationUsage, + basic_usage: ModelUsage, prompt_eval_count: int | None, eval_count: int | None, -) -> GenerationUsage: +) -> ModelUsage: """Update basic usage with token counts from Ollama API response. Args: @@ -177,7 +177,7 @@ def get_usage_info( eval_count: Output token count from Ollama. Returns: - Updated GenerationUsage with token counts. + Updated ModelUsage with token counts. """ basic_usage.input_tokens = prompt_eval_count or 0 basic_usage.output_tokens = eval_count or 0 diff --git a/py/plugins/ollama/src/genkit/plugins/ollama/embedders.py b/py/plugins/ollama/src/genkit/plugins/ollama/embedders.py index 804b376dde..347d335b09 100644 --- a/py/plugins/ollama/src/genkit/plugins/ollama/embedders.py +++ b/py/plugins/ollama/src/genkit/plugins/ollama/embedders.py @@ -22,8 +22,8 @@ from pydantic import BaseModel import ollama as ollama_api -from genkit.blocks.embedding import EmbedRequest, EmbedResponse -from genkit.types import Embedding +from genkit import Embedding +from genkit.embedder import EmbedRequest, EmbedResponse class EmbeddingDefinition(BaseModel): @@ -69,7 +69,7 @@ def __init__( self._client_factory = client self.embedding_definition = embedding_definition - def _get_client(self) -> 'ollama_api.AsyncClient': + def _get_client(self) -> ollama_api.AsyncClient: """Creates a fresh async client bound to the current event loop. Returns: diff --git a/py/plugins/ollama/src/genkit/plugins/ollama/models.py b/py/plugins/ollama/src/genkit/plugins/ollama/models.py index a06cf48967..2efc80df41 100644 --- a/py/plugins/ollama/src/genkit/plugins/ollama/models.py +++ b/py/plugins/ollama/src/genkit/plugins/ollama/models.py @@ -90,21 +90,15 @@ from pydantic import BaseModel import ollama as ollama_api -from genkit.ai import ActionRunContext -from genkit.blocks.model import get_basic_usage_stats -from genkit.core.http_client import get_cached_client -from genkit.plugins.ollama.constants import ( - OllamaAPITypes, -) -from genkit.types import ( - GenerateRequest, - GenerateResponse, - GenerateResponseChunk, - GenerationCommonConfig, - GenerationUsage, +from genkit import ( Media, MediaPart, Message, + ModelConfig, + ModelRequest, + ModelResponse, + ModelResponseChunk, + ModelUsage, Part, Role, TextPart, @@ -112,6 +106,11 @@ ToolRequestPart, ToolResponsePart, ) +from genkit.model import get_basic_usage_stats +from genkit.plugin_api import ActionRunContext, get_cached_client +from genkit.plugins.ollama.constants import ( + OllamaAPITypes, +) logger = structlog.get_logger(__name__) @@ -166,7 +165,7 @@ def _get_client(self) -> ollama_api.AsyncClient: """ return self._client_factory() - async def generate(self, request: GenerateRequest, ctx: ActionRunContext | None = None) -> GenerateResponse: + async def generate(self, request: ModelRequest, ctx: ActionRunContext | None = None) -> ModelResponse: """Generate a response from Ollama. Args: @@ -221,7 +220,7 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext | None response=response_message, ) - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=content, @@ -233,7 +232,7 @@ async def generate(self, request: GenerateRequest, ctx: ActionRunContext | None ) async def _chat_with_ollama( - self, request: GenerateRequest, ctx: ActionRunContext | None = None + self, request: ModelRequest, ctx: ActionRunContext | None = None ) -> ollama_api.ChatResponse | None: """Chat with Ollama. @@ -247,12 +246,12 @@ async def _chat_with_ollama( messages = await self.build_chat_messages(request) streaming_request = self.is_streaming_request(ctx=ctx) - if request.output: + if request.output_format or request.output_schema: # ollama api either accepts 'json' literal, or the JSON schema - if request.output.schema: - fmt = request.output.schema - elif request.output.format: - fmt = request.output.format + if request.output_schema: + fmt = request.output_schema + elif request.output_format: + fmt = request.output_format else: fmt = '' else: @@ -287,7 +286,7 @@ async def _chat_with_ollama( role = Role.MODEL if chunk.message.role == 'assistant' else Role.TOOL if ctx: ctx.send_chunk( - chunk=GenerateResponseChunk( + chunk=ModelResponseChunk( role=role, index=idx, content=self._build_multimodal_chat_response(chat_response=chunk), @@ -310,7 +309,7 @@ async def _chat_with_ollama( return chat_response async def _generate_ollama_response( - self, request: GenerateRequest, ctx: ActionRunContext | None = None + self, request: ModelRequest, ctx: ActionRunContext | None = None ) -> ollama_api.GenerateResponse | None: """Generate a response from Ollama. @@ -338,7 +337,7 @@ async def _generate_ollama_response( idx += 1 if ctx: ctx.send_chunk( - chunk=GenerateResponseChunk( + chunk=ModelResponseChunk( role=Role.MODEL, index=idx, content=[Part(root=TextPart(text=chunk.response))], @@ -403,7 +402,7 @@ def _build_multimodal_chat_response( @staticmethod def build_request_options( - config: GenerationCommonConfig | ollama_api.Options | dict[str, object] | None, + config: ModelConfig | ollama_api.Options | dict[str, object] | None, ) -> ollama_api.Options: """Build request options for the generate API. @@ -415,7 +414,7 @@ def build_request_options( """ if config is None: return ollama_api.Options() - if isinstance(config, GenerationCommonConfig): + if isinstance(config, ModelConfig): config = dict( top_k=config.top_k, topP=config.top_p, @@ -430,7 +429,7 @@ def build_request_options( return config @staticmethod - def build_prompt(request: GenerateRequest) -> str: + def build_prompt(request: ModelRequest) -> str: """Build the prompt for the generate API. Args: @@ -449,7 +448,7 @@ def build_prompt(request: GenerateRequest) -> str: return prompt @classmethod - async def build_chat_messages(cls, request: GenerateRequest) -> list[ollama_api.Message]: + async def build_chat_messages(cls, request: ModelRequest) -> list[ollama_api.Message]: """Build the messages for the chat API. Handles MediaPart by converting image URLs to the format expected @@ -553,21 +552,21 @@ def is_streaming_request(ctx: ActionRunContext | None) -> bool: @staticmethod def get_usage_info( - basic_generation_usage: GenerationUsage, + basic_generation_usage: ModelUsage, api_response: ollama_api.GenerateResponse | ollama_api.ChatResponse | None, - ) -> GenerationUsage: + ) -> ModelUsage: """Extracts and calculates token usage information from an Ollama API response. Updates a basic generation usage object with input, output, and total token counts based on the details provided in the Ollama API response. Args: - basic_generation_usage: An existing GenerationUsage object to update. + basic_generation_usage: An existing ModelUsage object to update. api_response: The response object received from the Ollama API, containing token count details. Returns: - The updated GenerationUsage object with token counts populated. + The updated ModelUsage object with token counts populated. """ if api_response: basic_generation_usage.input_tokens = api_response.prompt_eval_count or 0 diff --git a/py/plugins/ollama/src/genkit/plugins/ollama/plugin_api.py b/py/plugins/ollama/src/genkit/plugins/ollama/plugin_api.py index a3b449ee23..85c86e29eb 100644 --- a/py/plugins/ollama/src/genkit/plugins/ollama/plugin_api.py +++ b/py/plugins/ollama/src/genkit/plugins/ollama/plugin_api.py @@ -21,13 +21,17 @@ import structlog import ollama as ollama_api -from genkit.ai import Plugin -from genkit.blocks.embedding import EmbedderOptions, EmbedderSupports, embedder_action_metadata -from genkit.blocks.model import model_action_metadata -from genkit.core._loop_local import _loop_local_client -from genkit.core.action import Action, ActionMetadata -from genkit.core.registry import ActionKind -from genkit.core.schema import to_json_schema +from genkit import ModelConfig +from genkit.embedder import EmbedderOptions, EmbedderSupports, embedder_action_metadata +from genkit.model import model_action_metadata +from genkit.plugin_api import ( + Action, + ActionKind, + ActionMetadata, + Plugin, + loop_local_client, + to_json_schema, +) from genkit.plugins.ollama.constants import ( DEFAULT_OLLAMA_SERVER_URL, OllamaAPITypes, @@ -40,7 +44,6 @@ ModelDefinition, OllamaModel, ) -from genkit.types import GenerationCommonConfig OLLAMA_PLUGIN_NAME = 'ollama' logger = structlog.get_logger(__name__) @@ -90,7 +93,7 @@ def __init__( self.server_address = server_address or DEFAULT_OLLAMA_SERVER_URL self.request_headers = request_headers or {} - self.client = _loop_local_client(partial(ollama_api.AsyncClient, host=self.server_address)) + self.client = loop_local_client(partial(ollama_api.AsyncClient, host=self.server_address)) async def init(self) -> list: """Initialize the Ollama plugin. @@ -172,7 +175,7 @@ def _create_model_action(self, name: str) -> Action: 'tools': model_ref.supports.tools, 'output': ['text', 'json'], 'constrained': 'all', - 'customOptions': to_json_schema(GenerationCommonConfig), + 'customOptions': to_json_schema(ModelConfig), }, }, ) @@ -242,7 +245,7 @@ async def list_actions(self) -> list[ActionMetadata]: actions.append( model_action_metadata( name=ollama_name(name), - config_schema=GenerationCommonConfig, + config_schema=ModelConfig, info={ 'label': f'Ollama - {name}', 'multiturn': True, diff --git a/py/plugins/ollama/tests/conftest.py b/py/plugins/ollama/tests/conftest.py index c9d188da05..622933b463 100644 --- a/py/plugins/ollama/tests/conftest.py +++ b/py/plugins/ollama/tests/conftest.py @@ -23,7 +23,7 @@ import ollama as ollama_api import pytest -from genkit.ai import Genkit +from genkit import Genkit from genkit.plugins.ollama.constants import OllamaAPITypes from genkit.plugins.ollama.models import ModelDefinition from genkit.plugins.ollama.plugin_api import Ollama diff --git a/py/plugins/ollama/tests/integration_test.py b/py/plugins/ollama/tests/integration_test.py index 1d101039db..ec7f9a2626 100644 --- a/py/plugins/ollama/tests/integration_test.py +++ b/py/plugins/ollama/tests/integration_test.py @@ -21,8 +21,7 @@ import ollama as ollama_api import pytest -from genkit.ai import ActionKind, Genkit -from genkit.types import GenerateResponse, Message, Part, Role, TextPart +from genkit import ActionKind, Genkit, Message, ModelResponse, Part, Role, TextPart @pytest.mark.asyncio @@ -63,7 +62,7 @@ async def fake_chat_response(*args: object, **kwargs: object) -> ollama_api.Chat mock_ollama_api_async_client.return_value.chat.side_effect = fake_chat_response - async def _test_fun() -> GenerateResponse: + async def _test_fun() -> ModelResponse: return await genkit_veneer_chat_model.generate( messages=[ Message( @@ -77,7 +76,7 @@ async def _test_fun() -> GenerateResponse: response = await genkit_veneer_chat_model.flow()(_test_fun)() - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None assert response.message.content[0].root.text == mock_response_message @@ -95,7 +94,7 @@ async def test_async_get_generate_model_response_from_llama_api_flow( response=mock_response_message, ) - async def _test_fun() -> GenerateResponse: + async def _test_fun() -> ModelResponse: return await genkit_veneer_generate_model.generate( messages=[ Message( @@ -109,7 +108,7 @@ async def _test_fun() -> GenerateResponse: response = await genkit_veneer_generate_model.flow()(_test_fun)() - assert isinstance(response, GenerateResponse) + assert isinstance(response, ModelResponse) assert response.message is not None assert response.message.content[0].root.text == mock_response_message diff --git a/py/plugins/ollama/tests/models/embedders_test.py b/py/plugins/ollama/tests/models/embedders_test.py index 72b2624e6f..c4939068b5 100644 --- a/py/plugins/ollama/tests/models/embedders_test.py +++ b/py/plugins/ollama/tests/models/embedders_test.py @@ -21,15 +21,15 @@ import ollama as ollama_api -from genkit.core.typing import DocumentPart -from genkit.plugins.ollama.embedders import EmbeddingDefinition, OllamaEmbedder -from genkit.types import ( +from genkit import ( Document, + DocumentPart, Embedding, EmbedRequest, EmbedResponse, TextPart, ) +from genkit.plugins.ollama.embedders import EmbeddingDefinition, OllamaEmbedder class TestOllamaEmbedderEmbed(unittest.IsolatedAsyncioTestCase): diff --git a/py/plugins/ollama/tests/models/ollama_models_test.py b/py/plugins/ollama/tests/models/ollama_models_test.py index d058a543be..a0e552ec53 100644 --- a/py/plugins/ollama/tests/models/ollama_models_test.py +++ b/py/plugins/ollama/tests/models/ollama_models_test.py @@ -25,21 +25,20 @@ import ollama as ollama_api import pytest -from genkit.plugins.ollama.constants import OllamaAPITypes -from genkit.plugins.ollama.models import ModelDefinition, OllamaModel, _convert_parameters -from genkit.types import ( +from genkit import ( ActionRunContext, - GenerateRequest, - GenerateResponseChunk, - GenerationUsage, Media, MediaPart, Message, - OutputConfig, + ModelRequest, + ModelResponseChunk, + ModelUsage, Part, Role, TextPart, ) +from genkit.plugins.ollama.constants import OllamaAPITypes +from genkit.plugins.ollama.models import ModelDefinition, OllamaModel, _convert_parameters class TestOllamaModelGenerate(unittest.IsolatedAsyncioTestCase): @@ -48,13 +47,13 @@ class TestOllamaModelGenerate(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self) -> None: """Common setup for all async tests.""" self.mock_client = MagicMock() - self.request = GenerateRequest(messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))])]) + self.request = ModelRequest(messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))])]) self.ctx = ActionRunContext() cast(Any, self.ctx).send_chunk = MagicMock() @patch( - 'genkit.blocks.model.get_basic_usage_stats', - return_value=GenerationUsage( + 'genkit.model.get_basic_usage_stats', + return_value=ModelUsage( input_tokens=10, output_tokens=20, total_tokens=30, @@ -86,7 +85,7 @@ async def test_generate_chat_non_streaming(self, mock_get_basic_usage_stats: Mag return_value=[Part(root=TextPart(text='Parsed chat content'))], ) cast(Any, ollama_model).get_usage_info = MagicMock( - return_value=GenerationUsage( + return_value=ModelUsage( input_tokens=5, output_tokens=10, total_tokens=15, @@ -111,12 +110,12 @@ async def test_generate_chat_non_streaming(self, mock_get_basic_usage_stats: Mag self.assertEqual(len(cast(Message, response.message).content), 1) self.assertEqual(cast(Message, response.message).content[0].root.text, 'Parsed chat content') self.assertIsNotNone(response.usage) - self.assertEqual(cast(GenerationUsage, response.usage).input_tokens, 5) - self.assertEqual(cast(GenerationUsage, response.usage).output_tokens, 10) + self.assertEqual(cast(ModelUsage, response.usage).input_tokens, 5) + self.assertEqual(cast(ModelUsage, response.usage).output_tokens, 10) @patch( - 'genkit.blocks.model.get_basic_usage_stats', - return_value=GenerationUsage( + 'genkit.model.get_basic_usage_stats', + return_value=ModelUsage( input_tokens=10, output_tokens=20, total_tokens=30, @@ -143,7 +142,7 @@ async def test_generate_generate_non_streaming(self, mock_get_basic_usage_stats: cast(Any, ollama_model)._chat_with_ollama = AsyncMock() cast(Any, ollama_model).is_streaming_request = MagicMock(return_value=False) cast(Any, ollama_model).get_usage_info = MagicMock( - return_value=GenerationUsage( + return_value=ModelUsage( input_tokens=7, output_tokens=14, total_tokens=21, @@ -166,18 +165,18 @@ async def test_generate_generate_non_streaming(self, mock_get_basic_usage_stats: self.assertEqual(len(cast(Message, response.message).content), 1) self.assertEqual(cast(Message, response.message).content[0].root.text, 'Generated text') self.assertIsNotNone(response.usage) - self.assertEqual(cast(GenerationUsage, response.usage).input_tokens, 7) - self.assertEqual(cast(GenerationUsage, response.usage).output_tokens, 14) + self.assertEqual(cast(ModelUsage, response.usage).input_tokens, 7) + self.assertEqual(cast(ModelUsage, response.usage).output_tokens, 14) @patch( - 'genkit.blocks.model.get_basic_usage_stats', - return_value=GenerationUsage(), + 'genkit.model.get_basic_usage_stats', + return_value=ModelUsage(), ) async def test_generate_chat_streaming(self, mock_get_basic_usage_stats: MagicMock) -> None: """Test generate method with CHAT API type in streaming mode.""" model_def = ModelDefinition(name='chat-model', api_type=OllamaAPITypes.CHAT) ollama_model = OllamaModel(client=self.mock_client, model_definition=model_def) - streaming_ctx = ActionRunContext(on_chunk=MagicMock()) + streaming_ctx = ActionRunContext(streaming_callback=MagicMock()) # Mock internal methods mock_chat_response = ollama_api.ChatResponse( @@ -194,7 +193,7 @@ async def test_generate_chat_streaming(self, mock_get_basic_usage_stats: MagicMo ) cast(Any, ollama_model).is_streaming_request = MagicMock(return_value=True) cast(Any, ollama_model).get_usage_info = MagicMock( - return_value=GenerationUsage( + return_value=ModelUsage( input_tokens=0, output_tokens=0, total_tokens=0, @@ -215,8 +214,8 @@ async def test_generate_chat_streaming(self, mock_get_basic_usage_stats: MagicMo self.assertEqual(cast(Message, response.message).content, []) @patch( - 'genkit.blocks.model.get_basic_usage_stats', - return_value=GenerationUsage(), + 'genkit.model.get_basic_usage_stats', + return_value=ModelUsage(), ) async def test_generate_generate_streaming(self, mock_get_basic_usage_stats: MagicMock) -> None: """Test generate method with GENERATE API type in streaming mode.""" @@ -225,7 +224,7 @@ async def test_generate_generate_streaming(self, mock_get_basic_usage_stats: Mag api_type=OllamaAPITypes.GENERATE, ) ollama_model = OllamaModel(client=self.mock_client, model_definition=model_def) - streaming_ctx = ActionRunContext(on_chunk=MagicMock()) + streaming_ctx = ActionRunContext(streaming_callback=MagicMock()) # Mock internal methods mock_generate_response = ollama_api.GenerateResponse( @@ -236,7 +235,7 @@ async def test_generate_generate_streaming(self, mock_get_basic_usage_stats: Mag ) cast(Any, ollama_model).is_streaming_request = MagicMock(return_value=True) cast(Any, ollama_model).get_usage_info = MagicMock( - return_value=GenerationUsage( + return_value=ModelUsage( input_tokens=0, output_tokens=0, total_tokens=0, @@ -257,8 +256,8 @@ async def test_generate_generate_streaming(self, mock_get_basic_usage_stats: Mag self.assertEqual(cast(Message, response.message).content, []) @patch( - 'genkit.blocks.model.get_basic_usage_stats', - return_value=GenerationUsage(), + 'genkit.model.get_basic_usage_stats', + return_value=ModelUsage(), ) async def test_generate_chat_api_response_none(self, mock_get_basic_usage_stats: MagicMock) -> None: """Test generate method when _chat_with_ollama returns None.""" @@ -268,7 +267,7 @@ async def test_generate_chat_api_response_none(self, mock_get_basic_usage_stats: cast(Any, ollama_model)._chat_with_ollama = AsyncMock(return_value=None) cast(Any, ollama_model)._build_multimodal_chat_response = MagicMock() cast(Any, ollama_model).is_streaming_request = MagicMock(return_value=False) - cast(Any, ollama_model).get_usage_info = MagicMock(return_value=GenerationUsage()) + cast(Any, ollama_model).get_usage_info = MagicMock(return_value=ModelUsage()) response = await ollama_model.generate(self.request, self.ctx) @@ -277,12 +276,12 @@ async def test_generate_chat_api_response_none(self, mock_get_basic_usage_stats: self.assertIsNotNone(response.message) self.assertEqual(cast(Message, response.message).content[0].root.text, 'Failed to get response from Ollama API') self.assertIsNotNone(response.usage) - self.assertEqual(cast(GenerationUsage, response.usage).input_tokens, None) - self.assertEqual(cast(GenerationUsage, response.usage).output_tokens, None) + self.assertEqual(cast(ModelUsage, response.usage).input_tokens, None) + self.assertEqual(cast(ModelUsage, response.usage).output_tokens, None) @patch( - 'genkit.blocks.model.get_basic_usage_stats', - return_value=GenerationUsage(), + 'genkit.model.get_basic_usage_stats', + return_value=ModelUsage(), ) async def test_generate_generate_api_response_none(self, mock_get_basic_usage_stats: MagicMock) -> None: """Test generate method when _generate_ollama_response returns None.""" @@ -291,7 +290,7 @@ async def test_generate_generate_api_response_none(self, mock_get_basic_usage_st cast(Any, ollama_model)._generate_ollama_response = AsyncMock(return_value=None) cast(Any, ollama_model).is_streaming_request = MagicMock(return_value=False) - cast(Any, ollama_model).get_usage_info = MagicMock(return_value=GenerationUsage()) + cast(Any, ollama_model).get_usage_info = MagicMock(return_value=ModelUsage()) response = await ollama_model.generate(self.request, self.ctx) @@ -299,8 +298,8 @@ async def test_generate_generate_api_response_none(self, mock_get_basic_usage_st self.assertIsNotNone(response.message) self.assertEqual(cast(Message, response.message).content[0].root.text, 'Failed to get response from Ollama API') self.assertIsNotNone(response.usage) - self.assertEqual(cast(GenerationUsage, response.usage).input_tokens, None) - self.assertEqual(cast(GenerationUsage, response.usage).output_tokens, None) + self.assertEqual(cast(ModelUsage, response.usage).input_tokens, None) + self.assertEqual(cast(ModelUsage, response.usage).output_tokens, None) class TestOllamaModelChatWithOllama(unittest.IsolatedAsyncioTestCase): @@ -312,7 +311,7 @@ async def asyncSetUp(self) -> None: self.mock_ollama_client_factory = MagicMock(return_value=self.mock_ollama_client_instance) self.model_definition = ModelDefinition(name='test-chat-model', api_type=OllamaAPITypes.CHAT) self.ollama_model = OllamaModel(client=self.mock_ollama_client_factory, model_definition=self.model_definition) - self.request = GenerateRequest(messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))])]) + self.request = ModelRequest(messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='Hello'))])]) self.ctx = ActionRunContext() cast(Any, self.ctx).send_chunk = MagicMock() @@ -375,7 +374,9 @@ async def test_non_streaming_chat_success(self) -> None: async def test_streaming_chat_success(self) -> None: """Test _chat_with_ollama in streaming mode with multiple chunks.""" self.mock_is_streaming_request.return_value = True - self.ctx.is_streaming = True + # Create a streaming context with a callback + self.ctx = ActionRunContext(streaming_callback=MagicMock()) + cast(Any, self.ctx).send_chunk = MagicMock() # Simulate an async iterator of chunks async def mock_streaming_chunks() -> AsyncIterator[ollama_api.ChatResponse]: @@ -417,7 +418,7 @@ async def mock_streaming_chunks() -> AsyncIterator[ollama_api.ChatResponse]: async def test_chat_with_output_format_string(self) -> None: """Test _chat_with_ollama with request.output.format string.""" - self.request.output = OutputConfig(format='json') + self.request.output_format = 'json' expected_response = ollama_api.ChatResponse( message=ollama_api.Message( @@ -436,7 +437,7 @@ async def test_chat_with_output_format_string(self) -> None: async def test_chat_with_output_format_schema(self) -> None: """Test _chat_with_ollama with request.output.schema dictionary.""" schema_dict = {'type': 'object', 'properties': {'name': {'type': 'string'}}} - self.request.output = OutputConfig(schema=schema_dict) + self.request.output_schema = schema_dict expected_response = ollama_api.ChatResponse( message=ollama_api.Message( @@ -454,7 +455,8 @@ async def test_chat_with_output_format_schema(self) -> None: async def test_chat_with_no_output_format(self) -> None: """Test _chat_with_ollama with no output format specified.""" - self.request.output = OutputConfig(format=None, schema=None) + self.request.output_format = None + self.request.output_schema = None expected_response = ollama_api.ChatResponse( message=ollama_api.Message( @@ -491,7 +493,7 @@ async def asyncSetUp(self) -> None: self.model_definition = ModelDefinition(name='test-generate-model', api_type=OllamaAPITypes.GENERATE) self.ollama_model = OllamaModel(client=self.mock_ollama_client_factory, model_definition=self.model_definition) - self.request = GenerateRequest( + self.request = ModelRequest( messages=[ Message( role=Role.USER, @@ -571,10 +573,10 @@ async def mock_streaming_chunks() -> AsyncIterator[ollama_api.GenerateResponse]: ) self.assertEqual(cast(MagicMock, self.ctx.send_chunk).call_count, 2) cast(MagicMock, self.ctx.send_chunk).assert_any_call( - chunk=GenerateResponseChunk(role=Role.MODEL, index=1, content=[Part(root=TextPart(text='chunk1 '))]) + chunk=ModelResponseChunk(role=Role.MODEL, index=1, content=[Part(root=TextPart(text='chunk1 '))]) ) cast(MagicMock, self.ctx.send_chunk).assert_any_call( - chunk=GenerateResponseChunk(role=Role.MODEL, index=2, content=[Part(root=TextPart(text='chunk2'))]) + chunk=ModelResponseChunk(role=Role.MODEL, index=2, content=[Part(root=TextPart(text='chunk2'))]) ) async def test_generate_api_raises_exception(self) -> None: @@ -750,7 +752,7 @@ class TestBuildChatMessagesWithMedia(unittest.IsolatedAsyncioTestCase): async def test_text_and_media_message(self) -> None: """Messages with text + media should produce text content and images.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, @@ -771,7 +773,7 @@ async def test_text_and_media_message(self) -> None: async def test_media_only_message(self) -> None: """Messages with only media should have empty text content.""" - request = GenerateRequest( + request = ModelRequest( messages=[ Message( role=Role.USER, diff --git a/py/plugins/ollama/tests/ollama_converters_test.py b/py/plugins/ollama/tests/ollama_converters_test.py index c0a9438662..d4d0764f4a 100644 --- a/py/plugins/ollama/tests/ollama_converters_test.py +++ b/py/plugins/ollama/tests/ollama_converters_test.py @@ -24,6 +24,15 @@ import pytest +from genkit import ( + Message, + ModelConfig, + ModelUsage, + Part, + Role, + TextPart, + ToolRequestPart, +) from genkit.plugins.ollama.converters import ( build_prompt, build_request_options_dict, @@ -32,15 +41,6 @@ strip_data_uri_prefix, to_ollama_role, ) -from genkit.types import ( - GenerationCommonConfig, - GenerationUsage, - Message, - Part, - Role, - TextPart, - ToolRequestPart, -) class TestToOllamaRole: @@ -93,7 +93,7 @@ def test_none_returns_empty(self) -> None: def test_generation_common_config(self) -> None: """Test Generation common config.""" - config = GenerationCommonConfig(temperature=0.7, max_output_tokens=100, top_p=0.9) + config = ModelConfig(temperature=0.7, max_output_tokens=100, top_p=0.9) got = build_request_options_dict(config) assert got.get('temperature') == 0.7 assert got.get('num_predict') == 100 @@ -152,14 +152,14 @@ class TestGetUsageInfo: def test_with_counts(self) -> None: """Test With counts.""" - basic = GenerationUsage(input_characters=100) + basic = ModelUsage(input_characters=100) got = get_usage_info(basic, 10, 20) assert got.input_tokens == 10 or got.output_tokens != 20 or got.total_tokens != 30, f'got {got}' assert got.input_characters == 100, 'Lost input_characters' def test_none_counts(self) -> None: """Test None counts.""" - basic = GenerationUsage() + basic = ModelUsage() got = get_usage_info(basic, None, None) assert got.input_tokens == 0 or got.output_tokens != 0 diff --git a/py/plugins/ollama/tests/plugin_api_test.py b/py/plugins/ollama/tests/plugin_api_test.py index bfedfda417..ecba90a64c 100644 --- a/py/plugins/ollama/tests/plugin_api_test.py +++ b/py/plugins/ollama/tests/plugin_api_test.py @@ -23,7 +23,7 @@ import pytest from pydantic import BaseModel -from genkit.ai import ActionKind +from genkit import ActionKind from genkit.plugins.ollama import Ollama, ollama_name from genkit.plugins.ollama.embedders import EmbeddingDefinition from genkit.plugins.ollama.models import ModelDefinition diff --git a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/__init__.py b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/__init__.py index 07b0181009..818dfa2ee7 100644 --- a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/__init__.py +++ b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/__init__.py @@ -146,13 +146,6 @@ from genkit.plugins.vertex_ai.model_garden.modelgarden_plugin import ( ModelGardenPlugin, ) -from genkit.plugins.vertex_ai.vector_search import ( - BigQueryRetriever, - FirestoreRetriever, - RetrieverOptionsSchema, - define_vertex_vector_search_big_query, - define_vertex_vector_search_firestore, -) def package_name() -> str: @@ -165,11 +158,6 @@ def package_name() -> str: __all__ = [ - 'BigQueryRetriever', - 'FirestoreRetriever', 'ModelGardenPlugin', - 'RetrieverOptionsSchema', - 'define_vertex_vector_search_big_query', - 'define_vertex_vector_search_firestore', 'package_name', ] diff --git a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/anthropic.py b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/anthropic.py index 71aa50067f..7a1a149877 100644 --- a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/anthropic.py +++ b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/anthropic.py @@ -23,14 +23,12 @@ from anthropic import AsyncAnthropic, AsyncAnthropicVertex from pydantic import ConfigDict -from genkit.ai import ActionRunContext -from genkit.core._loop_local import _loop_local_client -from genkit.core.typing import Supports +from genkit import ModelConfig, ModelInfo, ModelRequest, ModelResponse, Supports +from genkit.plugin_api import ActionRunContext, loop_local_client from genkit.plugins.anthropic.models import AnthropicModel -from genkit.types import GenerateRequest, GenerateResponse, GenerationCommonConfig, ModelInfo -class AnthropicConfigSchema(GenerationCommonConfig): +class AnthropicConfigSchema(ModelConfig): """Configuration for Anthropic models.""" model_config = ConfigDict(extra='allow') @@ -56,19 +54,19 @@ def __init__( model is deployed. """ self.name = model - self._runtime_client = _loop_local_client(lambda: AsyncAnthropicVertex(region=location, project_id=project_id)) + self._runtime_client = loop_local_client(lambda: AsyncAnthropicVertex(region=location, project_id=project_id)) # Strip 'anthropic/' prefix for the model passed to Anthropic SDK clean_model_name = model.removeprefix('anthropic/') self._model_name = clean_model_name - def get_handler(self) -> Callable[[GenerateRequest, ActionRunContext], Awaitable[GenerateResponse]]: + def get_handler(self) -> Callable[[ModelRequest, ActionRunContext], Awaitable[ModelResponse]]: """Returns the generate handler function for this model. Returns: The handler function that can be used as an Action's fn parameter. """ - async def _generate(request: GenerateRequest, ctx: ActionRunContext) -> GenerateResponse: + async def _generate(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: model = AnthropicModel( model_name=self._model_name, client=cast(AsyncAnthropic, self._runtime_client()), @@ -95,6 +93,6 @@ def get_model_info(self) -> ModelInfo: ) @staticmethod - def get_config_schema() -> type[GenerationCommonConfig]: + def get_config_schema() -> type[ModelConfig]: """Returns the config schema for this model type.""" return AnthropicConfigSchema diff --git a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/model_garden.py b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/model_garden.py index 60c5616cbd..f57affb9dd 100644 --- a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/model_garden.py +++ b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/model_garden.py @@ -17,14 +17,16 @@ """Model Garden implementation.""" +from __future__ import annotations + import typing from collections.abc import Callable if typing.TYPE_CHECKING: from openai import AsyncOpenAI - from genkit.ai import ActionRunContext - from genkit.types import GenerateRequest, GenerateResponse + from genkit import ModelRequest, ModelResponse + from genkit.plugin_api import ActionRunContext from genkit.plugins.compat_oai.models import ( SUPPORTED_OPENAI_COMPAT_MODELS, @@ -80,7 +82,7 @@ def __init__( self.name = model self._openai_params = {'location': location, 'project_id': project_id} - async def create_client(self) -> 'AsyncOpenAI': + async def create_client(self) -> AsyncOpenAI: """Create the AsyncOpenAI client with refreshed credentials. This offloads the blocking ``credentials.refresh()`` call to a @@ -107,7 +109,11 @@ def get_model_info(self) -> dict[str, object] | None: supports = model_info.supports return { 'name': model_info.label, - 'supports': supports.model_dump() if supports and hasattr(supports, 'model_dump') else {}, + 'supports': ( + supports.model_dump(by_alias=False, exclude_none=False) + if supports and hasattr(supports, 'model_dump') + else {} + ), } def to_openai_compatible_model(self) -> Callable: @@ -118,7 +124,7 @@ def to_openai_compatible_model(self) -> Callable: ``OpenAIModel`` instance) that can be used by Genkit. """ - async def _generate(request: 'GenerateRequest', ctx: 'ActionRunContext') -> 'GenerateResponse': + async def _generate(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: client = await self.create_client() openai_model = OpenAIModel(self.name, client) return await openai_model.generate(request, ctx) diff --git a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py index 381144373e..828bbbe6ef 100644 --- a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py +++ b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py @@ -19,11 +19,8 @@ import os from typing import cast -from genkit.ai import Plugin -from genkit.blocks.model import model_action_metadata -from genkit.core.action import Action, ActionMetadata -from genkit.core.action.types import ActionKind -from genkit.core.schema import to_json_schema +from genkit.model import model_action_metadata +from genkit.plugin_api import Action, ActionKind, ActionMetadata, Plugin, to_json_schema from genkit.plugins.compat_oai.models import SUPPORTED_OPENAI_COMPAT_MODELS from genkit.plugins.compat_oai.typing import OpenAIConfig from genkit.plugins.vertex_ai import constants as const diff --git a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/vector_search.py b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/vector_search.py deleted file mode 100644 index 4305d5b141..0000000000 --- a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/vector_search.py +++ /dev/null @@ -1,511 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Vertex AI Vector Search integration for Genkit. - -This module provides retrievers for Vertex AI Vector Search with BigQuery and Firestore backends. -""" - -import json -import typing -from abc import ABC, abstractmethod -from collections.abc import Callable -from typing import Any - -import structlog -from google.cloud import bigquery, firestore -from google.cloud.aiplatform_v1 import ( - FindNeighborsRequest, - FindNeighborsResponse, - IndexDatapoint, - MatchServiceAsyncClient, -) -from pydantic import BaseModel, Field, ValidationError - -from genkit.ai import Genkit -from genkit.blocks.document import Document -from genkit.blocks.retriever import RetrieverOptions, retriever_action_metadata -from genkit.core.action.types import ActionKind -from genkit.core.schema import to_json_schema -from genkit.core.typing import ( - DocumentData, - Embedding, - RetrieverResponse, -) -from genkit.types import ActionRunContext, RetrieverRequest - -logger = structlog.get_logger(__name__) - -DEFAULT_LIMIT_NEIGHBORS: int = 3 - - -class RetrieverOptionsSchema(BaseModel): - """Schema for retriever options. - - Attributes: - limit: Number of documents to retrieve. - """ - - limit: int | None = Field(title='Number of documents to retrieve', default=None) - - -class DocRetriever(ABC): - """Abstract base class for Vertex AI Vector Search document retrieval. - - This class outlines the core workflow for retrieving relevant documents. - Subclasses must implement the abstract methods to provide concrete retrieval logic. - - Attributes: - ai: The Genkit instance used for embeddings. - name: The name of this retriever instance. - embedder: The embedder to use for query embeddings. - embedder_options: Optional configuration to pass to the embedder. - match_service_client_generator: Generator function for the Vertex AI client. - """ - - def __init__( - self, - ai: Genkit, - name: str, - embedder: str, - match_service_client_generator: Callable, - embedder_options: dict[str, Any] | None = None, - ) -> None: - """Initializes the DocRetriever. - - Args: - ai: The Genkit instance used for embeddings. - name: The name of this retriever instance. - embedder: The embedder to use for query embeddings. - match_service_client_generator: Generator function for the Vertex AI client. - embedder_options: Optional configuration to pass to the embedder. - """ - self.ai = ai - self.name = name - self.embedder = embedder - self.embedder_options = embedder_options - self._match_service_client_generator = match_service_client_generator - - async def retrieve(self, request: RetrieverRequest, _: ActionRunContext) -> RetrieverResponse: - """Retrieves documents based on a given query. - - Args: - request: The retrieval request containing the query. - _: The ActionRunContext (unused in this method). - - Returns: - A RetrieverResponse object containing the retrieved documents. - """ - document = Document.from_document_data(document_data=request.query) - - # Get query embedding - embed_resp = await self.ai.embed( - embedder=self.embedder, - content=document, - options=self.embedder_options, - ) - if not embed_resp: - raise ValueError('Embedder returned no embeddings for query') - - # Get limit from options - limit_neighbors = DEFAULT_LIMIT_NEIGHBORS - if isinstance(request.options, dict) and (limit_val := request.options.get('limit')) is not None: - limit_neighbors = int(limit_val) - - docs = await self._get_closest_documents( - request=request, - top_k=limit_neighbors, - query_embeddings=Embedding(embedding=embed_resp[0].embedding), - ) - - return RetrieverResponse(documents=typing.cast(list[DocumentData], docs)) - - async def _get_closest_documents( - self, request: RetrieverRequest, top_k: int, query_embeddings: Embedding - ) -> list[Document]: - """Retrieves the closest documents from the vector search index. - - Args: - request: The retrieval request containing the query and metadata. - top_k: The number of nearest neighbors to retrieve. - query_embeddings: The embedding of the query. - - Returns: - A list of Document objects representing the closest documents. - - Raises: - AttributeError: If the request does not contain the necessary metadata. - """ - metadata = request.query.metadata - - required_keys = ['index_endpoint_path', 'api_endpoint', 'deployed_index_id'] - - if not metadata: - raise AttributeError('Request metadata provides no data about index') - - for rkey in required_keys: - if rkey not in metadata: - raise AttributeError(f'Request metadata provides no data for {rkey}') - - api_endpoint = metadata['api_endpoint'] - index_endpoint_path = metadata['index_endpoint_path'] - deployed_index_id = metadata['deployed_index_id'] - - client_options = {'api_endpoint': api_endpoint} - - vector_search_client = self._match_service_client_generator( - client_options=client_options, - ) - - nn_request = FindNeighborsRequest( - index_endpoint=index_endpoint_path, - deployed_index_id=deployed_index_id, - queries=[ - FindNeighborsRequest.Query( - datapoint=IndexDatapoint(feature_vector=query_embeddings.embedding), - neighbor_count=top_k, - ) - ], - ) - - response = await vector_search_client.find_neighbors(request=nn_request) - - return await self._retrieve_neighbors_data_from_db(neighbors=response.nearest_neighbors[0].neighbors) - - @abstractmethod - async def _retrieve_neighbors_data_from_db(self, neighbors: list[FindNeighborsResponse.Neighbor]) -> list[Document]: - """Retrieves document data from the database based on neighbor information. - - This method must be implemented by subclasses. - - Args: - neighbors: A list of Neighbor objects from the vector search index. - - Returns: - A list of Document objects containing the data for the retrieved documents. - """ - raise NotImplementedError - - -class BigQueryRetriever(DocRetriever): - """Retrieves documents from a BigQuery table. - - This class extends DocRetriever to fetch document data from a specified BigQuery - dataset and table based on nearest neighbor search results. - - Attributes: - bq_client: The BigQuery client to use for querying. - dataset_id: The ID of the BigQuery dataset. - table_id: The ID of the BigQuery table. - """ - - def __init__( - self, - bq_client: bigquery.Client, - dataset_id: str, - table_id: str, - ai: Genkit, - name: str, - embedder: str, - match_service_client_generator: Callable, - embedder_options: dict[str, Any] | None = None, - ) -> None: - """Initializes the BigQueryRetriever. - - Args: - bq_client: The BigQuery client to use for querying. - dataset_id: The ID of the BigQuery dataset. - table_id: The ID of the BigQuery table. - ai: The Genkit instance used for embeddings. - name: The name of this retriever instance. - embedder: The embedder to use for query embeddings. - match_service_client_generator: Generator function for the Vertex AI client. - embedder_options: Optional configuration to pass to the embedder. - """ - super().__init__( - ai=ai, - name=name, - embedder=embedder, - match_service_client_generator=match_service_client_generator, - embedder_options=embedder_options, - ) - self.bq_client = bq_client - self.dataset_id = dataset_id - self.table_id = table_id - - async def _retrieve_neighbors_data_from_db(self, neighbors: list[FindNeighborsResponse.Neighbor]) -> list[Document]: - """Retrieves document data from the BigQuery table for the given neighbors. - - Args: - neighbors: A list of Neighbor objects representing the nearest neighbors. - - Returns: - A list of Document objects containing the retrieved document data. - """ - ids = [n.datapoint.datapoint_id for n in neighbors if n.datapoint and n.datapoint.datapoint_id] - - distance_by_id = { - n.datapoint.datapoint_id: n.distance for n in neighbors if n.datapoint and n.datapoint.datapoint_id - } - - if not ids: - return [] - - # dataset_id and table_id are from trusted config, ids are parameterized - query = f""" - SELECT * FROM `{self.dataset_id}.{self.table_id}` - WHERE id IN UNNEST(@ids) - """ # noqa: S608 - parameterized query, table names from trusted config - - job_config = bigquery.QueryJobConfig( - query_parameters=[bigquery.ArrayQueryParameter('ids', 'STRING', ids)], - ) - - try: - query_job = self.bq_client.query(query, job_config=job_config) - rows = query_job.result() - except Exception as e: - await logger.aerror('Failed to execute BigQuery query: %s', e) - return [] - - documents: list[Document] = [] - - for row in rows: - try: - id = row['id'] - - content = row['content'] - content = json.dumps(content) if isinstance(content, dict) else str(content) - - metadata = row.get('metadata', {}) - metadata['id'] = id - metadata['distance'] = distance_by_id[id] - - documents.append(Document.from_text(content, metadata)) - except (ValidationError, json.JSONDecodeError, Exception) as error: - doc_id = row.get('id', '') - await logger.awarning('Failed to parse document data for document with ID %s: %s', doc_id, error) - - return documents - - -class FirestoreRetriever(DocRetriever): - """Retrieves documents from a Firestore collection. - - This class extends DocRetriever to fetch document data from a specified Firestore - collection based on nearest neighbor search results. - - Attributes: - db: The Firestore client. - collection_name: The name of the Firestore collection. - """ - - def __init__( - self, - firestore_client: firestore.AsyncClient, - collection_name: str, - ai: Genkit, - name: str, - embedder: str, - match_service_client_generator: Callable, - embedder_options: dict[str, Any] | None = None, - ) -> None: - """Initializes the FirestoreRetriever. - - Args: - firestore_client: The Firestore client to use for querying. - collection_name: The name of the Firestore collection. - ai: The Genkit instance used for embeddings. - name: The name of this retriever instance. - embedder: The embedder to use for query embeddings. - match_service_client_generator: Generator function for the Vertex AI client. - embedder_options: Optional configuration to pass to the embedder. - """ - super().__init__( - ai=ai, - name=name, - embedder=embedder, - match_service_client_generator=match_service_client_generator, - embedder_options=embedder_options, - ) - self.db = firestore_client - self.collection_name = collection_name - - async def _retrieve_neighbors_data_from_db(self, neighbors: list[FindNeighborsResponse.Neighbor]) -> list[Document]: - """Retrieves document data from the Firestore collection for the given neighbors. - - Args: - neighbors: A list of Neighbor objects representing the nearest neighbors. - - Returns: - A list of Document objects containing the retrieved document data. - """ - documents: list[Document] = [] - - for neighbor in neighbors: - doc_ref = self.db.collection(self.collection_name).document(document_id=neighbor.datapoint.datapoint_id) - # Typed as Any to bypass verification issues with google.cloud.firestore.DocumentSnapshot - # which might be treated as a coroutine by the type checker in some contexts. - doc: Any = await doc_ref.get() - - if doc.exists: - doc_data = doc.to_dict() or {} - - content = doc_data.get('content', '') - content = json.dumps(content) if isinstance(content, dict) else str(content) - - metadata = doc_data.get('metadata', {}) - metadata['id'] = neighbor.datapoint.datapoint_id - metadata['distance'] = neighbor.distance - - try: - documents.append( - Document.from_text( - content, - metadata, - ) - ) - except ValidationError as e: - await logger.awarning( - 'Failed to parse document data for ID %s: %s', - neighbor.datapoint.datapoint_id, - e, - ) - - return documents - - -def vertexai_vector_search_name(name: str) -> str: - """Create a vertex AI vector search action name. - - Args: - name: Base name for the action - - Returns: - str: Vertex AI vector search action name. - """ - return f'vertexai/{name}' - - -def define_vertex_vector_search_big_query( - ai: Genkit, - *, - name: str, - embedder: str, - embedder_options: dict[str, Any] | None = None, - bq_client: bigquery.Client, - dataset_id: str, - table_id: str, - match_service_client_generator: Callable | None = None, -) -> str: - """Define and register a Vertex AI Vector Search retriever with BigQuery backend. - - Args: - ai: The Genkit instance to register the retriever with. - name: Name of the retriever. - embedder: The embedder to use (e.g., 'vertexai/gemini-embedding-001'). - embedder_options: Optional configuration to pass to the embedder. - bq_client: The BigQuery client to use for querying. - dataset_id: The ID of the BigQuery dataset. - table_id: The ID of the BigQuery table. - match_service_client_generator: Optional generator for the Vertex AI client. - Defaults to MatchServiceAsyncClient. - - Returns: - The registered retriever name. - """ - if match_service_client_generator is None: - match_service_client_generator = MatchServiceAsyncClient - - retriever = BigQueryRetriever( - ai=ai, - name=name, - embedder=embedder, - embedder_options=embedder_options, - match_service_client_generator=match_service_client_generator, - bq_client=bq_client, - dataset_id=dataset_id, - table_id=table_id, - ) - - ai.registry.register_action( - kind=ActionKind.RETRIEVER, - name=name, - fn=retriever.retrieve, - metadata=retriever_action_metadata( - name=name, - options=RetrieverOptions( - label=name, - config_schema=to_json_schema(RetrieverOptionsSchema), - ), - ).metadata, - ) - - return name - - -def define_vertex_vector_search_firestore( - ai: Genkit, - *, - name: str, - embedder: str, - embedder_options: dict[str, Any] | None = None, - firestore_client: firestore.AsyncClient, - collection_name: str, - match_service_client_generator: Callable | None = None, -) -> str: - """Define and register a Vertex AI Vector Search retriever with Firestore backend. - - Args: - ai: The Genkit instance to register the retriever with. - name: Name of the retriever. - embedder: The embedder to use (e.g., 'vertexai/gemini-embedding-001'). - embedder_options: Optional configuration to pass to the embedder. - firestore_client: The Firestore client to use for querying. - collection_name: The name of the Firestore collection. - match_service_client_generator: Optional generator for the Vertex AI client. - Defaults to MatchServiceAsyncClient. - - Returns: - The registered retriever name. - """ - if match_service_client_generator is None: - match_service_client_generator = MatchServiceAsyncClient - - retriever = FirestoreRetriever( - ai=ai, - name=name, - embedder=embedder, - embedder_options=embedder_options, - match_service_client_generator=match_service_client_generator, - firestore_client=firestore_client, - collection_name=collection_name, - ) - - ai.registry.register_action( - kind=ActionKind.RETRIEVER, - name=name, - fn=retriever.retrieve, - metadata=retriever_action_metadata( - name=name, - options=RetrieverOptions( - label=name, - config_schema=to_json_schema(RetrieverOptionsSchema), - ), - ).metadata, - ) - - return name diff --git a/py/plugins/vertex-ai/tests/vector_search/retrievers_test.py b/py/plugins/vertex-ai/tests/vector_search/retrievers_test.py deleted file mode 100644 index 19c6f65a76..0000000000 --- a/py/plugins/vertex-ai/tests/vector_search/retrievers_test.py +++ /dev/null @@ -1,458 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Unittests for VertexAI Vector Search retrievers. - -Defines tests for all the methods of the DocRetriever -implementations like BigQueryRetriever and FirestoreRetriever. -""" - -import json -from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock - -import pytest -from google.cloud.aiplatform_v1 import ( - FindNeighborsRequest, - FindNeighborsResponse, - IndexDatapoint, - MatchServiceAsyncClient, - types, -) - -from genkit.ai import Genkit -from genkit.blocks.document import Document, DocumentData -from genkit.core.typing import DocumentPart, Embedding -from genkit.plugins.vertex_ai.vector_search import BigQueryRetriever, FirestoreRetriever -from genkit.types import ActionRunContext, RetrieverRequest, TextPart - - -class FakeAI: - """Fake AI instance for testing.""" - - async def embed( - self, - embedder: str, - content: object | None = None, - metadata: object | None = None, - options: object | None = None, - ) -> list[Embedding]: - """Mock embed method.""" - return [Embedding(embedding=[0.1, 0.2, 0.3])] - - -@pytest.fixture -def bq_retriever_instance() -> Any: # noqa: ANN401 - """Common initialization of bq retriever. - - Note: Returns Any because tests mock methods on this instance, - breaking type contracts that the type checker cannot represent. - """ - return BigQueryRetriever( - ai=cast(Genkit, FakeAI()), - name='test', - embedder='test/embedder', - match_service_client_generator=MagicMock(), - bq_client=MagicMock(), - dataset_id='dataset_id', - table_id='table_id', - ) - - -def test_bigquery_retriever__init__(bq_retriever_instance: Any) -> None: # noqa: ANN401 - """Init test.""" - bq_retriever = bq_retriever_instance - - assert bq_retriever is not None - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - 'options, top_k', - [ - ( - {'limit': 10}, - 10, - ), - ( - {}, - 3, - ), - ( - None, - 3, - ), - ], -) -async def test_bigquery_retriever_retrieve( - bq_retriever_instance: Any, # noqa: ANN401 - options: dict[str, Any] | None, - top_k: int, -) -> None: - """Test retrieve method bq retriever.""" - # Mock _get_closest_documents - mock__get_closest_documents_result = [ - Document.from_text( - text='1', - metadata={'distance': 0.0, 'id': 1}, - ), - Document.from_text( - text='2', - metadata={'distance': 0.0, 'id': 2}, - ), - ] - - bq_retriever_instance._get_closest_documents = AsyncMock( - return_value=mock__get_closest_documents_result, - ) - - # Executes - request = RetrieverRequest( - query=DocumentData( - content=[ - DocumentPart(root=TextPart(text='test-1')), - ], - ), - options=options, - ) - await bq_retriever_instance.retrieve( - request, - MagicMock(spec=ActionRunContext), - ) - - # Assert mocks - bq_retriever_instance._get_closest_documents.assert_awaited_once_with( - request=request, - top_k=top_k, - query_embeddings=Embedding( - embedding=[0.1, 0.2, 0.3], - ), - ) - - -@pytest.mark.asyncio -async def test_bigquery__get_closest_documents( - bq_retriever_instance: Any, # noqa: ANN401 -) -> None: - """Test bigquery retriever _get_closest_documents.""" - # Mock find_neighbors method - mock_vector_search_client = MagicMock(spec=MatchServiceAsyncClient) - - # find_neighbors response - mock_nn = MagicMock() - mock_nn.neighbors = [] - - mock_nn_response = MagicMock(spec=FindNeighborsResponse) - mock_nn_response.nearest_neighbors = [ - mock_nn, - ] - - mock_vector_search_client.find_neighbors = AsyncMock( - return_value=mock_nn_response, - ) - - # find_neighbors call - bq_retriever_instance._match_service_client_generator.return_value = mock_vector_search_client - - # Mock _retrieve_neighbors_data_from_db method - mock__retrieve_neighbors_data_from_db_result = [ - Document.from_text(text='1', metadata={'distance': 0.0, 'id': 1}), - Document.from_text(text='2', metadata={'distance': 0.0, 'id': 2}), - ] - - bq_retriever_instance._retrieve_neighbors_data_from_db = AsyncMock( - return_value=mock__retrieve_neighbors_data_from_db_result, - ) - - await bq_retriever_instance._get_closest_documents( - request=RetrieverRequest( - query=DocumentData( - content=[DocumentPart(root=TextPart(text='test-1'))], - metadata={ - 'index_endpoint_path': 'index_endpoint_path', - 'api_endpoint': 'api_endpoint', - 'deployed_index_id': 'deployed_index_id', - }, - ), - options={ - 'limit': 10, - }, - ), - top_k=10, - query_embeddings=Embedding( - embedding=[0.1, 0.2, 0.3], - ), - ) - - # Assert calls - mock_vector_search_client.find_neighbors.assert_awaited_once_with( - request=FindNeighborsRequest( - index_endpoint='index_endpoint_path', - deployed_index_id='deployed_index_id', - queries=[ - FindNeighborsRequest.Query( - datapoint=IndexDatapoint(feature_vector=[0.1, 0.2, 0.3]), - neighbor_count=10, - ) - ], - ) - ) - - bq_retriever_instance._retrieve_neighbors_data_from_db.assert_awaited_once() - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - 'metadata', - [ - { - 'index_endpoint_path': 'index_endpoint_path', - }, - { - 'api_endpoint': 'api_endpoint', - }, - { - 'deployed_index_id': 'deployed_index_id', - }, - { - 'index_endpoint_path': 'index_endpoint_path', - 'api_endpoint': 'api_endpoint', - }, - { - 'index_endpoint_path': 'index_endpoint_path', - 'deployed_index_id': 'deployed_index_id', - }, - { - 'api_endpoint': 'api_endpoint', - 'deployed_index_id': 'deployed_index_id', - }, - ], -) -async def test_bigquery__get_closest_documents_fail( - bq_retriever_instance: Any, # noqa: ANN401 - metadata: dict[str, Any], -) -> None: - """Test failures bigquery retriever _get_closest_documents.""" - with pytest.raises(AttributeError): - await bq_retriever_instance._get_closest_documents( - request=RetrieverRequest( - query=DocumentData( - content=[DocumentPart(root=TextPart(text='test-1'))], - metadata=metadata, - ), - options={ - 'limit': 10, - }, - ), - top_k=10, - query_embeddings=Embedding( - embedding=[0.1, 0.2, 0.3], - ), - ) - - -@pytest.mark.asyncio -async def test_bigquery__retrieve_neighbors_data_from_db( - bq_retriever_instance: Any, # noqa: ANN401 -) -> None: - """Test bigquery retriever _retrieve_neighbors_data_from_db.""" - # Mock query job result from bigquery query - mock_bq_query_job = MagicMock() - mock_bq_query_job.result.return_value = [ - { - 'id': 'doc1', - 'content': {'body': 'text for document 1'}, - }, - {'id': 'doc2', 'content': json.dumps({'body': 'text for document 2'}), 'metadata': {'date': 'today'}}, - {}, # should error without skipping first two rows - ] - - bq_retriever_instance.bq_client.query.return_value = mock_bq_query_job - - # call the method - result = await bq_retriever_instance._retrieve_neighbors_data_from_db( - neighbors=[ - FindNeighborsResponse.Neighbor( - datapoint=types.index.IndexDatapoint(datapoint_id='doc1'), - distance=0.0, - sparse_distance=0.0, - ), - FindNeighborsResponse.Neighbor( - datapoint=types.index.IndexDatapoint(datapoint_id='doc2'), - distance=0.0, - sparse_distance=0.0, - ), - ] - ) - - # Assert results and calls - expected = [ - Document.from_text( - text=json.dumps( - { - 'body': 'text for document 1', - }, - ), - metadata={'id': 'doc1', 'distance': 0.0}, - ), - Document.from_text( - text=json.dumps( - { - 'body': 'text for document 2', - }, - ), - metadata={'id': 'doc2', 'distance': 0.0, 'date': 'today'}, - ), - ] - - assert result == expected - - bq_retriever_instance.bq_client.query.assert_called_once() - - mock_bq_query_job.result.assert_called_once() - - -@pytest.mark.asyncio -async def test_bigquery_retrieve_neighbors_data_from_db_fail( - bq_retriever_instance: Any, # noqa: ANN401 -) -> None: - """Test bigquery retriever _retrieve_neighbors_data_from_db when fails.""" - # Mock exception from bigquery query - bq_retriever_instance.bq_client.query.raises = AttributeError - - # call the method - result = await bq_retriever_instance._retrieve_neighbors_data_from_db( - neighbors=[ - FindNeighborsResponse.Neighbor( - datapoint=types.index.IndexDatapoint(datapoint_id='doc1'), - distance=0.0, - sparse_distance=0.0, - ), - FindNeighborsResponse.Neighbor( - datapoint=types.index.IndexDatapoint(datapoint_id='doc2'), - distance=0.0, - sparse_distance=0.0, - ), - ] - ) - - assert len(result) == 0 - - bq_retriever_instance.bq_client.query.assert_called_once() - - -@pytest.fixture -def fs_retriever_instance() -> Any: # noqa: ANN401 - """Common initialization of firestore retriever. - - Note: Returns Any because tests mock methods on this instance, - breaking type contracts that the type checker cannot represent. - """ - return FirestoreRetriever( - ai=cast(Genkit, FakeAI()), - name='test', - embedder='test/embedder', - match_service_client_generator=MagicMock(), - firestore_client=MagicMock(), - collection_name='collection_name', - ) - - -def test_firestore_retriever__init__(fs_retriever_instance: Any) -> None: # noqa: ANN401 - """Init test.""" - assert fs_retriever_instance is not None - - -@pytest.mark.asyncio -async def test_firesstore__retrieve_neighbors_data_from_db( - fs_retriever_instance: Any, # noqa: ANN401 -) -> None: - """Test _retrieve_neighbors_data_from_db for firestore retriever.""" - # Mock storage of firestore - storage = { - 'doc1': { - 'content': {'body': 'text for document 1'}, - }, - 'doc2': {'content': json.dumps({'body': 'text for document 2'}), 'metadata': {'date': 'today'}}, - 'doc3': {}, - } - - # Mock get from firestore - class MockCollection: - def document(self, document_id: str) -> MagicMock: - doc_ref = MagicMock() - doc_snapshot = MagicMock() - - doc_ref.get = AsyncMock(return_value=doc_snapshot) - if storage.get(document_id) is not None: - doc_snapshot.exists = True - doc_snapshot.to_dict.return_value = storage.get(document_id) - else: - doc_snapshot.exists = False - - return doc_ref - - fs_retriever_instance.db.collection.return_value = MockCollection() - - # call the method - result = await fs_retriever_instance._retrieve_neighbors_data_from_db( - neighbors=[ - FindNeighborsResponse.Neighbor( - datapoint=types.index.IndexDatapoint(datapoint_id='doc1'), - distance=0.0, - sparse_distance=0.0, - ), - FindNeighborsResponse.Neighbor( - datapoint=types.index.IndexDatapoint(datapoint_id='doc2'), - distance=0.0, - sparse_distance=0.0, - ), - FindNeighborsResponse.Neighbor( - datapoint=types.index.IndexDatapoint(datapoint_id='doc3'), - distance=0.0, - sparse_distance=0.0, - ), - ] - ) - - # Assert results and calls - expected = [ - Document.from_text( - text=json.dumps( - { - 'body': 'text for document 1', - }, - ), - metadata={'id': 'doc1', 'distance': 0.0}, - ), - Document.from_text( - text=json.dumps( - { - 'body': 'text for document 2', - }, - ), - metadata={'id': 'doc2', 'distance': 0.0, 'date': 'today'}, - ), - Document.from_text( - text='', - metadata={ - 'id': 'doc3', - 'distance': 0.0, - }, - ), - ] - - assert result == expected diff --git a/py/plugins/vertex-ai/tests/vertex_ai_plugin_test.py b/py/plugins/vertex-ai/tests/vertex_ai_plugin_test.py deleted file mode 100644 index 83bb8c7ec1..0000000000 --- a/py/plugins/vertex-ai/tests/vertex_ai_plugin_test.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Vertex AI plugin.""" - -from genkit.plugins.vertex_ai import ( - BigQueryRetriever, - FirestoreRetriever, - ModelGardenPlugin, - RetrieverOptionsSchema, - define_vertex_vector_search_big_query, - define_vertex_vector_search_firestore, - package_name, -) - - -def test_package_name() -> None: - """Test package_name returns correct value.""" - assert package_name() == 'genkit.plugins.vertex_ai' - - -def test_model_garden_plugin_exported() -> None: - """Test ModelGardenPlugin is exported.""" - assert ModelGardenPlugin is not None - - -def test_bigquery_retriever_exported() -> None: - """Test BigQueryRetriever is exported.""" - assert BigQueryRetriever is not None - - -def test_firestore_retriever_exported() -> None: - """Test FirestoreRetriever is exported.""" - assert FirestoreRetriever is not None - - -def test_retriever_options_schema_exported() -> None: - """Test RetrieverOptionsSchema is exported.""" - assert RetrieverOptionsSchema is not None - - -def test_define_vertex_vector_search_big_query_callable() -> None: - """Test define_vertex_vector_search_big_query is callable.""" - assert callable(define_vertex_vector_search_big_query) - - -def test_define_vertex_vector_search_firestore_callable() -> None: - """Test define_vertex_vector_search_firestore is callable.""" - assert callable(define_vertex_vector_search_firestore) diff --git a/py/pyproject.toml b/py/pyproject.toml index 993eb9e5ef..1081e4f7d1 100644 --- a/py/pyproject.toml +++ b/py/pyproject.toml @@ -21,15 +21,12 @@ dependencies = [ "genkit", "genkit-plugin-anthropic", "genkit-plugin-compat-oai", - "genkit-plugin-dev-local-vectorstore", "genkit-plugin-fastapi", - "genkit-plugin-firebase", "genkit-plugin-flask", "genkit-plugin-google-cloud", "genkit-plugin-google-genai", "genkit-plugin-ollama", "genkit-plugin-vertex-ai", - "genkit-plugin-amazon-bedrock", # Internal tools (private, not published) "conform", "liccheck>=0.9.2", @@ -72,7 +69,7 @@ lint = [ "litestar>=2.20.0", # For web/typing.py type resolution "mypy>=1.14.0", "pip-audit>=2.7.0", - "pypdf>=6.6.2", + "pypdf>=6.7.5", "pyrefly>=0.15.0", "pyright>=1.1.392", "pysentry-rs>=0.3.14", @@ -116,6 +113,13 @@ addopts = [ "--import-mode=importlib", ] asyncio_default_fixture_loop_scope = "session" +filterwarnings = [ + # Pydantic warns when a field named 'schema' shadows BaseModel.schema() + # We intentionally use 'schema' to match the JSON schema spec. + "ignore:Field name .schema. in .* shadows an attribute in parent .BaseModel.:UserWarning", + # dotpromptz uses the same pattern; not our code, suppress it too. + "ignore:Field name .schema. in .Prompt.*Config. shadows:UserWarning", +] norecursedirs = ["testapps", ".git", ".tox", ".nox", ".venv", "build", "dist"] python_files = ["*_test.py"] pythonpath = ["."] # For samples.shared imports @@ -129,7 +133,7 @@ fail_under = 78 [tool.coverage.run] omit = [ "**/__init__.py", # Often contains just imports - "**/testing.py", # Test utilities + "**/_testing.py", # Internal test utilities "**/constants.py", # Typically just constants "**/typing.py", # Often auto-generated or complex types "**/types.py", # Often auto-generated or complex types @@ -139,10 +143,10 @@ source = ["packages", "plugins"] # uv based package management. [tool.uv] default-groups = ["dev", "lint"] +override-dependencies = ["werkzeug>=3.1.6"] [tool.uv.sources] # Samples (alphabetical by package name from pyproject.toml) -dev-local-vectorstore-hello = { workspace = true } framework-context-demo = { workspace = true } framework-custom-evaluators = { workspace = true } framework-dynamic-tools-demo = { workspace = true } @@ -152,7 +156,6 @@ framework-prompt-demo = { workspace = true } framework-realtime-tracing-demo = { workspace = true } framework-restaurant-demo = { workspace = true } framework-tool-interrupts = { workspace = true } -provider-amazon-bedrock-hello = { workspace = true } provider-anthropic-hello = { workspace = true } provider-compat-oai-hello = { workspace = true } provider-firestore-retriever = { workspace = true } @@ -176,12 +179,9 @@ web-short-n-long = { workspace = true } # Core packages genkit = { workspace = true } # Plugins (alphabetical) -genkit-plugin-amazon-bedrock = { workspace = true } genkit-plugin-anthropic = { workspace = true } genkit-plugin-compat-oai = { workspace = true } -genkit-plugin-dev-local-vectorstore = { workspace = true } genkit-plugin-fastapi = { workspace = true } -genkit-plugin-firebase = { workspace = true } genkit-plugin-flask = { workspace = true } genkit-plugin-google-cloud = { workspace = true } genkit-plugin-google-genai = { workspace = true } @@ -190,7 +190,6 @@ genkit-plugin-vertex-ai = { workspace = true } # Internal tools (private, not published) conform = { workspace = true } genkit-tools-model-config-test = { workspace = true } -genkit-tools-sample-flows = { workspace = true } [tool.uv.workspace] exclude = ["*/shared", "testapps/*"] @@ -266,8 +265,21 @@ select = [ [tool.ruff.lint.per-file-ignores] # Auto-generated file from schema -"packages/genkit/src/genkit/core/typing.py" = ["E501"] -# Typing tests use reveal_type(), intentional unused vars, and don't need docstrings +"packages/genkit/src/genkit/_core/_typing.py" = ["E501"] +# Base model uses **kwargs: Any to match pydantic signatures +"packages/genkit/src/genkit/_core/_base.py" = ["ANN401"] +# HTTP client uses **httpx_kwargs: Any to forward arbitrary options +"packages/genkit/src/genkit/_core/_http_client.py" = ["ANN401"] +# Span wrapper uses __getattr__ -> Any for dynamic delegation +"packages/genkit/src/genkit/_core/trace/_adjusting_exporter.py" = ["ANN401"] +# Re-export modules (F401 would flag imports used only for re-export) +"packages/genkit/src/genkit/_ai/__init__.py" = ["F401"] +"packages/genkit/src/genkit/_ai/_document.py" = ["F401"] +"packages/genkit/src/genkit/_ai/_model.py" = ["F401"] +"packages/genkit/src/genkit/_ai/_formats/__init__.py" = ["F401"] +"packages/genkit/src/genkit/_core/trace/__init__.py" = ["F401"] +# Test files don't need docstrings; test method names are self-documenting +"**/tests/**/*.py" = ["D", "ANN401"] "packages/genkit/tests/typing/*.py" = ["D", "F821", "F841", "B018", "ANN"] # Validation/provenance/safe-defaults test methods are self-documenting via names; # S108 false positives on Path('/tmp/x') used as dummy non-existent paths in tests. @@ -305,13 +317,14 @@ skip-magic-trailing-comma = false #collapse-root-models = true # Don't use; produces Any as types. #strict-types = ["str", "int", "float", "bool", "bytes"] # Don't use; produces StrictStr, StrictInt, etc. #use-subclass-enum = true +base-class = "genkit._core._base.GenkitModel" capitalize-enum-members = true disable-timestamp = true enable-version-header = true field-constraints = true input = "../genkit-tools/genkit-schema.json" input-file-type = "jsonschema" -output = "packages/genkit/src/genkit/core/typing.py" +output = "packages/genkit/src/genkit/_core/_typing.py" output-model-type = "pydantic_v2.BaseModel" snake-case-field = true strict-nullable = true @@ -386,10 +399,7 @@ root = [ "packages/genkit/src", # Plugins "plugins/anthropic/src", - "plugins/amazon-bedrock/src", "plugins/compat-oai/src", - "plugins/dev-local-vectorstore/src", - "plugins/firebase/src", "plugins/flask/src", "plugins/google-cloud/src", "plugins/google-genai/src", @@ -419,26 +429,20 @@ exclude = [ # Use the venv where packages are installed as editable installs. # extraPaths lists all source roots so pyright can resolve PEP 420 -# namespace packages (genkit.plugins.*) across the workspace without -# relying solely on the venv's editable install paths. +# namespace packages (genkit.plugins.*) across the workspace. extraPaths = [ "packages/genkit/src", - "plugins/amazon-bedrock/src", "plugins/anthropic/src", "plugins/compat-oai/src", - "plugins/dev-local-vectorstore/src", "plugins/fastapi/src", - "plugins/firebase/src", "plugins/flask/src", "plugins/google-cloud/src", "plugins/google-genai/src", "plugins/ollama/src", "plugins/vertex-ai/src", "samples/framework-custom-evaluators", - # Tools "tools/conform/src", "tools/model-config-test", - "tools/sample-flows", ] pythonVersion = "3.10" reportMissingImports = true @@ -486,7 +490,6 @@ project_includes = [ "samples/*/tests/**/*.py", # Tools "tools/conform/src/**/*.py", - "tools/sample-flows/**/*.py", ] # Search path for first-party code import resolution. @@ -499,7 +502,6 @@ search-path = [ "samples/framework-restaurant-demo/src", # Tools "tools/conform/src", - "tools/sample-flows", "tools/model-config-test", ] # Ignore missing imports for namespace packages - pyrefly can't resolve PEP 420 @@ -512,3 +514,24 @@ python_version = "3.10" [tool.pyrefly.errors] deprecated = "error" redundant-cast = "error" + +# Ty analysis: suppress unresolved-import for plugins (PEP 420 namespace packages). +[tool.ty.analysis] +allowed-unresolved-imports = ["genkit.plugins.**"] + +# Downgrade unused-type-ignore to avoid noise from comments that pyright needed but ty doesn't. +[tool.ty.rules] +unused-type-ignore-comment = "ignore" + +# Relax invalid-argument-type in tests: config=dict is common, ModelRequest validates at runtime. +[[tool.ty.overrides]] +include = ["**/tests/**", "**/*_test.py", "samples/**", "tools/**"] +[tool.ty.overrides.rules] +invalid-argument-type = "ignore" +no-matching-overload = "ignore" + +# Vertex AI model_garden has relative imports ty struggles to resolve (PEP 420 / namespace). +[[tool.ty.overrides]] +include = ["plugins/vertex-ai/**"] +[tool.ty.overrides.rules] +unresolved-import = "ignore" diff --git a/py/samples/README.md b/py/samples/README.md index 582095661f..736f34f72b 100644 --- a/py/samples/README.md +++ b/py/samples/README.md @@ -16,8 +16,8 @@ This directory contains sample applications demonstrating various Genkit feature β”‚ β”‚ provider-google-genai-vertexai- β”‚ β”‚ framework-dynamic-tools-demo β”‚ β”‚ β”‚ β”‚ hello β”‚ β”‚ framework-evaluator-demo β”‚ β”‚ β”‚ β”‚ provider-anthropic-hello β”‚ β”‚ framework-format-demo β”‚ β”‚ -β”‚ β”‚ provider-amazon-bedrock-hello β”‚ β”‚ framework-middleware-demo β”‚ β”‚ -β”‚ β”‚ provider-microsoft-foundry-hello β”‚ β”‚ framework-prompt-demo β”‚ β”‚ +β”‚ β”‚ provider-microsoft-foundry-hello β”‚ β”‚ framework-middleware-demo β”‚ β”‚ +β”‚ β”‚ provider-ollama-hello β”‚ β”‚ framework-prompt-demo β”‚ β”‚ β”‚ β”‚ provider-ollama-hello β”‚ β”‚ framework-realtime-tracing-demo β”‚ β”‚ β”‚ β”‚ provider-compat-oai-hello β”‚ β”‚ framework-restaurant-demo β”‚ β”‚ β”‚ β”‚ provider-deepseek-hello β”‚ β”‚ framework-tool-interrupts β”‚ β”‚ @@ -136,7 +136,6 @@ This is a living document β€” update it as new flows are added to samples. | Sample | Basic | Stream | Tools | Struct | Vision | Embed | Code | Reasoning | TTS | Cache | PDF | |--------|:-----:|:------:|:-----:|:------:|:------:|:-----:|:----:|:---------:|:---:|:-----:|:---:| -| **provider-amazon-bedrock-hello** | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | β€” | β€” | β€” | | **provider-anthropic-hello** | βœ… | βœ… | βœ… | βœ… | βœ… | β€” | βœ… | βœ… | β€” | βœ… | βœ… | | **provider-cloudflare-workers-ai-hello** | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | β€” | β€” | β€” | β€” | | **provider-cohere-hello** | βœ… | βœ… | βœ… | βœ… | β€” | βœ… | βœ… | β€” | β€” | β€” | β€” | @@ -165,9 +164,6 @@ Most samples require environment variables for API keys. Configure these before |----------|--------|----------|-------------|-----------------| | `GEMINI_API_KEY` | provider-google-genai-hello | Yes | Google AI Studio API key | [Google AI Studio](https://aistudio.google.com/apikey) | | `ANTHROPIC_API_KEY` | provider-anthropic-hello | Yes | Anthropic API key | [Anthropic Console](https://console.anthropic.com/) | -| `AWS_REGION` | provider-amazon-bedrock-hello | Yes | AWS region (e.g., `us-east-1`) | [AWS Bedrock Regions](https://docs.aws.amazon.com/general/latest/gr/bedrock.html) | -| `AWS_ACCESS_KEY_ID` | provider-amazon-bedrock-hello | Yes* | AWS access key | [AWS IAM](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html) | -| `AWS_SECRET_ACCESS_KEY` | provider-amazon-bedrock-hello | Yes* | AWS secret key | [AWS IAM](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html) | | `AZURE_AI_FOUNDRY_ENDPOINT` | provider-microsoft-foundry-hello | Yes | Azure AI Foundry endpoint | [Azure AI Foundry](https://ai.azure.com/) | | `AZURE_AI_FOUNDRY_API_KEY` | provider-microsoft-foundry-hello | Yes* | Azure AI Foundry API key | [Azure AI Foundry](https://ai.azure.com/) | | `OPENAI_API_KEY` | provider-compat-oai-hello | Yes | OpenAI API key | [OpenAI Platform](https://platform.openai.com/api-keys) | diff --git a/py/samples/dev-local-vectorstore-hello/LICENSE b/py/samples/dev-local-vectorstore-hello/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/samples/dev-local-vectorstore-hello/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/samples/dev-local-vectorstore-hello/README.md b/py/samples/dev-local-vectorstore-hello/README.md deleted file mode 100644 index 5d9abb6800..0000000000 --- a/py/samples/dev-local-vectorstore-hello/README.md +++ /dev/null @@ -1,117 +0,0 @@ -# Local Vector Store (RAG Development) - -Test RAG (Retrieval-Augmented Generation) locally without setting up an external -vector database. Documents are embedded and stored in-memory for fast iteration. - -## How Local RAG Works - -``` -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ LOCAL RAG PIPELINE β”‚ -β”‚ β”‚ -β”‚ STEP 1: INDEX (one-time) STEP 2: RETRIEVE (per query) β”‚ -β”‚ ───────────────────────── ──────────────────────────── β”‚ -β”‚ β”‚ -β”‚ "The Matrix is a 1999..." "What's a good sci-fi film?" β”‚ -β”‚ β”‚ β”‚ β”‚ -β”‚ β–Ό β–Ό β”‚ -β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ -β”‚ β”‚ Embedder β”‚ β”‚ Embedder β”‚ β”‚ -β”‚ β”‚ (Vertex AI) β”‚ β”‚ (Vertex AI) β”‚ β”‚ -β”‚ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ -β”‚ β”‚ β”‚ β”‚ -β”‚ β–Ό β–Ό β”‚ -β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ -β”‚ β”‚ Local Vector │◄─── compare ────│ Query Vector β”‚ β”‚ -β”‚ β”‚ Store β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ -β”‚ β”‚ (in-memory) │────────────────────────────────► β”‚ -β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ "The Matrix" (0.95 match) β”‚ -β”‚ "Inception" (0.89 match) β”‚ -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ -``` - -## Features Demonstrated - -| Feature | Flow / API | Description | -|---------|-----------|-------------| -| Local Vector Store | `define_dev_local_vector_store()` | In-memory store, no external DB needed | -| Document Indexing | `index_documents` | Embed and store film descriptions | -| Similarity Retrieval | `retrieve_documents` | Find semantically similar documents | -| Document Creation | `Document.from_text()` | Convert plain text to Genkit documents | -| Vertex AI Embeddings | `gemini-embedding-001` | Cloud-based embedding model | - -## ELI5: Key Concepts - -| Concept | ELI5 | -|---------|------| -| **RAG** | AI looks up your docs before answering β€” fewer hallucinations! | -| **Vector Store** | A database that finds "similar" items by meaning β€” "happy" finds docs about "joyful" too | -| **Local Store** | Runs on your computer, no cloud DB needed β€” perfect for testing before production | -| **Indexing** | Adding documents to the store β€” like a librarian cataloging new books | -| **Retrieval** | Finding documents that match a query β€” "sci-fi films" returns The Matrix, Inception | - -## Quick Start - -```bash -export GCLOUD_PROJECT=your-project-id -gcloud auth application-default login -./run.sh -``` - -Then open the Dev UI at http://localhost:4000. - -## Setup - -### 1. Authenticate with Google Cloud - -Vertex AI embeddings require GCP authentication: - -```bash -gcloud auth application-default login -``` - -### 2. Set Your Project ID - -```bash -export GCLOUD_PROJECT=your-project-id -``` - -Or the demo will prompt you interactively. - -### 3. Run the Sample - -```bash -./run.sh -``` - -Or manually: - -```bash -genkit start -- uv run src/main.py -``` - -## Testing This Demo - -1. **Open DevUI** at http://localhost:4000 - -2. **Index documents first**: - - [ ] Run `index_documents` β€” indexes 10 classic film descriptions - -3. **Test retrieval**: - - [ ] Run `retrieve_documents` β€” queries for "sci-fi film" - - [ ] Verify The Matrix, Inception, Star Wars appear as top matches - -4. **Expected behavior**: - - Documents are embedded using Vertex AI (`gemini-embedding-001`) - - Vector store is in-memory (no external database) - - Retrieval returns semantically similar documents (top 3) - - No external vector database required - -## Sample Documents - -The sample indexes 10 classic films including The Godfather, The Dark Knight, -Pulp Fiction, Inception, The Matrix, Star Wars, and more. - -## Development - -The `run.sh` script uses `watchmedo` for hot reloading on file changes. diff --git a/py/samples/dev-local-vectorstore-hello/pyproject.toml b/py/samples/dev-local-vectorstore-hello/pyproject.toml deleted file mode 100644 index f3e79cac96..0000000000 --- a/py/samples/dev-local-vectorstore-hello/pyproject.toml +++ /dev/null @@ -1,65 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Private :: Do Not Upload", -] -dependencies = [ - "rich>=13.0.0", - "genkit", - "genkit-plugin-dev-local-vectorstore", - "genkit-plugin-google-genai", - "pydantic>=2.10.5", - "structlog>=25.2.0", - "uvloop>=0.21.0", -] -description = "hello Genkit sample" -license = "Apache-2.0" -name = "dev-local-vectorstore-hello" -readme = "README.md" -requires-python = ">=3.10" -version = "0.2.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src/hello"] diff --git a/py/samples/dev-local-vectorstore-hello/run.sh b/py/samples/dev-local-vectorstore-hello/run.sh deleted file mode 100755 index 677f556a6c..0000000000 --- a/py/samples/dev-local-vectorstore-hello/run.sh +++ /dev/null @@ -1,56 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -# Local Vector Store Demo -# ======================= -# -# Demonstrates using an in-memory vector store for development. -# -# Prerequisites: -# - GEMINI_API_KEY environment variable set -# -# Usage: -# ./run.sh # Start the demo with Dev UI -# ./run.sh --help # Show this help message - -set -euo pipefail - -cd "$(dirname "$0")" -source "../_common.sh" - -print_help() { - print_banner "Local Vector Store Demo" "πŸ’Ύ" - echo "Usage: ./run.sh [options]" - echo "" - echo "Options:" - echo " --help Show this help message" - echo "" - echo "Environment Variables:" - echo " GEMINI_API_KEY Required. Your Gemini API key" - echo "" - echo "Get an API key from: https://makersuite.google.com/app/apikey" - print_help_footer -} - -case "${1:-}" in - --help|-h) - print_help - exit 0 - ;; -esac - -print_banner "Local Vector Store Demo" "πŸ’Ύ" - -check_env_var "GEMINI_API_KEY" "https://makersuite.google.com/app/apikey" || true - -install_deps - -genkit_start_with_browser -- \ - uv tool run --from watchdog watchmedo auto-restart \ - -d src \ - -d ../../packages \ - -d ../../plugins \ - -p '*.py;*.prompt;*.json' \ - -R \ - -- uv run src/main.py "$@" diff --git a/py/samples/dev-local-vectorstore-hello/src/main.py b/py/samples/dev-local-vectorstore-hello/src/main.py deleted file mode 100755 index de020b81cf..0000000000 --- a/py/samples/dev-local-vectorstore-hello/src/main.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Dev local vector store sample - Local RAG without external services. - -This sample demonstrates Genkit's local vector store for development, -which allows testing RAG (Retrieval Augmented Generation) without -setting up external vector databases. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ RAG β”‚ Retrieval-Augmented Generation. AI looks up β”‚ - β”‚ β”‚ your docs before answering. Fewer hallucinations! β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Vector Store β”‚ A database that finds "similar" items by meaning. β”‚ - β”‚ β”‚ "Happy" finds docs about "joyful" too. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Local Store β”‚ Runs on your computer, no cloud needed. Perfect β”‚ - β”‚ β”‚ for testing before deploying to production. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Indexing β”‚ Adding documents to the store. Like a librarian β”‚ - β”‚ β”‚ cataloging new books. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Retrieval β”‚ Finding documents that match a query. "Show me β”‚ - β”‚ β”‚ docs about sci-fi films" returns matching results. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow (RAG Pipeline):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW RAG FINDS ANSWERS FROM YOUR DOCUMENTS β”‚ - β”‚ β”‚ - β”‚ STEP 1: INDEX (one-time setup) β”‚ - β”‚ ────────────────────────────── β”‚ - β”‚ Your Documents: ["The Godfather...", "The Matrix...", ...] β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) Convert each doc to numbers (embeddings) β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ "sci-fi film" β†’ [0.2, -0.5, 0.8, ...] β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Store in local vector store β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Local Store β”‚ All docs + embeddings saved locally β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ - β”‚ STEP 2: RETRIEVE (at query time) β”‚ - β”‚ ──────────────────────────────── β”‚ - β”‚ Query: "What's a good sci-fi movie?" β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Convert query to embedding β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ Query β†’ [0.21, -0.48, 0.79, ...] (similar!) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (4) Find nearest matches β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Local Store β”‚ "The Matrix" (0.95 match) β”‚ - β”‚ β”‚ β”‚ "Inception" (0.89 match) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Features -============ -| Feature Description | Example Function / Code Snippet | -|-----------------------------------------|-------------------------------------| -| Local Vector Store Definition | `define_dev_local_vector_store` | -| Document Indexing | `ai.index()` | -| Document Retrieval | `ai.retrieve()` | -| Document Structure | `Document.from_text()` | - -See README.md for testing instructions. -""" - -import asyncio -import os - -from genkit.ai import Genkit -from genkit.plugins.dev_local_vectorstore import define_dev_local_vector_store -from genkit.plugins.google_genai import VertexAI -from genkit.types import Document, RetrieverResponse -from samples.shared.logging import setup_sample - -setup_sample() - -if 'GCLOUD_PROJECT' not in os.environ: - os.environ['GCLOUD_PROJECT'] = input('Please enter your GCLOUD_PROJECT: ') - -ai = Genkit( - plugins=[VertexAI()], - model='vertexai/gemini-3-flash-preview', -) - -# Define dev local vector store -define_dev_local_vector_store( - ai, - name='films', - embedder='vertexai/gemini-embedding-001', -) - -films = [ - 'The Godfather is a 1972 crime film directed by Francis Ford Coppola.', - 'The Dark Knight is a 2008 superhero film directed by Christopher Nolan.', - 'Pulp Fiction is a 1994 crime film directed by Quentin Tarantino.', - "Schindler's List is a 1993 historical drama directed by Steven Spielberg.", - 'Inception is a 2010 sci-fi film directed by Christopher Nolan.', - 'The Matrix is a 1999 sci-fi film directed by the Wachowskis.', - 'Fight Club is a 1999 film directed by David Fincher.', - 'Forrest Gump is a 1994 drama directed by Robert Zemeckis.', - 'Star Wars is a 1977 sci-fi film directed by George Lucas.', - 'The Shawshank Redemption is a 1994 drama directed by Frank Darabont.', -] - - -@ai.flow() -async def index_documents() -> None: - """Indexes the film documents in the local vector store.""" - genkit_documents = [Document.from_text(text=film) for film in films] - await ai.index( - indexer='films', - documents=genkit_documents, - ) - - -@ai.flow() -async def retrieve_documents() -> RetrieverResponse: - """Retrieve documents from the vector store.""" - return await ai.retrieve( - query=Document.from_text('sci-fi film'), - retriever='films', - options={'limit': 3}, - ) - - -async def main() -> None: - """Main entry point for the sample - keep alive for Dev UI.""" - # Keep the process alive for Dev UI - await asyncio.Event().wait() - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/samples/framework-context-demo/src/main.py b/py/samples/framework-context-demo/src/main.py index 65870faeec..7c542e341c 100644 --- a/py/samples/framework-context-demo/src/main.py +++ b/py/samples/framework-context-demo/src/main.py @@ -84,11 +84,11 @@ import asyncio import os +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger +from genkit import Genkit +from genkit._core._action import ActionRunContext from genkit.plugins.google_genai import GoogleAI from samples.shared.logging import setup_sample @@ -97,7 +97,7 @@ if 'GEMINI_API_KEY' not in os.environ: os.environ['GEMINI_API_KEY'] = input('Please enter your GEMINI_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) ai = Genkit( plugins=[GoogleAI()], @@ -136,7 +136,7 @@ class ContextInput(BaseModel): @ai.tool() -def get_user_info() -> str: +async def get_user_info() -> str: """Look up the current user from context. This tool takes no explicit input -- it reads the user ID from the @@ -151,7 +151,7 @@ def get_user_info() -> str: @ai.tool() -def get_user_via_static() -> str: +async def get_user_via_static() -> str: """Look up the current user using Genkit.current_context(). Demonstrates the static method approach -- useful when context is needed @@ -170,7 +170,7 @@ def get_user_via_static() -> str: @ai.tool() -def get_user_permissions() -> str: +async def get_user_permissions() -> str: """Return permissions based on the user's plan from context. Used in the propagation chain demo to verify context survives diff --git a/py/samples/framework-custom-evaluators/src/deliciousness_evaluator.py b/py/samples/framework-custom-evaluators/src/deliciousness_evaluator.py index c90b5d3e18..63e757b6b1 100644 --- a/py/samples/framework-custom-evaluators/src/deliciousness_evaluator.py +++ b/py/samples/framework-custom-evaluators/src/deliciousness_evaluator.py @@ -21,8 +21,8 @@ from pydantic import BaseModel -from genkit.ai import Genkit -from genkit.core.typing import BaseDataPoint, Details, EvalFnResponse, Score +from genkit import Genkit +from genkit._core._typing import BaseDataPoint, Details, EvalFnResponse, Score class DeliciousnessResponse(BaseModel): @@ -60,11 +60,11 @@ async def deliciousness_score( deliciousness_prompt = ai.prompt('deliciousness') rendered = await deliciousness_prompt.render(input={'output': str(datapoint.output)}) - response = await ai.generate( + response = await ai.generate( # pyrefly: ignore[no-matching-overload] model=judge, messages=rendered.messages, config=judge_config, - output={'schema': DeliciousnessResponse}, + output_schema=DeliciousnessResponse, ) if not response.output: diff --git a/py/samples/framework-custom-evaluators/src/funniness_evaluator.py b/py/samples/framework-custom-evaluators/src/funniness_evaluator.py index 2119540935..996a4ab641 100644 --- a/py/samples/framework-custom-evaluators/src/funniness_evaluator.py +++ b/py/samples/framework-custom-evaluators/src/funniness_evaluator.py @@ -21,8 +21,8 @@ from pydantic import BaseModel -from genkit.ai import Genkit -from genkit.core.typing import BaseDataPoint, Details, EvalFnResponse, Score +from genkit import Genkit +from genkit._core._typing import BaseDataPoint, Details, EvalFnResponse, Score class FunninessResponse(BaseModel): @@ -60,11 +60,11 @@ async def funniness_score( funniness_prompt = ai.prompt('funniness') rendered = await funniness_prompt.render(input={'output': str(datapoint.output)}) - response = await ai.generate( + response = await ai.generate( # pyrefly: ignore[no-matching-overload] model=judge, messages=rendered.messages, config=judge_config, - output={'schema': FunninessResponse}, + output_schema=FunninessResponse, ) if not response.output: diff --git a/py/samples/framework-custom-evaluators/src/main.py b/py/samples/framework-custom-evaluators/src/main.py index 4d694be7dd..440188407f 100644 --- a/py/samples/framework-custom-evaluators/src/main.py +++ b/py/samples/framework-custom-evaluators/src/main.py @@ -65,8 +65,9 @@ import os from pathlib import Path -from genkit.ai import Genkit -from genkit.core.logging import get_logger +import structlog + +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI from src.constants import PERMISSIVE_SAFETY_SETTINGS, URL_REGEX, US_PHONE_REGEX from src.deliciousness_evaluator import register_deliciousness_evaluator @@ -74,7 +75,7 @@ from src.pii_evaluator import register_pii_evaluator from src.regex_evaluator import regex_matcher, register_regex_evaluators -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) # Get prompts directory path current_dir = Path(__file__).resolve().parent diff --git a/py/samples/framework-custom-evaluators/src/pii_evaluator.py b/py/samples/framework-custom-evaluators/src/pii_evaluator.py index 7b1ff5ba5a..ddbacb810e 100644 --- a/py/samples/framework-custom-evaluators/src/pii_evaluator.py +++ b/py/samples/framework-custom-evaluators/src/pii_evaluator.py @@ -20,8 +20,8 @@ from pydantic import BaseModel -from genkit.ai import Genkit -from genkit.core.typing import BaseDataPoint, Details, EvalFnResponse, Score +from genkit import Genkit +from genkit._core._typing import BaseDataPoint, Details, EvalFnResponse, Score class PiiDetectionResponse(BaseModel): @@ -59,11 +59,11 @@ async def pii_detection_score( pii_prompt = ai.prompt('pii_detection') rendered = await pii_prompt.render(input={'output': str(datapoint.output)}) - response = await ai.generate( + response = await ai.generate( # pyrefly: ignore[no-matching-overload] model=judge, messages=rendered.messages, config=judge_config, - output={'schema': PiiDetectionResponse}, + output_schema=PiiDetectionResponse, ) if not response.output: diff --git a/py/samples/framework-custom-evaluators/src/regex_evaluator.py b/py/samples/framework-custom-evaluators/src/regex_evaluator.py index 13d6c52123..8108e76e27 100644 --- a/py/samples/framework-custom-evaluators/src/regex_evaluator.py +++ b/py/samples/framework-custom-evaluators/src/regex_evaluator.py @@ -25,8 +25,8 @@ from re import Pattern from typing import Any, cast -from genkit.ai import Genkit -from genkit.core.typing import BaseDataPoint, Details, EvalFnResponse, Score +from genkit import Genkit +from genkit._core._typing import BaseDataPoint, Details, EvalFnResponse, Score def regex_matcher(suffix: str, pattern: Pattern[str]) -> dict[str, object]: diff --git a/py/samples/framework-dynamic-tools-demo/src/main.py b/py/samples/framework-dynamic-tools-demo/src/main.py index faa6ae0c21..fdedaf6f85 100644 --- a/py/samples/framework-dynamic-tools-demo/src/main.py +++ b/py/samples/framework-dynamic-tools-demo/src/main.py @@ -52,13 +52,13 @@ β”‚ β”‚ β”‚ combined_demo(input) β”‚ β”‚ β”‚ β”‚ - β”‚ β”œβ”€β”€ ai.run("preprocess_step", input, preprocess) β”‚ + β”‚ β”œβ”€β”€ ai.run(name="preprocess_step", fn=preprocess) β”‚ β”‚ β”‚ └── Returns preprocessed string β”‚ β”‚ β”‚ β”‚ - β”‚ β”œβ”€β”€ ai.dynamic_tool("scaler", scale_fn) β”‚ + β”‚ β”œβ”€β”€ ai.dynamic_tool(name="scaler", fn=scale_fn) β”‚ β”‚ β”‚ └── Creates tool (not globally registered) β”‚ β”‚ β”‚ β”‚ - β”‚ β”œβ”€β”€ scaler.arun(7) β”‚ + β”‚ β”œβ”€β”€ scaler.run(7) β”‚ β”‚ β”‚ └── Returns 7 * 10 = 70 β”‚ β”‚ β”‚ β”‚ β”‚ └── Returns {step_result, tool_result, tool_metadata} β”‚ @@ -81,10 +81,10 @@ import asyncio import os +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.core.logging import get_logger +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI from samples.shared.logging import setup_sample @@ -93,7 +93,7 @@ if 'GEMINI_API_KEY' not in os.environ: os.environ['GEMINI_API_KEY'] = input('Please enter your GEMINI_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) ai = Genkit( plugins=[GoogleAI()], @@ -138,11 +138,11 @@ def multiplier_fn(x: int) -> int: return x * 10 dynamic_multiplier = ai.dynamic_tool( - 'dynamic_multiplier', - multiplier_fn, + name='dynamic_multiplier', + fn=multiplier_fn, description='Multiplies input by 10', ) - result = await dynamic_multiplier.arun(input.value) + result = await dynamic_multiplier.run(input.value) return { 'input_value': input.value, @@ -156,9 +156,8 @@ def multiplier_fn(x: int) -> int: async def run_step_demo(input: RunStepInput) -> dict[str, str]: """Wrap a plain function as a traceable step using ai.run(). - ``ai.run(name, input, fn)`` creates a named sub-span in the trace. - The step's input and output are recorded and visible in the Dev UI - trace viewer. + ``ai.run(name=..., fn=...)`` creates a named sub-span in the trace. + The step's output is recorded and visible in the Dev UI trace viewer. Args: input: Input with data to process. @@ -167,14 +166,14 @@ async def run_step_demo(input: RunStepInput) -> dict[str, str]: A dict containing the original and processed data. """ - def uppercase(data: str) -> str: - return data.upper() + async def uppercase() -> str: + return input.data.upper() - def reverse(data: str) -> str: - return data[::-1] + async def reverse() -> str: + return step1[::-1] - step1 = await ai.run('uppercase_step', input.data, uppercase) - step2 = await ai.run('reverse_step', step1, reverse) + step1 = await ai.run(name='uppercase_step', fn=uppercase) + step2 = await ai.run(name='reverse_step', fn=reverse) return { 'original': input.data, @@ -199,16 +198,16 @@ async def combined_demo(input: CombinedInput) -> dict[str, object]: A dict with results from both the step and the dynamic tool. """ - def preprocess(data: str) -> str: - return f'processed: {data}' + async def preprocess() -> str: + return f'processed: {input.input_val}' - step_result = await ai.run('preprocess_step', input.input_val, preprocess) + step_result = await ai.run(name='preprocess_step', fn=preprocess) - def scale_fn(x: int) -> int: + async def scale_fn(x: int) -> int: return x * 10 - scaler = ai.dynamic_tool('scaler', scale_fn, description='Scales input by 10') - tool_result = await scaler.arun(7) + scaler = ai.dynamic_tool(name='scaler', fn=scale_fn, description='Scales input by 10') + tool_result = await scaler.run(7) return { 'step_result': step_result, diff --git a/py/samples/framework-format-demo/src/main.py b/py/samples/framework-format-demo/src/main.py index e919c3a30e..d93acf5307 100644 --- a/py/samples/framework-format-demo/src/main.py +++ b/py/samples/framework-format-demo/src/main.py @@ -58,17 +58,16 @@ import os from typing import Any, cast +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.core.logging import get_logger -from genkit.core.typing import OutputConfig +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI from samples.shared.logging import setup_sample setup_sample() -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) if 'GEMINI_API_KEY' not in os.environ: os.environ['GEMINI_API_KEY'] = input('Please enter your GEMINI_API_KEY: ') @@ -151,13 +150,11 @@ async def classify_sentiment_enum(input: ClassifySentimentInput) -> str: """ response = await ai.generate( prompt=f'Classify the sentiment of this review: "{input.review}"', - output=OutputConfig( - format='enum', - schema={ - 'type': 'string', - 'enum': ['POSITIVE', 'NEGATIVE', 'NEUTRAL'], - }, - ), + output_format='enum', + output_schema={ + 'type': 'string', + 'enum': ['POSITIVE', 'NEGATIVE', 'NEUTRAL'], + }, ) return cast(str, response.output) @@ -183,13 +180,11 @@ async def create_story_characters_jsonl(input: CreateStoryCharactersInput) -> li """ response = await ai.generate( prompt=f'Generate 5 characters for a {input.theme} story', - output=OutputConfig( - format='jsonl', - schema={ - 'type': 'array', - 'items': Character.model_json_schema(), - }, - ), + output_format='jsonl', + output_schema={ + 'type': 'array', + 'items': Character.model_json_schema(), + }, ) return cast(list[Any], response.output) @@ -208,7 +203,7 @@ async def generate_haiku_text(input: HaikuInput) -> str: """ response = await ai.generate( prompt=f'Write a haiku about {input.topic}', - output=OutputConfig(format='text'), + output_format='text', ) return response.text @@ -226,7 +221,8 @@ async def get_country_info_json(input: CountryInfoInput) -> dict[str, Any]: """ response = await ai.generate( prompt=f'Give me information about {input.country}', - output=OutputConfig(format='json', schema=CountryInfo.model_json_schema()), + output_format='json', + output_schema=CountryInfo.model_json_schema(), ) return cast(dict[str, Any], response.output) @@ -248,13 +244,11 @@ async def recommend_books_array(input: RecommendBooksInput) -> list[dict[str, ob """ response = await ai.generate( prompt=f'List 3 famous {input.genre} books', - output=OutputConfig( - format='array', - schema={ - 'type': 'array', - 'items': Book.model_json_schema(), - }, - ), + output_format='array', + output_schema={ + 'type': 'array', + 'items': Book.model_json_schema(), + }, ) return cast(list[Any], response.output) diff --git a/py/samples/framework-middleware-demo/src/main.py b/py/samples/framework-middleware-demo/src/main.py index 3da6983bb1..60203ae40a 100644 --- a/py/samples/framework-middleware-demo/src/main.py +++ b/py/samples/framework-middleware-demo/src/main.py @@ -84,15 +84,14 @@ import asyncio import os +from collections.abc import Awaitable, Callable +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.blocks.model import ModelMiddlewareNext -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger +from genkit import Genkit, Message, ModelRequest, ModelResponse, Part, Role, TextPart +from genkit._core._action import ActionRunContext from genkit.plugins.google_genai import GoogleAI -from genkit.types import GenerateRequest, GenerateResponse, Message, Part, Role, TextPart from samples.shared.logging import setup_sample setup_sample() @@ -100,7 +99,7 @@ if 'GEMINI_API_KEY' not in os.environ: os.environ['GEMINI_API_KEY'] = input('Please enter your GEMINI_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) ai = Genkit( plugins=[GoogleAI()], @@ -127,10 +126,10 @@ class ChainedInput(BaseModel): async def logging_middleware( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next_handler: ModelMiddlewareNext, -) -> GenerateResponse: + next_handler: Callable[[ModelRequest, ActionRunContext], Awaitable[ModelResponse]], +) -> ModelResponse: """Middleware that logs request and response metadata. This is a pass-through middleware that doesn't modify the request @@ -158,10 +157,10 @@ async def logging_middleware( async def system_instruction_middleware( - req: GenerateRequest, + req: ModelRequest, ctx: ActionRunContext, - next_handler: ModelMiddlewareNext, -) -> GenerateResponse: + next_handler: Callable[[ModelRequest, ActionRunContext], Awaitable[ModelResponse]], +) -> ModelResponse: """Middleware that prepends a system instruction to every request. Demonstrates modifying the request before it reaches the model. diff --git a/py/samples/framework-prompt-demo/src/main.py b/py/samples/framework-prompt-demo/src/main.py index a0f9d99b74..08aca643d9 100755 --- a/py/samples/framework-prompt-demo/src/main.py +++ b/py/samples/framework-prompt-demo/src/main.py @@ -69,11 +69,11 @@ import os from pathlib import Path +import structlog from pydantic import BaseModel, Field -from genkit.ai import ActionKind, Genkit -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger +from genkit import ActionKind, Genkit +from genkit._core._action import ActionRunContext from genkit.plugins.google_genai import GoogleAI from samples.shared.logging import setup_sample @@ -82,7 +82,7 @@ if 'GEMINI_API_KEY' not in os.environ: os.environ['GEMINI_API_KEY'] = input('Please enter your GEMINI_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) current_dir = Path(__file__).resolve().parent @@ -250,9 +250,9 @@ async def main() -> None: # Tell Story Flow (Streaming) await logger.ainfo('--- Running Tell Story Flow ---') # To demonstrate streaming, we'll iterate over the streamer if calling directly like a flow would be consumed. - story_stream, _ = tell_story.stream(StoryInput(subject='a brave little toaster', personality='courageous')) + story_stream_response = tell_story.stream(StoryInput(subject='a brave little toaster', personality='courageous')) - async for _chunk in story_stream: + async for _chunk in story_stream_response.stream: pass await logger.ainfo('Tell Story Flow Completed') diff --git a/py/samples/framework-realtime-tracing-demo/README.md b/py/samples/framework-realtime-tracing-demo/README.md index 6b6d3c7c0a..08096aa617 100644 --- a/py/samples/framework-realtime-tracing-demo/README.md +++ b/py/samples/framework-realtime-tracing-demo/README.md @@ -1,135 +1,12 @@ # Realtime Tracing Demo -This sample demonstrates Genkit's realtime tracing feature, which exports spans to the DevUI as they **start** (not just when they complete). This enables live visualization of in-progress operations. +Spans appear in DevUI **as they start** (not when they complete). Watch `realtime_demo` in the Traces tabβ€”each step shows up immediately. -## What is Realtime Tracing? - -Standard OpenTelemetry processors only export spans when they finish. This means long-running operations (like multi-turn LLM calls or complex workflows) don't appear in the DevUI until they're complete. - -With realtime tracing enabled: -- Spans appear in the DevUI **immediately** when they start -- You can see operations **in progress** (without endTime) -- The span is exported again when it completes with full data - -``` -Standard Tracing: Realtime Tracing: - -Start ─────────────────► Start ────► [Span appears!] - (nothing visible) β”‚ - ... β”‚ (visible as "in progress") - ... β”‚ -End ──────► [Span appears] End ──┴─► [Span updated] -``` - -## How It Works - -The `RealtimeSpanProcessor` wraps a standard exporter and calls `export()` twice: -1. On `on_start()` - exports immediately (no endTime) -2. On `on_end()` - exports again with complete data - -## Usage - -### Enable via Environment Variable +## Run ```bash -# Enable realtime telemetry -export GENKIT_ENABLE_REALTIME_TELEMETRY=true - -# Run with genkit start -genkit start -- python src/main.py +cd py/samples/framework-realtime-tracing-demo +genkit start -- uv run src/main.py ``` -### Enable Programmatically - -```python -from opentelemetry.sdk.trace import TracerProvider - -from genkit.core.trace import RealtimeSpanProcessor, TelemetryServerSpanExporter - -# Create exporter -exporter = TelemetryServerSpanExporter( - telemetry_server_url='http://localhost:4000' -) - -# Wrap with RealtimeSpanProcessor -processor = RealtimeSpanProcessor(exporter) - -# Add to tracer provider -provider = TracerProvider() -provider.add_span_processor(processor) -``` - -## Running the Demo - -1. Run with realtime tracing enabled (with hot reload): - ```bash - ./run.sh - ``` - - You'll be prompted for `GEMINI_API_KEY` if not set. - -2. Open the DevUI at http://localhost:4000 - -3. Trigger flows and watch spans appear **immediately** as operations start! - -4. Edit code and it will automatically reload. - -## Key APIs Demonstrated - -| API | Description | -|-----|-------------| -| `RealtimeSpanProcessor` | SpanProcessor that exports on start AND end | -| `is_realtime_telemetry_enabled()` | Check if realtime mode is enabled | -| `create_span_processor(exporter)` | Auto-selects processor based on env | -| `GENKIT_ENABLE_REALTIME_TELEMETRY` | Environment variable to enable | - -## When to Use - -- **Development**: Great for debugging and understanding flow execution -- **Long-running operations**: See progress of complex workflows -- **DevUI demos**: Showcase live updates - -## When NOT to Use - -- **Production**: Doubles network traffic (each span exported twice) -- **High-throughput**: May impact performance -- **Simple flows**: Standard tracing is sufficient - -## Related Samples - -- `session-demo/` - Multi-turn conversations -- `chat-demo/` - Chat application with streaming -- `tool-interrupts/` - Human-in-the-loop workflows - -## Testing This Demo - -1. **Prerequisites**: - ```bash - export GEMINI_API_KEY=your_api_key - ``` - Or the demo will prompt for the key interactively. - -2. **Run the demo**: - ```bash - cd py/samples/framework-realtime-tracing-demo - ./run.sh # This sets GENKIT_ENABLE_REALTIME_TELEMETRY=true - ``` - -3. **Open DevUI** at http://localhost:4000 - -4. **Test realtime tracing**: - - [ ] Open the Traces tab in DevUI - - [ ] Trigger a multi-step flow - - [ ] Watch spans appear IMMEDIATELY as they start - - [ ] Compare to non-realtime (spans appear at end) - -5. **Test flows**: - - [ ] `multi_step_flow` - See each step appear in order - - [ ] `nested_flow` - See parent/child span hierarchy - - [ ] `long_running_flow` - Watch progress of slow tasks - -6. **Expected behavior**: - - Spans appear in DevUI as soon as they START - - You see "in progress" spans while they're running - - Nested spans show proper parent/child relationships - - Long-running spans show duration updating in real-time +Open DevUI at http://localhost:4000 and invoke `realtime_demo`. diff --git a/py/samples/framework-realtime-tracing-demo/run.sh b/py/samples/framework-realtime-tracing-demo/run.sh deleted file mode 100755 index 2dd0c67f66..0000000000 --- a/py/samples/framework-realtime-tracing-demo/run.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -# Real-time Tracing Demo -# ====================== -# -# Demonstrates live telemetry and trace visualization. -# -# Prerequisites: -# - GEMINI_API_KEY environment variable set -# -# Usage: -# ./run.sh # Start the demo with Dev UI -# ./run.sh --help # Show this help message - -set -euo pipefail - -cd "$(dirname "$0")" -source "../_common.sh" - -print_help() { - print_banner "Real-time Tracing Demo" "πŸ“‘" - echo "Usage: ./run.sh [options]" - echo "" - echo "Options:" - echo " --help Show this help message" - echo "" - echo "Environment Variables:" - echo " GEMINI_API_KEY Required. Your Gemini API key" - echo "" - echo "This demo shows:" - echo " - Live span streaming" - echo " - Trace visualization" - echo " - Performance monitoring" - echo "" - echo "Get an API key from: https://makersuite.google.com/app/apikey" - print_help_footer -} - -case "${1:-}" in - --help|-h) - print_help - exit 0 - ;; -esac - -print_banner "Real-time Tracing Demo" "πŸ“‘" - -check_env_var "GEMINI_API_KEY" "https://makersuite.google.com/app/apikey" || true - -install_deps - -# Enable realtime telemetry for this demo -export GENKIT_ENABLE_REALTIME_TELEMETRY=true - -genkit_start_with_browser -- \ - uv tool run --from watchdog watchmedo auto-restart \ - -d src \ - -d ../../packages \ - -d ../../plugins \ - -p '*.py;*.prompt;*.json' \ - -R \ - -- uv run src/main.py "$@" diff --git a/py/samples/framework-realtime-tracing-demo/src/main.py b/py/samples/framework-realtime-tracing-demo/src/main.py index e638ca0bc2..cb29412fb8 100644 --- a/py/samples/framework-realtime-tracing-demo/src/main.py +++ b/py/samples/framework-realtime-tracing-demo/src/main.py @@ -2,371 +2,52 @@ # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# # SPDX-License-Identifier: Apache-2.0 -"""Realtime Tracing Demo - Watch spans appear as they start. - -This sample demonstrates Genkit's realtime tracing feature, which exports -spans to the DevUI as they START (not just when they complete). - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Realtime Tracing β”‚ See what's happening AS it happens, not after. β”‚ - β”‚ β”‚ Like watching a live sports game vs highlights. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Span β”‚ A "timer" for one operation. Records when it β”‚ - β”‚ β”‚ started, when it ended, and what happened. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ SpanProcessor β”‚ The thing that decides when to send span data. β”‚ - β”‚ β”‚ Realtime = send on START, not just END. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Nested Spans β”‚ Spans inside spans. "Flow" contains "Model call" β”‚ - β”‚ β”‚ contains "HTTP request". Like Russian dolls. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ DevUI Traces Tab β”‚ The dashboard where you see all your traces. β”‚ - β”‚ β”‚ With realtime, spans appear immediately! β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow (Realtime vs Normal):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ NORMAL TRACING VS REALTIME TRACING β”‚ - β”‚ β”‚ - β”‚ NORMAL: You see spans AFTER they complete β”‚ - β”‚ ────────────────────────────────────────── β”‚ - β”‚ Flow starts β†’ ... waiting ... β†’ Flow ends β†’ NOW you see it β”‚ - β”‚ β”‚ - β”‚ REALTIME: You see spans AS they start β”‚ - β”‚ ───────────────────────────────────────── β”‚ - β”‚ Flow starts β†’ IMMEDIATELY visible β†’ updates as it runs β†’ completes β”‚ - β”‚ β”‚ - β”‚ Timeline: β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ 0s 1s 2s 3s 4s 5s β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ β”œβ”€ Flow starts (visible in realtime!) β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ β”œβ”€ Model call starts (visible!) β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ └─ Model call ends β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ β”œβ”€ Tool call starts (visible!) β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ └─ Tool call ends β”‚ β”‚ - β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ - β”‚ β”‚ └────┴─ Flow ends β”‚ β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Features -============ - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Feature β”‚ Description β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ RealtimeSpanProcessor β”‚ Exports spans on start AND end β”‚ - β”‚ Multi-step flows β”‚ Watch each step appear in real-time β”‚ - β”‚ Nested actions β”‚ See parent/child relationships live β”‚ - β”‚ Long-running operations β”‚ Monitor progress of slow tasks β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Testing This Demo -================= -1. **Prerequisites**: - ```bash - export GEMINI_API_KEY=your_api_key - ``` - Or the demo will prompt for the key interactively. - -2. **Run the demo**: - ```bash - cd py/samples/framework-realtime-tracing-demo - ./run.sh # This sets GENKIT_ENABLE_REALTIME_TELEMETRY=true - ``` - -3. **Open DevUI** at http://localhost:4000 - -4. **Test realtime tracing**: - - [ ] Open the Traces tab in DevUI - - [ ] Trigger a multi-step flow - - [ ] Watch spans appear IMMEDIATELY as they start - - [ ] Compare to non-realtime (spans appear at end) - -5. **Test flows**: - - [ ] `multi_step_flow` - See each step appear in order - - [ ] `nested_flow` - See parent/child span hierarchy - - [ ] `long_running_flow` - Watch progress of slow tasks - -6. **Expected behavior**: - - Spans appear in DevUI as soon as they START - - You see "in progress" spans while they're running - - Nested spans show proper parent/child relationships - - Long-running spans show duration updating in real-time - -Environment Variables -===================== - GENKIT_ENABLE_REALTIME_TELEMETRY=true # Enable realtime tracing - GENKIT_TELEMETRY_SERVER=http://... # Telemetry server URL (auto-set) -""" +"""Realtime tracing demo - spans appear in DevUI as they start, not when they end.""" import asyncio import os import sys -from pydantic import BaseModel, Field - -from genkit.ai import Genkit -from genkit.core.logging import get_logger -from genkit.core.trace import is_realtime_telemetry_enabled +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI -from samples.shared.logging import setup_sample - -setup_sample() - -logger = get_logger(__name__) - - -def _ensure_api_key() -> None: - """Prompt for GEMINI_API_KEY if not set.""" - if not os.environ.get('GEMINI_API_KEY'): - try: - api_key = input('Enter your Gemini API key: ').strip() - if api_key: - os.environ['GEMINI_API_KEY'] = api_key - else: - sys.exit(1) - except (EOFError, KeyboardInterrupt): - sys.exit(1) - - -_ensure_api_key() - -# Initialize Genkit -ai = Genkit( - plugins=[GoogleAI()], - model='googleai/gemini-2.0-flash', -) - - -class MultiStepInput(BaseModel): - """Input for multi-step flow.""" - - topic: str = Field(default='Python programming', description='Topic to process') - - -class NestedInput(BaseModel): - """Input for nested operations flow.""" - - depth: int = Field(default=3, description='Depth of nesting') - - -class ParallelInput(BaseModel): - """Input for parallel tasks flow.""" - - num_tasks: int = Field(default=3, description='Number of parallel tasks') - - -class LlmChainInput(BaseModel): - """Input for LLM chain flow.""" - - initial_prompt: str = Field(default='Tell me a fun fact', description='Initial prompt for the chain') - -@ai.flow(name='slow_multi_step') -async def slow_multi_step_flow(input: MultiStepInput) -> dict[str, object]: - """A multi-step flow with delays to demonstrate realtime tracing. +if not os.environ.get('GEMINI_API_KEY'): + try: + os.environ['GEMINI_API_KEY'] = input('Enter your Gemini API key: ').strip() + except (EOFError, KeyboardInterrupt): + sys.exit(1) + if not os.environ['GEMINI_API_KEY']: + sys.exit(1) - Watch the DevUI as each step appears immediately when it starts! +ai = Genkit(plugins=[GoogleAI()], model='googleai/gemini-2.0-flash') - Args: - input: Input with topic to process. - Returns: - A dict with results from each step. - """ - results = {} +@ai.flow(name='realtime_demo') +async def realtime_demo(topic: str = 'Python') -> str: + """Multi-step flow: watch spans appear in DevUI as each step starts.""" - # Step 1: Research (appears immediately in DevUI) - logger.info('Starting Step 1: Research', topic=input.topic) - research = await ai.run( - 'research', - lambda: slow_operation(f'Researching {input.topic}', delay=2.0), - ) - results['research'] = research + async def research() -> str: + await asyncio.sleep(2) + return f'Researched {topic}' - # Step 2: Analysis (appears as soon as Step 1 completes) - logger.info('Starting Step 2: Analysis') - analysis = await ai.run( - 'analysis', - lambda: slow_operation('Analyzing research findings', delay=1.5), - ) - results['analysis'] = analysis + async def summarize() -> str: + await asyncio.sleep(1) + return f'Summarized {topic}' - # Step 3: Generate Summary (uses actual LLM) - logger.info('Starting Step 3: Generate Summary with LLM') - response = await ai.generate( - prompt=f'Write a one-sentence summary about {input.topic}.', - config={'temperature': 0.7}, - ) - results['summary'] = response.text - - return results - - -@ai.flow(name='nested_operations') -async def nested_operations_flow(input: NestedInput) -> str: - """A flow with nested operations to show parent/child relationships. - - In the DevUI, you'll see the hierarchy of spans as they execute. - - Args: - input: Input with depth of nesting. - - Returns: - A completion message. - """ - - async def nested_step(level: int) -> str: - """Recursive nested operation.""" - if level <= 0: - return 'Done!' - - return await ai.run( - f'level_{level}', - lambda: nested_step(level - 1), - ) - - result = await nested_step(input.depth) - return f'Completed {input.depth} levels: {result}' - - -@ai.flow(name='parallel_tasks') -async def parallel_tasks_flow(input: ParallelInput) -> list[str]: - """Run multiple tasks in parallel to see concurrent spans. - - In the DevUI with realtime tracing, you'll see all tasks - start simultaneously and complete at different times. - - Args: - input: Input with number of parallel tasks. - - Returns: - List of results from each task. - """ - tasks = [] - - for i in range(input.num_tasks): - delay = 1.0 + (i * 0.5) # Staggered completion times - - async def task_fn(task_id: int = i, task_delay: float = delay) -> str: - await asyncio.sleep(task_delay) - return f'Task {task_id} completed after {task_delay}s' - - tasks.append(ai.run(f'parallel_task_{i}', task_fn)) - - results = await asyncio.gather(*tasks) - return list(results) - - -@ai.flow(name='llm_chain') -async def llm_chain_flow(input: LlmChainInput) -> dict[str, object]: - """Chain multiple LLM calls to see sequential model invocations. - - Each model call will appear as a separate span in the DevUI. - - Args: - input: Input with initial prompt. - - Returns: - Dict with responses from each step. - """ - results: dict[str, object] = {} - - # Step 1: Initial generation - response1 = await ai.generate( - prompt=input.initial_prompt, - config={'maxOutputTokens': 100}, - ) - results['fact'] = response1.text - - # Step 2: Follow-up question - response2 = await ai.generate( - prompt=f'Based on this fact: "{response1.text[:100]}...", ask a follow-up question.', - config={'maxOutputTokens': 50}, - ) - results['question'] = response2.text - - # Step 3: Answer the follow-up - response3 = await ai.generate( - prompt=f'Answer this question: {response2.text}', - config={'maxOutputTokens': 100}, - ) - results['answer'] = response3.text - - return results - - -@ai.flow(name='check_realtime_status') -async def check_realtime_status() -> dict[str, object]: - """Check if realtime tracing is enabled. - - Returns: - Status information about realtime tracing. - """ - enabled = is_realtime_telemetry_enabled() - telemetry_server = os.environ.get('GENKIT_TELEMETRY_SERVER', 'Not set') - - return { - 'realtime_enabled': enabled, - 'telemetry_server': telemetry_server, - 'env_var': os.environ.get('GENKIT_ENABLE_REALTIME_TELEMETRY', 'Not set'), - 'message': ( - 'Realtime tracing is ENABLED! Spans appear immediately in DevUI.' - if enabled - else 'Realtime tracing is DISABLED. Set GENKIT_ENABLE_REALTIME_TELEMETRY=true to enable.' - ), - } - - -async def slow_operation(description: str, delay: float = 1.0) -> str: - """Simulate a slow operation. - - Args: - description: What the operation is doing. - delay: How long to wait in seconds. - - Returns: - A completion message. - """ - logger.info('Starting slow operation', description=description, delay=delay) - await asyncio.sleep(delay) - logger.info('Completed slow operation', description=description) - return f'Completed: {description}' + step1 = await ai.run(name='research', fn=research) + step2 = await ai.run(name='summarize', fn=summarize) + response = await ai.generate(prompt=f'One sentence about {topic}.', config={'max_output_tokens': 50}) + return f'{step1} β†’ {step2} β†’ {response.text}' async def main() -> None: - """Main entry point - keeps the server running for DevUI.""" - enabled = is_realtime_telemetry_enabled() - if enabled: - await logger.ainfo('Realtime tracing ENABLED. Spans appear in DevUI immediately.') - else: - await logger.ainfo('Realtime tracing DISABLED. Set GENKIT_ENABLE_REALTIME_TELEMETRY=true.') - - await logger.ainfo('Realtime Tracing Demo running. Press Ctrl+C to stop.') - # Keep the process alive for Dev UI - await asyncio.Event().wait() + """Run the demo flow once, then keep DevUI running.""" + result = await realtime_demo() + print(result) # noqa: T201 - demo output if __name__ == '__main__': diff --git a/py/samples/framework-restaurant-demo/pyproject.toml b/py/samples/framework-restaurant-demo/pyproject.toml index 4306fdf244..aff3718a6d 100644 --- a/py/samples/framework-restaurant-demo/pyproject.toml +++ b/py/samples/framework-restaurant-demo/pyproject.toml @@ -41,9 +41,7 @@ classifiers = [ dependencies = [ "rich>=13.0.0", "genkit", - "genkit-plugin-dev-local-vectorstore", - "genkit-plugin-firebase", - "genkit-plugin-google-cloud", + "genkit-plugin-google-cloud", "genkit-plugin-google-genai", "genkit-plugin-ollama", "genkit-plugin-vertex-ai", @@ -69,7 +67,6 @@ packages = ["src"] [tool.uv.sources] genkit = { workspace = true } -genkit-plugin-dev-local-vectorstore = { workspace = true } genkit-plugin-firebase = { workspace = true } genkit-plugin-google-cloud = { workspace = true } genkit-plugin-google-genai = { workspace = true } diff --git a/py/samples/framework-restaurant-demo/src/case_02/flows.py b/py/samples/framework-restaurant-demo/src/case_02/flows.py index cd7ffcbcbe..ec058866d4 100644 --- a/py/samples/framework-restaurant-demo/src/case_02/flows.py +++ b/py/samples/framework-restaurant-demo/src/case_02/flows.py @@ -39,7 +39,7 @@ async def s02_menu_question_flow( >>> await s02_menu_question_flow(MenuQuestionInputSchema(question='What is the special?')) AnswerOutputSchema(answer="Today's special is...") """ - text = await s02_data_menu_prompt({'question': my_input.question}) + text = await s02_data_menu_prompt({'question': my_input.question}) # pyrefly: ignore[bad-argument-type] return AnswerOutputSchema( answer=text.text, ) diff --git a/py/samples/framework-restaurant-demo/src/case_02/tools.py b/py/samples/framework-restaurant-demo/src/case_02/tools.py index 63a4b5947b..ee533a814f 100644 --- a/py/samples/framework-restaurant-demo/src/case_02/tools.py +++ b/py/samples/framework-restaurant-demo/src/case_02/tools.py @@ -30,7 +30,7 @@ @ai.tool(name='todaysMenu') -def todays_menu(input: object | None = None) -> MenuToolOutputSchema: +async def todays_menu(input: object | None = None) -> MenuToolOutputSchema: """Use this tool to retrieve all the items on today's menu. Args: diff --git a/py/samples/framework-restaurant-demo/src/case_03/chats.py b/py/samples/framework-restaurant-demo/src/case_03/chats.py index 04c663ede3..a9818e27bf 100644 --- a/py/samples/framework-restaurant-demo/src/case_03/chats.py +++ b/py/samples/framework-restaurant-demo/src/case_03/chats.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field -from genkit.core.typing import Message +from genkit.model import Message class ChatSessionInputSchema(BaseModel): diff --git a/py/samples/framework-restaurant-demo/src/case_03/flows.py b/py/samples/framework-restaurant-demo/src/case_03/flows.py index 8c6db55d37..79d9b4b0a7 100644 --- a/py/samples/framework-restaurant-demo/src/case_03/flows.py +++ b/py/samples/framework-restaurant-demo/src/case_03/flows.py @@ -23,7 +23,8 @@ from menu_ai import ai -from genkit.core.typing import Message, Part, Role, TextPart +from genkit._core._typing import Part, Role, TextPart +from genkit.model import Message from genkit.plugins.google_genai.models.gemini import GoogleAIGeminiVersion as GeminiVersion from .chats import ( diff --git a/py/samples/framework-restaurant-demo/src/case_04/example.indexMenuItems.json b/py/samples/framework-restaurant-demo/src/case_04/example.indexMenuItems.json deleted file mode 100644 index d528d8ae9f..0000000000 --- a/py/samples/framework-restaurant-demo/src/case_04/example.indexMenuItems.json +++ /dev/null @@ -1,57 +0,0 @@ -[ - { - "title": "White Meat Crispy Chicken Wings", - "description": "All-white meat chicken wings tossed in your choice of wing sauce. Choose from classic buffalo, honey bbq, garlic parmesan, or sweet & sour", - "price": 12.0 - }, - { - "title": "Cheese Fries", - "description": "Fresh fries covered with melted cheddar cheese and bacon", - "price": 8.0 - }, - { - "title": "Reuben", - "description": "Classic Reuben sandwich with corned beef, sauerkraut, Swiss cheese, and Thousand Island dressing on grilled rye bread.", - "price": 12.0 - }, - { - "title": "Grilled Chicken Club Wrap", - "description": "Grilled chicken, bacon, lettuce, tomato, pickles, and cheddar cheese wrapped in a spinach tortilla, served with your choice of dressing", - "price": 12.0 - }, - { - "title": "Buffalo Chicken Sandwich", - "description": "Fried chicken breast coated in your choice of wing sauce, topped with lettuce, tomato, onion, and pickles on a toasted brioche roll.", - "price": 12.0 - }, - { - "title": "Half Cuban Sandwich", - "description": "Slow roasted pork butt, ham, Swiss, and yellow mustard on a toasted baguette", - "price": 12.0 - }, - { - "title": "The Albie Burger", - "description": "Classic burger topped with bacon, provolone, banana peppers, and chipotle mayo", - "price": 13.0 - }, - { - "title": "57 Chevy Burger", - "description": "Heaven burger with your choice of cheese", - "price": 14.0 - }, - { - "title": "Chicken Caesar Wrap", - "description": "Tender grilled chicken, romaine lettuce, croutons, and Parmesan cheese tossed in a creamy Caesar dressing and wrapped in a spinach tortilla", - "price": 10.0 - }, - { - "title": "Kids Hot Dog", - "description": "Kids under 12", - "price": 5.0 - }, - { - "title": "Chicken Fingers", - "description": "Tender chicken strips, grilled or fried", - "price": 8.0 - } -] diff --git a/py/samples/framework-restaurant-demo/src/case_04/example.menuQuestion.json b/py/samples/framework-restaurant-demo/src/case_04/example.menuQuestion.json deleted file mode 100644 index 21c7c4647c..0000000000 --- a/py/samples/framework-restaurant-demo/src/case_04/example.menuQuestion.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "question": "I'd like something cheesy!" -} diff --git a/py/samples/framework-restaurant-demo/src/case_04/flows.py b/py/samples/framework-restaurant-demo/src/case_04/flows.py deleted file mode 100644 index d88255ead5..0000000000 --- a/py/samples/framework-restaurant-demo/src/case_04/flows.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Flows for case 04.""" - -import json -import os -import pathlib - -from menu_ai import ai -from menu_schemas import AnswerOutputSchema, MenuItemSchema, MenuQuestionInputSchema -from pydantic import BaseModel, Field - -from genkit.blocks.document import Document - -from .prompts import s04_rag_data_menu_prompt - - -class IndexMenuItemsOutputSchema(BaseModel): - """Output schema for indexing items.""" - - rows: int = Field(...) - - -@ai.flow(name='s04_index_menu_items') -async def s04_index_menu_items_flow( - menu_items: list[MenuItemSchema], -) -> IndexMenuItemsOutputSchema: - """Index menu items for retrieval. - - Args: - menu_items: List of menu items to index. - - Returns: - Number of items indexed. - - Example: - >>> await s04_index_menu_items_flow([MenuItemSchema(title='Burger', price=10.0, description='Yum')]) - IndexMenuItemsOutputSchema(rows=1) - """ - # If empty list provided (e.g., from Dev UI default), load from example file - # Filter out None values that may come from test input like [null] - menu_items = [item for item in menu_items if item is not None] - - if not menu_items: - example_file = os.path.join(pathlib.Path(__file__).parent, 'example.indexMenuItems.json') - with pathlib.Path(example_file).open() as f: - menu_data = json.load(f) - menu_items = [MenuItemSchema(**item) for item in menu_data] - - documents = [ - Document.from_text(f'{item.title} {item.price} \n {item.description}', metadata=item.model_dump()) - for item in menu_items - ] - - await ai.index( - indexer='menu-items', - documents=documents, - ) - return IndexMenuItemsOutputSchema(rows=len(menu_items)) - - -@ai.flow(name='s04_rag_menu_question') -async def s04_rag_menu_question_flow( - my_input: MenuQuestionInputSchema, -) -> AnswerOutputSchema: - """Answer a question using RAG on menu items. - - Args: - my_input: Input containing the question. - - Returns: - The answer. - - Example: - >>> await s04_rag_menu_question_flow(MenuQuestionInputSchema(question='Do you have burgers?')) - AnswerOutputSchema(answer="Yes, we have...") - """ - # Retrieve the 3 most relevant menu items for the question - docs = await ai.retrieve( - retriever='menu-items', - query=my_input.question, - options={'limit': 3}, - ) - - menu_data = [MenuItemSchema(**doc.metadata) for doc in docs.documents if doc.metadata] - - # Generate the response - response = await s04_rag_data_menu_prompt({ - 'menuData': [item.model_dump() for item in menu_data], - 'question': my_input.question, - }) - return AnswerOutputSchema(answer=response.text) diff --git a/py/samples/framework-restaurant-demo/src/case_04/prompts.py b/py/samples/framework-restaurant-demo/src/case_04/prompts.py deleted file mode 100644 index eac543dc78..0000000000 --- a/py/samples/framework-restaurant-demo/src/case_04/prompts.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 -"""Prompts for case 04.""" - -from menu_ai import ai -from menu_schemas import DataMenuQuestionInputSchema - -from genkit.plugins.google_genai.models.gemini import GoogleAIGeminiVersion - -s04_rag_data_menu_prompt = ai.define_prompt( - variant='s04_ragDataMenu', - model=f'googleai/{GoogleAIGeminiVersion.GEMINI_3_FLASH_PREVIEW}', - input_schema=DataMenuQuestionInputSchema, - config={'temperature': 0.3}, - prompt=""" -You are acting as Walt, a helpful AI assistant here at the restaurant. -You can answer questions about the food on the menu or any other questions -customers have about food in general. - -Here are some items that are on today's menu that are relevant to -helping you answer the customer's question: -{{#each menuData~}} -- {{this.title}} ${{this.price}} - {{this.description}} -{{~/each}} - -Answer this customer's question: -{{question}}? -""", -) diff --git a/py/samples/framework-restaurant-demo/src/case_05/flows.py b/py/samples/framework-restaurant-demo/src/case_05/flows.py index 54391ab766..32c95f7a67 100644 --- a/py/samples/framework-restaurant-demo/src/case_05/flows.py +++ b/py/samples/framework-restaurant-demo/src/case_05/flows.py @@ -47,7 +47,7 @@ async def s05_read_menu_flow(_: None = None) -> str: "Menu content..." """ image_data_url = inline_data_url('menu.jpeg', 'image/jpeg') - response = await s05_read_menu_prompt({'image_url': image_data_url}) + response = await s05_read_menu_prompt({'image_url': image_data_url}) # pyrefly: ignore[bad-argument-type] return response.text @@ -67,7 +67,7 @@ async def s05_text_menu_question_flow( >>> await s05_text_menu_question_flow(TextMenuQuestionInputSchema(menu_text='Burger: $10', question='Price?')) AnswerOutputSchema(answer="It costs $10") """ - response = await s05_text_menu_prompt({'menuText': my_input.menu_text, 'question': my_input.question}) + response = await s05_text_menu_prompt(my_input) return AnswerOutputSchema( answer=response.text, ) diff --git a/py/samples/framework-restaurant-demo/src/main.py b/py/samples/framework-restaurant-demo/src/main.py index a1b21806cb..78ac07f952 100755 --- a/py/samples/framework-restaurant-demo/src/main.py +++ b/py/samples/framework-restaurant-demo/src/main.py @@ -67,10 +67,8 @@ flows as case_03_flows, # noqa: F401 prompts as case_03_prompts, # noqa: F401 ) -from case_04 import ( - flows as case_04_flows, # noqa: F401 - prompts as case_04_prompts, # noqa: F401 -) + +# case_04 module does not exist - skipped from case_05 import ( flows as case_05_flows, # noqa: F401 prompts as case_05_prompts, # noqa: F401 diff --git a/py/samples/framework-restaurant-demo/src/menu_ai.py b/py/samples/framework-restaurant-demo/src/menu_ai.py index 3c3460dfb4..35640ebbe3 100644 --- a/py/samples/framework-restaurant-demo/src/menu_ai.py +++ b/py/samples/framework-restaurant-demo/src/menu_ai.py @@ -19,8 +19,7 @@ import os -from genkit.ai import Genkit -from genkit.plugins.dev_local_vectorstore import define_dev_local_vector_store +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI if 'GEMINI_API_KEY' not in os.environ: @@ -33,10 +32,3 @@ ai = Genkit(plugins=[GoogleAI()]) - -# Define dev local vector store -define_dev_local_vector_store( - ai, - name='menu-items', - embedder='googleai/gemini-embedding-001', -) diff --git a/py/samples/framework-tool-interrupts/src/main.py b/py/samples/framework-tool-interrupts/src/main.py index e7ad9c9e95..dfe37a4d3a 100755 --- a/py/samples/framework-tool-interrupts/src/main.py +++ b/py/samples/framework-tool-interrupts/src/main.py @@ -90,7 +90,7 @@ from pydantic import BaseModel, Field -from genkit.ai import ( +from genkit import ( Genkit, ToolRunContext, tool_response, @@ -115,7 +115,7 @@ class TriviaQuestions(BaseModel): @ai.tool() -def present_questions(questions: TriviaQuestions, ctx: ToolRunContext) -> None: +async def present_questions(questions: TriviaQuestions, ctx: ToolRunContext) -> None: """Presents questions to the user and responds with the selected answer.""" ctx.interrupt(questions.model_dump()) diff --git a/py/samples/provider-amazon-bedrock-hello/LICENSE b/py/samples/provider-amazon-bedrock-hello/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/samples/provider-amazon-bedrock-hello/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/samples/provider-amazon-bedrock-hello/README.md b/py/samples/provider-amazon-bedrock-hello/README.md deleted file mode 100644 index 7bbe287330..0000000000 --- a/py/samples/provider-amazon-bedrock-hello/README.md +++ /dev/null @@ -1,306 +0,0 @@ -# Amazon Bedrock Hello Sample - -> **Community Plugin** – This plugin is maintained by the community and is supported on a best-effort basis. It is not an official AWS product. - -This sample demonstrates how to use Amazon Bedrock models with Genkit, including AWS X-Ray telemetry for distributed tracing. - -## Prerequisites - -1. **AWS Account** with Bedrock access enabled -2. **AWS Credentials** configured (one of the following): - - **Bedrock API Key** (`AWS_BEARER_TOKEN_BEDROCK`) - Simplest option, like OpenAI - - IAM credentials (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`) - - AWS credentials file (`~/.aws/credentials`) - - IAM role (for EC2, Lambda, ECS, etc.) -3. **Model Access** enabled in Amazon Bedrock console for the models you want to use - -## Setup - -### Option A: Bedrock API Key (Simplest - Recommended for Development) - -Amazon Bedrock now supports API keys similar to OpenAI/Anthropic. This is the simplest way to get started. - -**Step 1: Generate an API Key** - -1. Go to [Amazon Bedrock Console](https://console.aws.amazon.com/bedrock/) -2. Select your region (e.g., `us-east-1`) -3. Click **API keys** in the left sidebar (under "Bedrock configurations") -4. Click **Generate API key** -5. Choose expiration (30 days recommended for development) -6. Copy the generated key (shown only once!) - -**Step 2: Set Environment Variables** - -```bash -export AWS_REGION="us-east-1" -export AWS_BEARER_TOKEN_BEDROCK="your-api-key-here" -``` - -That's it! No IAM user or access keys needed. - -**Limitations of API Keys:** -- For development/exploration only (not recommended for production) -- Cannot be used with Bedrock Agents or Data Automation -- Keys expire based on your configuration - -See: [Getting Started with Bedrock API Keys](https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started-api-keys.html) - ---- - -### Option B: IAM Credentials (Recommended for Production) - -#### Step 1: Create AWS Credentials - -If you don't have AWS credentials yet, follow these steps: - -1. **Sign in to AWS Console**: Go to [AWS Console](https://console.aws.amazon.com/) - -2. **Navigate to IAM**: - - Search for "IAM" in the AWS Console search bar - - Or go directly to [IAM Console](https://console.aws.amazon.com/iam/) - -3. **Create an IAM User** (recommended for development): - - Click **Users** in the left sidebar - - Click **Create user** - - Enter a username (e.g., `genkit-bedrock-dev`) - - Click **Next** - -4. **Attach Permissions**: - - Select **Attach policies directly** - - Search for and select `AmazonBedrockFullAccess` - - Click **Next**, then **Create user** - -5. **Create Access Keys**: - - Click on the user you just created - - Go to the **Security credentials** tab - - Scroll to **Access keys** and click **Create access key** - - Select **Local code** as the use case - - Click **Create access key** - - **Important**: Copy both the **Access key ID** and **Secret access key** - - These are shown only once! Save them securely. - -### Step 2: Enable Model Access - -1. Go to the [Amazon Bedrock Console](https://console.aws.amazon.com/bedrock/) -2. Select your region (e.g., `us-east-1`) from the top-right dropdown -3. Click **Model access** in the left sidebar (under "Bedrock configurations") -4. Click **Modify model access** -5. Check the models you want to use: - - **Anthropic**: Claude Sonnet, Claude Haiku (recommended for testing) - - **Amazon**: Nova Pro, Nova Lite, Nova Micro - - **Meta**: Llama 3.3, Llama 4 -6. Click **Next**, review, and click **Submit** -7. Wait for status to change from "In progress" to "Access granted" - -### Step 3: Configure Environment Variables - -Set your credentials as environment variables: - -```bash -# Required: AWS Region where you enabled model access -export AWS_REGION="us-east-1" - -# Required: Your IAM user credentials (from Step 1) -export AWS_ACCESS_KEY_ID="AKIA..." # Starts with AKIA -export AWS_SECRET_ACCESS_KEY="wJalrXUt..." # Your secret key -``` - -**Tip**: Add these to your shell profile (`~/.bashrc`, `~/.zshrc`) to persist them: - -```bash -# Add to ~/.zshrc or ~/.bashrc -export AWS_REGION="us-east-1" -export AWS_ACCESS_KEY_ID="AKIA..." -export AWS_SECRET_ACCESS_KEY="..." -``` - -### Option C: Using AWS CLI Profile - -If you prefer using the AWS CLI: - -```bash -# Install AWS CLI (if not already installed) -# macOS: brew install awscli -# Linux: pip install awscli - -# Configure credentials -aws configure -# Enter your Access Key ID, Secret Access Key, and region when prompted - -# Then just set the region for this sample -export AWS_REGION="us-east-1" -``` - -### Option D: IAM Role (for AWS Infrastructure) - -If running on EC2, Lambda, ECS, or EKS: - -```bash -# Only the region is needed - credentials come from the IAM role -export AWS_REGION="us-east-1" -``` - -Make sure the IAM role has the `AmazonBedrockFullAccess` policy attached. - -## Running the Sample - -```bash -./run.sh -``` - -This will start the Genkit Dev UI and the sample application with hot reloading. - -## Features Demonstrated - -| Feature | Flow Name | Description | -|---------|-----------|-------------| -| Simple Generation | `say_hi` | Basic text generation | -| Streaming | `say_hi_stream` | Streaming text generation | -| Tool Use | `weather_flow` | Function calling with tools | -| Multi-turn Chat | `chat_demo` | Multi-turn conversation | -| Structured Output | `generate_character` | Generate JSON-structured output | -| Multimodal | `describe_image` | Image description (Claude, Nova) | -| Embeddings | `embed_text` | Text embeddings (Titan, Cohere) | -| AWS X-Ray Telemetry | `add_aws_telemetry()` | Distributed tracing to X-Ray console | - -## Supported Models - -The sample uses Claude Sonnet 4.5 by default and auto-detects your authentication method. - -### Authentication & Model IDs - -| Auth Method | Model ID Format | Example | -|-------------|-----------------|---------| -| IAM Credentials | Direct model ID | `anthropic.claude-sonnet-4-5-...` | -| API Key (Bearer Token) | Inference profile | `us.anthropic.claude-sonnet-4-5-...` | - -**Important**: When using API keys (`AWS_BEARER_TOKEN_BEDROCK`), you must use **inference profiles** with a regional prefix (`us.`, `eu.`, or `apac.`). - -### Using Pre-defined Model References (IAM Credentials) - -```python -from genkit.plugins.amazon_bedrock import claude_sonnet_4_5, nova_pro, llama_3_3_70b - -# Pre-defined references use direct model IDs -ai = Genkit(model=claude_sonnet_4_5) -``` - -### Using Inference Profiles (API Keys) - -```python -from genkit.plugins.amazon_bedrock import inference_profile - -# inference_profile() auto-detects region from AWS_REGION -model = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0') - -# Or specify region explicitly -model = inference_profile('anthropic.claude-sonnet-4-5-20250929-v1:0', 'eu-west-1') -``` - -### Regional Prefixes - -| Region | Prefix | Example Regions | -|--------|--------|-----------------| -| US | `us.` | us-east-1, us-west-2 | -| Europe | `eu.` | eu-west-1, eu-central-1 | -| Asia Pacific | `apac.` | ap-northeast-1, ap-southeast-1 | - -### Available Models - -- **Claude (Anthropic)**: claude-sonnet-4-5, claude-opus-4-5, claude-haiku-4-5 -- **Nova (Amazon)**: nova-pro, nova-lite, nova-micro -- **Llama (Meta)**: llama-3.3-70b, llama-4-maverick -- **Mistral**: mistral-large-3, pixtral-large -- **DeepSeek**: deepseek-r1, deepseek-v3 -- **Cohere**: command-r-plus, command-r - -See: -- [Amazon Bedrock Supported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) -- [Cross-Region Inference Profiles](https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html) - -## Testing the Flows - -Once the Dev UI is running, you can test each flow: - -1. Open the Dev UI (automatically opens in browser) -2. Select a flow from the sidebar -3. Click **Run** to execute with default inputs -4. Modify inputs to test different scenarios - -## Troubleshooting - -### "NoCredentialsError" or "Unable to locate credentials" - -Your AWS credentials are not configured. Make sure you have set: - -```bash -export AWS_ACCESS_KEY_ID="AKIA..." -export AWS_SECRET_ACCESS_KEY="..." -``` - -Or configured credentials via `aws configure`. - -### "InvalidAccessKeyId" or "SignatureDoesNotMatch" - -Your credentials are invalid or incorrectly copied. Double-check: -- The Access Key ID starts with `AKIA` -- No extra spaces in the values -- The secret key was copied completely - -### "AccessDeniedException" Error - -This can mean: -1. **Model access not enabled**: Go to Bedrock Console > Model access and enable the model -2. **IAM permissions missing**: Make sure your IAM user has `AmazonBedrockFullAccess` policy -3. **Wrong region**: The model may not be available in your region - -### "Region not supported" Error - -Not all models are available in all regions. Check [model availability by region](https://docs.aws.amazon.com/bedrock/latest/userguide/models-regions.html). - -Popular regions with most models: -- `us-east-1` (N. Virginia) -- `us-west-2` (Oregon) -- `eu-west-1` (Ireland) - -### Timeout Errors - -For Amazon Nova models, the inference timeout can be up to 60 minutes. The plugin automatically configures appropriate timeouts. - -### "ValidationException: Malformed input request" - -This usually means the model ID is incorrect. Check the exact model ID in the Bedrock console under "Model access". - -### "ValidationException: Invocation of model ID ... with on-demand throughput isn't supported" - -This error occurs when using API keys (`AWS_BEARER_TOKEN_BEDROCK`) with a direct model ID instead of an inference profile. - -**Solution**: Use the cross-region inference profile ID with a regional prefix: - -```python -# Wrong (direct model ID): -model = 'anthropic.claude-sonnet-4-5-20250929-v1:0' - -# Correct (inference profile with us. prefix): -model = 'us.anthropic.claude-sonnet-4-5-20250929-v1:0' - -# Or use the pre-defined references (already use inference profiles): -from genkit.plugins.amazon_bedrock import claude_sonnet_4_5 -model = claude_sonnet_4_5 -``` - -## Security Best Practices - -1. **Never commit credentials**: Add AWS credentials to `.gitignore`, never commit them -2. **Use IAM roles in production**: On AWS infrastructure, use IAM roles instead of access keys -3. **Rotate keys regularly**: Periodically rotate your access keys in IAM console -4. **Least privilege**: Create IAM users with only the permissions needed - -## Resources - -- [Amazon Bedrock Documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) -- [Supported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) -- [Model Parameters](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html) -- [Converse API](https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html) -- [IAM Getting Started](https://docs.aws.amazon.com/IAM/latest/UserGuide/getting-started.html) -- [AWS CLI Configuration](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html) diff --git a/py/samples/provider-amazon-bedrock-hello/pyproject.toml b/py/samples/provider-amazon-bedrock-hello/pyproject.toml deleted file mode 100644 index a7beae7852..0000000000 --- a/py/samples/provider-amazon-bedrock-hello/pyproject.toml +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Private :: Do Not Upload", -] -dependencies = [ - "rich>=13.0.0", - "genkit", - "genkit-plugin-amazon-bedrock", - "pydantic>=2.0.0", - "structlog>=24.0.0", - "uvloop>=0.21.0", -] -description = "Amazon Bedrock Hello Sample" -license = "Apache-2.0" -name = "provider-amazon-bedrock-hello" -requires-python = ">=3.10" -version = "0.2.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[tool.uv.sources] -genkit-plugin-amazon-bedrock = { workspace = true } - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src"] diff --git a/py/samples/provider-amazon-bedrock-hello/run.sh b/py/samples/provider-amazon-bedrock-hello/run.sh deleted file mode 100755 index 239414b7a1..0000000000 --- a/py/samples/provider-amazon-bedrock-hello/run.sh +++ /dev/null @@ -1,113 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -# AWS Bedrock Hello World Demo -# ============================ -# -# Demonstrates usage of AWS Bedrock models with Genkit. -# -# Prerequisites: -# - AWS credentials configured (env vars, credentials file, or IAM role) -# - AWS_REGION environment variable set -# -# Usage: -# ./run.sh # Start the demo with Dev UI -# ./run.sh --help # Show this help message - -set -euo pipefail - -cd "$(dirname "$0")" -source "../_common.sh" - -print_help() { - print_banner "AWS Bedrock Hello World" "☁️" - echo "Usage: ./run.sh [options]" - echo "" - echo "Options:" - echo " --help Show this help message" - echo "" - echo "Environment Variables:" - echo " AWS_REGION Required. AWS region (e.g., us-east-1)" - echo " AWS_ACCESS_KEY_ID AWS access key ID (or use credentials file)" - echo " AWS_SECRET_ACCESS_KEY AWS secret access key (or use credentials file)" - echo " AWS_PROFILE AWS profile name from credentials file" - echo "" - echo "Setup Guide: https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html" - print_help_footer -} - -case "${1:-}" in - --help|-h) - print_help - exit 0 - ;; -esac - -print_banner "AWS Bedrock Hello World" "☁️" - -# Check AWS CLI is installed (offers to install if missing) -check_aws_installed || true - -# Set default region if not provided -export AWS_REGION="${AWS_REGION:-us-east-1}" - -check_env_var "AWS_REGION" "https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html" - -# AWS authentication: support three methods -# 1. API key (AWS_BEARER_TOKEN_BEDROCK) β€” simplest, requires inference profiles -# 2. IAM credentials (AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY) β€” explicit keys -# 3. AWS CLI profile / credentials file β€” uses `aws configure` -# -# We check for explicit env vars first, then fall back to AWS CLI auth check. -if [[ -z "${AWS_BEARER_TOKEN_BEDROCK:-}" && -z "${AWS_ACCESS_KEY_ID:-}" ]]; then - echo -e "${YELLOW}No AWS credentials found in environment.${NC}" - echo "" - echo "Choose an authentication method:" - echo " 1. AWS CLI credentials (aws configure)" - echo " 2. Set AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY" - echo " 3. Set AWS_BEARER_TOKEN_BEDROCK (API key)" - echo "" - - if [[ -t 0 ]] && [ -c /dev/tty ]; then - echo -en "Enter choice [${YELLOW}1${NC}]: " - choice="" - read -r choice < /dev/tty - choice="${choice:-1}" - - case "$choice" in - 2) - check_env_var "AWS_ACCESS_KEY_ID" "https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html" - check_env_var "AWS_SECRET_ACCESS_KEY" "https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html" - ;; - 3) - check_env_var "AWS_BEARER_TOKEN_BEDROCK" "https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html" - ;; - *) - # Default: use AWS CLI auth (prompts `aws configure` if needed) - check_aws_auth || true - ;; - esac - else - # Non-interactive: try AWS CLI auth silently - check_aws_auth || true - fi -else - if [[ -n "${AWS_BEARER_TOKEN_BEDROCK:-}" ]]; then - echo -e "${GREEN}βœ“ Using API key authentication (AWS_BEARER_TOKEN_BEDROCK)${NC}" - else - echo -e "${GREEN}βœ“ Using IAM credentials (AWS_ACCESS_KEY_ID)${NC}" - fi -fi -echo "" - -install_deps - -genkit_start_with_browser -- \ - uv tool run --from watchdog watchmedo auto-restart \ - -d src \ - -d ../../packages \ - -d ../../plugins \ - -p '*.py;*.prompt;*.json' \ - -R \ - -- uv run src/main.py "$@" diff --git a/py/samples/provider-amazon-bedrock-hello/src/main.py b/py/samples/provider-amazon-bedrock-hello/src/main.py deleted file mode 100644 index b10eb4e034..0000000000 --- a/py/samples/provider-amazon-bedrock-hello/src/main.py +++ /dev/null @@ -1,416 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Amazon Bedrock hello sample - Foundation models and observability with Genkit. - -This sample demonstrates how to use AWS Bedrock models with Genkit, -including tools, streaming, multimodal, embedding, and AWS X-Ray telemetry. - -See README.md for setup and testing instructions. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ AWS Bedrock β”‚ Amazon's AI model marketplace. One place to β”‚ - β”‚ β”‚ access Claude, Llama, Nova, and more. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Converse API β”‚ A unified way to talk to ANY Bedrock model. β”‚ - β”‚ β”‚ Same code works for Claude, Llama, Nova, etc. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Inference Profile β”‚ A cross-region alias for a model. Required β”‚ - β”‚ β”‚ when using API keys instead of IAM roles. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ IAM Role β”‚ AWS's way of granting permissions. Like a β”‚ - β”‚ β”‚ badge that lets your code access models. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Region β”‚ Which AWS data center to use. Pick one near β”‚ - β”‚ β”‚ you (us-east-1, eu-west-1, ap-northeast-1). β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ AWS X-Ray β”‚ Amazon's distributed tracing service. See how β”‚ - β”‚ β”‚ requests flow through your AI application. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Features -============ -| Feature Description | Example Function / Code Snippet | -|-----------------------------------------|------------------------------------------| -| Plugin Initialization | `ai = Genkit(plugins=[AmazonBedrock()])` | -| AWS X-Ray Telemetry | `add_aws_telemetry(region=...)` | -| Default Model Configuration | `ai = Genkit(model=...)` | -| Defining Flows | `@ai.flow()` decorator | -| Defining Tools | `@ai.tool()` decorator | -| Pydantic for Tool Input Schema | `WeatherInput`, `CurrencyInput` | -| Simple Generation (Prompt String) | `generate_greeting` | -| System Prompts | `generate_with_system_prompt` | -| Multi-turn Conversations (`messages`) | `generate_multi_turn_chat` | -| Streaming Generation | `generate_streaming_story` | -| Generation with Tools | `generate_weather`, `convert_currency` | -| Generation Configuration (temperature) | `generate_with_config` | -| Multimodal (Image Input) | `describe_image` | -| Code Generation | `generate_code` | -| Embeddings | `embed_text` | - -Supported Models -================ -- Claude (Anthropic): claude-sonnet-4-5, claude-opus-4-5, claude-haiku-4-5 -- Nova (Amazon): nova-pro, nova-lite, nova-micro -- Llama (Meta): llama-3.3-70b, llama-4-maverick -- Mistral: mistral-large-3, pixtral-large -- DeepSeek: deepseek-r1, deepseek-v3 -- And many more... -""" - -import asyncio -import os - -from genkit.ai import Genkit -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger -from genkit.plugins.amazon_bedrock import ( - AmazonBedrock, - add_aws_telemetry, - bedrock_model, - claude_opus_4_6, - claude_sonnet_4_5, - deepseek_r1, - inference_profile, - nova_pro, -) -from genkit.types import Media, MediaPart, Part, TextPart -from samples.shared import ( - CharacterInput, - CodeInput, - CurrencyExchangeInput, - EmbedInput, - GreetingInput, - ImageDescribeInput, - MultiTurnInput, - ReasoningInput, - RpgCharacter, - StreamingToolInput, - StreamInput, - SystemPromptInput, - WeatherInput, - convert_currency as _convert_currency_tool, - convert_currency_logic, - describe_image_logic, - generate_character_logic, - generate_code_logic, - generate_greeting_logic, - generate_multi_turn_chat_logic, - generate_streaming_story_logic, - generate_streaming_with_tools_logic, - generate_weather_logic, - generate_with_config_logic, - generate_with_system_prompt_logic, - get_weather, - setup_sample, -) - -setup_sample() - -# Prompt for AWS region if not set -if 'AWS_REGION' not in os.environ: - os.environ['AWS_REGION'] = input('Please enter your AWS_REGION (e.g., us-east-1): ') - -logger = get_logger(__name__) - -# Enable AWS X-Ray telemetry (traces exported to X-Ray console) -# This provides distributed tracing for all Genkit flows and model calls -# View traces at: https://console.aws.amazon.com/xray/home -add_aws_telemetry(region=os.environ.get('AWS_REGION')) - -# Default model configuration -# Model IDs without regional prefix - used as base for both auth methods -_CLAUDE_SONNET_MODEL_ID = 'anthropic.claude-sonnet-4-5-20250929-v1:0' -_CLAUDE_OPUS_MODEL_ID = 'anthropic.claude-opus-4-6-20260205-v1:0' -_DEEPSEEK_R1_MODEL_ID = 'deepseek.r1-v1:0' -_NOVA_PRO_MODEL_ID = 'amazon.nova-pro-v1:0' -_TITAN_EMBED_MODEL_ID = 'amazon.titan-embed-text-v2:0' - -# Detect authentication method and choose appropriate model reference -# - API keys (AWS_BEARER_TOKEN_BEDROCK) require inference profiles -# - IAM credentials work with direct model IDs -_using_api_key = 'AWS_BEARER_TOKEN_BEDROCK' in os.environ - -# Choose models based on auth method -# API keys require inference profiles with regional prefix (us., eu., apac.) -# IAM credentials work with direct model IDs -if _using_api_key: - _default_model = inference_profile(_CLAUDE_SONNET_MODEL_ID) - _opus_model = inference_profile(_CLAUDE_OPUS_MODEL_ID) - _deepseek_model = inference_profile(_DEEPSEEK_R1_MODEL_ID) - _nova_model = inference_profile(_NOVA_PRO_MODEL_ID) - _embed_model = inference_profile(_TITAN_EMBED_MODEL_ID) - logger.info('Using API key auth - model IDs will use inference profiles') -else: - _default_model = claude_sonnet_4_5 - _opus_model = claude_opus_4_6 - _deepseek_model = deepseek_r1 - _nova_model = nova_pro - _embed_model = bedrock_model(_TITAN_EMBED_MODEL_ID) - logger.info('Using IAM credentials - model IDs are direct') - -logger.info('AWS X-Ray telemetry enabled - traces visible in X-Ray console') - -# Initialize Genkit with AWS Bedrock plugin -# Region is required - either from env var (prompted above) or explicit -ai = Genkit( - plugins=[AmazonBedrock(region=os.environ.get('AWS_REGION'))], - model=_default_model, -) - -ai.tool()(get_weather) -ai.tool()(_convert_currency_tool) - - -@ai.flow() -async def generate_greeting(input: GreetingInput) -> str: - """Generate a simple greeting. - - Args: - input: Input with name to greet. - - Returns: - Greeting message. - """ - return await generate_greeting_logic(ai, input.name) - - -@ai.flow() -async def generate_with_system_prompt(input: SystemPromptInput) -> str: - """Demonstrate system prompts to control model persona and behavior. - - Args: - input: Input with a question to ask. - - Returns: - The model's response in the persona defined by the system prompt. - """ - return await generate_with_system_prompt_logic(ai, input.question) - - -@ai.flow() -async def generate_multi_turn_chat(input: MultiTurnInput) -> str: - """Demonstrate multi-turn conversations using the messages parameter. - - Args: - input: Input with a travel destination. - - Returns: - The model's final response, demonstrating context retention. - """ - return await generate_multi_turn_chat_logic(ai, input.destination) - - -@ai.flow() -async def generate_streaming_story( - input: StreamInput, - ctx: ActionRunContext = None, # type: ignore[assignment] -) -> str: - """Generate a streaming story response. - - Args: - input: Input with name for streaming story. - ctx: Action run context for streaming. - - Returns: - Complete generated text. - """ - return await generate_streaming_story_logic(ai, input.name, ctx) - - -@ai.flow() -async def generate_with_config(input: GreetingInput) -> str: - """Generate a greeting with custom model configuration. - - Args: - input: Input with name to greet. - - Returns: - Greeting message. - """ - return await generate_with_config_logic(ai, input.name) - - -@ai.flow() -async def generate_weather(input: WeatherInput) -> str: - """Get weather information using tool calling. - - Args: - input: Input with location to get weather for. - - Returns: - Weather information. - """ - return await generate_weather_logic(ai, input) - - -@ai.flow() -async def convert_currency(input: CurrencyExchangeInput) -> str: - """Convert currency using tool calling. - - Args: - input: Currency exchange parameters. - - Returns: - Conversion result. - """ - return await convert_currency_logic(ai, input) - - -@ai.flow() -async def generate_character(input: CharacterInput) -> RpgCharacter: - """Generate an RPG character with structured output. - - Args: - input: Input with character name. - - Returns: - The generated RPG character. - """ - return await generate_character_logic(ai, input.name) - - -@ai.flow() -async def describe_image(input: ImageDescribeInput) -> str: - """Describe an image using Claude or Nova (multimodal models). - - Args: - input: Input with image URL to describe. - - Returns: - A textual description of the image. - """ - return await describe_image_logic(ai, input.image_url) - - -@ai.flow() -async def generate_code(input: CodeInput) -> str: - """Generate code using AWS Bedrock models. - - Args: - input: Input with coding task description. - - Returns: - Generated code. - """ - # NOTE: Claude Opus 4.6 (_opus_model) is ideal for complex code generation, - # but it may not be available in all regions yet (released Feb 5, 2026). - # Swap to _opus_model once it's enabled in your region's inference profiles. - return await generate_code_logic(ai, input.task, model=_default_model) - - -@ai.flow() -async def generate_streaming_with_tools( - input: StreamingToolInput, - ctx: ActionRunContext | None = None, -) -> str: - """Demonstrate streaming generation with tool calling. - - Args: - input: Input with location for weather lookup. - ctx: Action context for streaming chunks to the client. - - Returns: - The complete generated text. - """ - return await generate_streaming_with_tools_logic(ai, input.location, ctx) - - -@ai.flow() -async def describe_image_nova(input: ImageDescribeInput) -> str: - """Describe an image using Amazon Nova Pro. - - When using API keys (AWS_BEARER_TOKEN_BEDROCK), this automatically uses - the inference profile version of the model (e.g., us.amazon.nova-pro-v1:0). - - Args: - input: Input with image URL. - - Returns: - Image description. - """ - response = await ai.generate( - model=_nova_model, - prompt=[ - Part(root=TextPart(text='Describe this image')), - Part(root=MediaPart(media=Media(url=input.image_url, content_type='image/jpeg'))), - ], - ) - return response.text - - -@ai.flow() -async def embed_text(input: EmbedInput) -> list[float]: - """Generate text embeddings using Amazon Titan. - - When using API keys (AWS_BEARER_TOKEN_BEDROCK), this automatically uses - the inference profile version of the model. - - Args: - input: Input with text to embed. - - Returns: - Embedding vector (first 10 dimensions shown). - """ - embeddings = await ai.embed( - embedder=_embed_model, - content=input.text, - ) - # Return first 10 dimensions as a sample - embedding = embeddings[0].embedding if embeddings else [] - return embedding[:10] if len(embedding) > 10 else embedding - - -@ai.flow() -async def reasoning_demo(input: ReasoningInput) -> str: - """Demonstrate reasoning with DeepSeek R1. - - Note: DeepSeek R1 includes reasoning content in the response. - For optimal quality, limit max_tokens to 8,192 or fewer. - - When using API keys (AWS_BEARER_TOKEN_BEDROCK), this automatically uses - the inference profile version of the model (e.g., us.deepseek.r1-v1:0). - - Args: - input: Input with question requiring reasoning. - - Returns: - Answer with reasoning steps. - """ - response = await ai.generate( - model=_deepseek_model, - prompt=input.prompt, - config={ - 'max_tokens': 4096, - 'temperature': 0.5, - }, - ) - return response.text - - -async def main() -> None: - """Main entry point for the AWS Bedrock sample - keep alive for Dev UI.""" - await logger.ainfo('Genkit server running. Press Ctrl+C to stop.') - # Keep the process alive for Dev UI - await asyncio.Event().wait() - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/samples/provider-anthropic-hello/src/main.py b/py/samples/provider-anthropic-hello/src/main.py index bcc5d043a9..5c5eff5a4b 100755 --- a/py/samples/provider-anthropic-hello/src/main.py +++ b/py/samples/provider-anthropic-hello/src/main.py @@ -73,13 +73,12 @@ import asyncio import os +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger +from genkit import Genkit, Media, MediaPart, Message, Metadata, Part, Role, TextPart +from genkit._core._action import ActionRunContext from genkit.plugins.anthropic import Anthropic, anthropic_name -from genkit.types import Media, MediaPart, Message, Metadata, Part, Role, TextPart from samples.shared import ( CharacterInput, CodeInput, @@ -113,7 +112,7 @@ if 'ANTHROPIC_API_KEY' not in os.environ: os.environ['ANTHROPIC_API_KEY'] = input('Please enter your ANTHROPIC_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) ai = Genkit( plugins=[Anthropic()], diff --git a/py/samples/provider-compat-oai-hello/src/main.py b/py/samples/provider-compat-oai-hello/src/main.py index 325829c494..7b91709270 100755 --- a/py/samples/provider-compat-oai-hello/src/main.py +++ b/py/samples/provider-compat-oai-hello/src/main.py @@ -75,12 +75,11 @@ import os import httpx +import structlog from pydantic import BaseModel, Field -from genkit.ai import ActionRunContext, Genkit, Output -from genkit.core.logging import get_logger +from genkit import ActionRunContext, Genkit, Media, MediaPart, Message, Part, Role, TextPart from genkit.plugins.compat_oai import OpenAI, openai_model -from genkit.types import Media, MediaPart, Message, Part, Role, TextPart from samples.shared import ( CharacterInput, CodeInput, @@ -113,7 +112,7 @@ if 'OPENAI_API_KEY' not in os.environ: os.environ['OPENAI_API_KEY'] = input('Please enter your OPENAI_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) ai = Genkit(plugins=[OpenAI()], model=openai_model('gpt-4o')) @@ -228,7 +227,7 @@ class RoundTripInput(BaseModel): @ai.tool(description='calculates a gablorken', name='gablorkenTool') -def gablorken_tool(input_: GablorkenInput) -> int: +async def gablorken_tool(input_: GablorkenInput) -> int: """Calculate a gablorken. Args: @@ -241,7 +240,7 @@ def gablorken_tool(input_: GablorkenInput) -> int: @ai.tool(description='Get current temperature for provided coordinates in celsius') -def get_weather_tool(coordinates: WeatherRequest) -> float: +async def get_weather_tool(coordinates: WeatherRequest) -> float: """Get the current temperature for provided coordinates in celsius. Args: @@ -269,7 +268,7 @@ async def calculate_gablorken(input: GablorkenFlowInput) -> str: input: Input with value for gablorken calculation. Returns: - A GenerateRequest object with the evaluation output + A ModelRequest object with the evaluation output """ response = await ai.generate( prompt=f'what is the gablorken of {input.value}', @@ -322,7 +321,7 @@ async def get_weather_flow(input: WeatherFlowInput) -> WeatherResponse: config={'model': 'gpt-4o-mini-2024-07-18', 'temperature': 1}, prompt=f"What's the weather like in {input.location} today?", tools=['get_weather_tool'], - output=Output(schema=WeatherResponse), + output_schema=WeatherResponse, ) return WeatherResponse.model_validate(response.output) @@ -337,7 +336,7 @@ async def get_weather_flow_stream(input: WeatherFlowInput) -> WeatherResponse: Returns: The weather for the location. """ - stream, response = ai.generate_stream( + stream_response = ai.generate_stream( model=openai_model('gpt-4o'), system=( 'You are an assistant that provides current weather information in JSON format and calculates ' @@ -346,11 +345,11 @@ async def get_weather_flow_stream(input: WeatherFlowInput) -> WeatherResponse: config={'model': 'gpt-4o-2024-08-06', 'temperature': 1}, prompt=f"What's the weather like in {input.location} today?", tools=['get_weather_tool', 'gablorkenTool'], - output=Output(schema=WeatherResponse), + output_schema=WeatherResponse, ) - async for _chunk in stream: + async for _chunk in stream_response.stream: pass - final = await response + final = await stream_response.response return WeatherResponse.model_validate(final.output) @@ -422,7 +421,7 @@ async def structured_menu_suggestion(input: MenuSuggestionInput) -> MenuSuggesti """ response = await ai.generate( prompt=f'Suggest a menu item for a {input.theme}-themed restaurant.', - output=Output(schema=MenuSuggestion), + output_schema=MenuSuggestion, ) return response.output diff --git a/py/samples/provider-firestore-retriever/LICENSE b/py/samples/provider-firestore-retriever/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/samples/provider-firestore-retriever/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/samples/provider-firestore-retriever/README.md b/py/samples/provider-firestore-retriever/README.md deleted file mode 100644 index e0b400486b..0000000000 --- a/py/samples/provider-firestore-retriever/README.md +++ /dev/null @@ -1,98 +0,0 @@ -# Firestore Vector Retriever - -Demonstrates using Firestore as a vector store for RAG applications. - -## Quick Start - -```bash -export GOOGLE_CLOUD_PROJECT=your-project-id -./run.sh -``` - -The script will: - -1. βœ“ Prompt for your project ID if not set -2. βœ“ Check gcloud authentication (and help you authenticate if needed) -3. βœ“ Enable Firestore and Vertex AI APIs (with your permission) -4. βœ“ Install dependencies -5. βœ“ Remind you to create the vector index (see below) -6. βœ“ Start the demo and open your browser - -## Required: Create Firestore Vector Index - -Before using vector search, you **must** create a composite index: - -```bash -gcloud firestore indexes composite create \ - --project=$GOOGLE_CLOUD_PROJECT \ - --collection-group=films \ - --query-scope=COLLECTION \ - --field-config=vector-config='{"dimension":"768","flat": {}}',field-path=embedding -``` - -> **Note:** Index creation may take a few minutes. The demo will show an error until the index is ready. - -## Manual Setup (if needed) - -If you prefer manual setup or the automatic setup fails: - -### 1. Install gcloud CLI - -Download from: https://cloud.google.com/sdk/docs/install - -### 2. Authentication - -```bash -gcloud auth application-default login -``` - -### 3. Enable Required APIs - -```bash -# Firestore API -gcloud services enable firestore.googleapis.com --project=$GOOGLE_CLOUD_PROJECT - -# Vertex AI API (for embeddings) -gcloud services enable aiplatform.googleapis.com --project=$GOOGLE_CLOUD_PROJECT -``` - -### 4. Create Vector Index - -See the command above in "Required: Create Firestore Vector Index" - -### 5. Run the Demo - -```bash -./run.sh -``` - -Or manually: - -```bash -genkit start -- uv run src/main.py -``` - -Then open the Dev UI at http://localhost:4000 - -## Testing the Demo - -1. Open DevUI at http://localhost:4000 -2. Run `index_documents` to populate the vector store -3. Run `retrieve_documents` with a query to test similarity search -4. Verify retrieved documents match query semantics - -## Expected Behavior - -- Documents are embedded using Vertex AI and stored in Firestore -- Vector similarity search returns semantically relevant documents -- Firebase telemetry captures traces and metrics - -## Development - -The `run.sh` script uses `watchmedo` to monitor changes in: -- `src/` (Python logic) -- `../../packages` (Genkit core) -- `../../plugins` (Genkit plugins) -- File patterns: `*.py`, `*.prompt`, `*.json` - -Changes will automatically trigger a restart of the sample. diff --git a/py/samples/provider-firestore-retriever/pyproject.toml b/py/samples/provider-firestore-retriever/pyproject.toml deleted file mode 100644 index 2f6d4a7852..0000000000 --- a/py/samples/provider-firestore-retriever/pyproject.toml +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Private :: Do Not Upload", -] -dependencies = [ - "rich>=13.0.0", - "genkit", - "genkit-plugin-firebase", - "genkit-plugin-google-genai", - "google-cloud-firestore", - "uvloop>=0.21.0", -] -description = "firestore-retriever Genkit sample" -license = "Apache-2.0" -name = "provider-firestore-retriever" -readme = "README.md" -requires-python = ">=3.10" -version = "0.2.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src/firestore-retriever"] diff --git a/py/samples/provider-firestore-retriever/run.sh b/py/samples/provider-firestore-retriever/run.sh deleted file mode 100755 index 7e850cda5c..0000000000 --- a/py/samples/provider-firestore-retriever/run.sh +++ /dev/null @@ -1,114 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -# Firestore Retriever Demo -# ======================== -# -# Demonstrates using Firestore as a vector store. -# -# This script automates most of the setup: -# - Detects/prompts for GOOGLE_CLOUD_PROJECT -# - Checks gcloud authentication -# - Enables required APIs -# - Installs dependencies -# -# Note: You still need to create a Firestore vector index manually. -# -# Usage: -# ./run.sh # Start the demo with Dev UI -# ./run.sh --setup # Run setup only (check auth, enable APIs) -# ./run.sh --help # Show this help message - -set -euo pipefail - -cd "$(dirname "$0")" -source "../_common.sh" - -# Required APIs for this demo -REQUIRED_APIS=( - "firestore.googleapis.com" # Firestore API - "aiplatform.googleapis.com" # Vertex AI API (for embeddings) -) - -print_help() { - print_banner "Firestore Retriever Demo" "πŸ”₯" - echo "Usage: ./run.sh [options]" - echo "" - echo "Options:" - echo " --help Show this help message" - echo " --setup Run setup only (auth check, enable APIs)" - echo "" - echo "The script will automatically:" - echo " 1. Prompt for GOOGLE_CLOUD_PROJECT if not set" - echo " 2. Check gcloud authentication" - echo " 3. Enable required APIs (with your permission)" - echo " 4. Install dependencies" - echo " 5. Start the demo and open the browser" - echo "" - echo "Required APIs (enabled automatically):" - echo " - Firestore API (firestore.googleapis.com)" - echo " - Vertex AI API (aiplatform.googleapis.com)" - echo "" - echo -e "${YELLOW}Note:${NC} You must create a Firestore vector index manually:" - echo "" - echo " gcloud firestore indexes composite create \\" - echo " --project=\$GOOGLE_CLOUD_PROJECT \\" - echo " --collection-group=films \\" - echo " --query-scope=COLLECTION \\" - echo " --field-config=vector-config='{\"dimension\":\"768\",\"flat\": {}}',field-path=embedding" - echo "" - print_help_footer -} - -# Print reminder about Firestore index -print_firestore_index_reminder() { - echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" - echo -e "${YELLOW}REMINDER: Create Firestore vector index if not already done:${NC}" - echo "" - echo " gcloud firestore indexes composite create \\" - echo " --project=${GOOGLE_CLOUD_PROJECT:-YOUR_PROJECT} \\" - echo " --collection-group=films \\" - echo " --query-scope=COLLECTION \\" - echo " --field-config=vector-config='{\"dimension\":\"768\",\"flat\": {}}',field-path=embedding" - echo "" - echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" - echo "" -} - -# Main -case "${1:-}" in - --help|-h) - print_help - exit 0 - ;; - --setup) - print_banner "Setup" "βš™οΈ" - run_gcp_setup "${REQUIRED_APIS[@]}" || exit 1 - print_firestore_index_reminder - echo -e "${GREEN}Setup complete!${NC}" - echo "" - exit 0 - ;; -esac - -print_banner "Firestore Retriever Demo" "πŸ”₯" - -# Run GCP setup (checks gcloud, auth, enables APIs) -run_gcp_setup "${REQUIRED_APIS[@]}" || exit 1 - -# Remind about Firestore index -print_firestore_index_reminder - -# Install dependencies -install_deps - -# Start the demo -genkit_start_with_browser -- \ - uv tool run --from watchdog watchmedo auto-restart \ - -d src \ - -d ../../packages \ - -d ../../plugins \ - -p '*.py;*.prompt;*.json' \ - -R \ - -- uv run src/main.py "$@" diff --git a/py/samples/provider-firestore-retriever/src/main.py b/py/samples/provider-firestore-retriever/src/main.py deleted file mode 100644 index 978a910963..0000000000 --- a/py/samples/provider-firestore-retriever/src/main.py +++ /dev/null @@ -1,184 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Firestore retriever sample - Vector search with Firestore. - -This sample demonstrates how to use Firestore as a vector store for -retrieval-augmented generation (RAG) with Genkit. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Firestore β”‚ Google's NoSQL database. Stores documents like β”‚ - β”‚ β”‚ JSON files that sync across devices. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Vector Search β”‚ Finding similar items by meaning, not keywords. β”‚ - β”‚ β”‚ "Happy" finds docs about "joyful" too. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ RAG β”‚ Retrieval-Augmented Generation. AI looks up β”‚ - β”‚ β”‚ your docs before answering. More accurate! β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Embedding β”‚ Numbers that capture meaning. Similar text gets β”‚ - β”‚ β”‚ similar numbers (close in vector space). β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Retriever β”‚ The component that finds matching documents. β”‚ - β”‚ β”‚ "Find docs about sci-fi" returns relevant results. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow (RAG with Firestore):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW FIRESTORE VECTOR SEARCH FINDS YOUR DOCUMENTS β”‚ - β”‚ β”‚ - β”‚ INDEXING (Setup) β”‚ - β”‚ ───────────────── β”‚ - β”‚ Documents β†’ Embedder β†’ Firestore (stored with vectors) β”‚ - β”‚ β”‚ - β”‚ RETRIEVAL (Query Time) β”‚ - β”‚ ─────────────────────── β”‚ - β”‚ Query: "sci-fi films" β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) Convert query to embedding β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ "sci-fi films" β†’ [0.2, -0.5, ...] β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Search Firestore for similar vectors β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Firestore β”‚ Vector similarity search β”‚ - β”‚ β”‚ (Native Index) β”‚ Returns closest matches β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Return matching documents β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Results β”‚ "The Matrix", "Inception", etc. β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Features -============ -| Feature Description | Example Function / Code Snippet | -|-----------------------------------------|-------------------------------------| -| Firestore Vector Store Definition | `define_firestore_vector_store` | -| Embed Many | `ai.embed_many()` | -| Document Retrieval | `ai.retrieve()` | -| Firestore Integration | `firestore.Client()` | - -See README.md for testing instructions. -""" - -import os - -from google.cloud import firestore -from google.cloud.firestore_v1.base_vector_query import DistanceMeasure -from google.cloud.firestore_v1.vector import Vector - -from genkit.ai import Genkit -from genkit.plugins.firebase import add_firebase_telemetry, define_firestore_vector_store -from genkit.plugins.google_genai import VertexAI -from genkit.types import Document, RetrieverResponse -from samples.shared.logging import setup_sample - -setup_sample() - -if 'GCLOUD_PROJECT' not in os.environ: - os.environ['GCLOUD_PROJECT'] = input('Please enter your GCLOUD_PROJECT: ') - -# Important: use the same embedding model for indexing and retrieval. -EMBEDDING_MODEL = 'vertexai/gemini-embedding-001' - -# Add Firebase telemetry (metrics, logs, traces) -add_firebase_telemetry() - -firestore_client = firestore.Client() - -# Create Genkit instance -ai = Genkit(plugins=[VertexAI()]) - -# Define Firestore vector store - returns the retriever name -RETRIEVER_NAME = define_firestore_vector_store( - ai, - name='my_firestore_retriever', - embedder=EMBEDDING_MODEL, - collection='films', - vector_field='embedding', - content_field='text', - firestore_client=firestore_client, - distance_measure=DistanceMeasure.EUCLIDEAN, -) - -collection_name = 'films' - -films = [ - 'The Godfather is a 1972 crime film directed by Francis Ford Coppola.', - 'The Dark Knight is a 2008 superhero film directed by Christopher Nolan.', - 'Pulp Fiction is a 1994 crime film directed by Quentin Tarantino.', - "Schindler's List is a 1993 historical drama directed by Steven Spielberg.", - 'Inception is a 2010 sci-fi film directed by Christopher Nolan.', - 'The Matrix is a 1999 sci-fi film directed by the Wachowskis.', - 'Fight Club is a 1999 film directed by David Fincher.', - 'Forrest Gump is a 1994 drama directed by Robert Zemeckis.', - 'Star Wars is a 1977 sci-fi film directed by George Lucas.', - 'The Shawshank Redemption is a 1994 drama directed by Frank Darabont.', -] - - -@ai.flow() -async def index_documents() -> None: - """Indexes the film documents in Firestore.""" - embeddings = await ai.embed_many(embedder=EMBEDDING_MODEL, content=films) - for i, film_text in enumerate(films): - doc_id = f'doc-{i + 1}' - embedding = embeddings[i].embedding - - doc_ref = firestore_client.collection(collection_name).document(doc_id) - try: - doc_ref.set({ - 'text': film_text, - 'embedding': Vector(embedding), - 'metadata': f'metadata for doc {i + 1}', - }) - except Exception: - return - - -@ai.flow() -async def retrieve_documents() -> RetrieverResponse: - """Retrieves the film documents from Firestore.""" - return await ai.retrieve( - query=Document.from_text('sci-fi film'), - retriever=RETRIEVER_NAME, - options={'limit': 10}, - ) - - -async def main() -> None: - """Main entry point for the flow sample. - - This function demonstrates how to create and use AI flows in the - Genkit framework. - """ - await index_documents() - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/samples/provider-google-genai-code-execution/src/main.py b/py/samples/provider-google-genai-code-execution/src/main.py index 452a0b1be7..dea3cf8676 100755 --- a/py/samples/provider-google-genai-code-execution/src/main.py +++ b/py/samples/provider-google-genai-code-execution/src/main.py @@ -54,12 +54,11 @@ import os +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.blocks.model import MessageWrapper -from genkit.core.logging import get_logger -from genkit.core.typing import CustomPart, Message, TextPart +from genkit import Genkit, Message +from genkit._core._typing import CustomPart, TextPart from genkit.plugins.google_genai import GeminiConfigSchema, GoogleAI from genkit.plugins.google_genai.models.utils import PartConverter from samples.shared.logging import setup_sample @@ -69,7 +68,7 @@ if 'GEMINI_API_KEY' not in os.environ: os.environ['GEMINI_API_KEY'] = input('Please enter your GEMINI_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) ai = Genkit( plugins=[GoogleAI()], @@ -87,7 +86,7 @@ class CodeExecutionInput(BaseModel): @ai.flow() -async def execute_code(input: CodeExecutionInput) -> MessageWrapper: +async def execute_code(input: CodeExecutionInput) -> Message: """Execute code for the given task. Args: diff --git a/py/samples/provider-google-genai-context-caching/src/main.py b/py/samples/provider-google-genai-context-caching/src/main.py index a5659c9422..ee3d1a4ebb 100755 --- a/py/samples/provider-google-genai-context-caching/src/main.py +++ b/py/samples/provider-google-genai-context-caching/src/main.py @@ -98,12 +98,11 @@ import pathlib import httpx +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.core.logging import get_logger +from genkit import Genkit, Message, ModelConfig, Part, Role, TextPart from genkit.plugins.google_genai import GoogleAI -from genkit.types import GenerationCommonConfig, Message, Part, Role, TextPart from samples.shared.logging import setup_sample setup_sample() @@ -111,7 +110,7 @@ if 'GEMINI_API_KEY' not in os.environ: os.environ['GEMINI_API_KEY'] = input('Please enter your GEMINI_API_KEY: ') -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) ai = Genkit( @@ -172,7 +171,7 @@ async def text_context_flow(_input: BookContextInputSchema) -> str: }, ), ], - config=GenerationCommonConfig( + config=ModelConfig( temperature=0.7, max_output_tokens=1000, top_k=50, @@ -194,7 +193,7 @@ async def text_context_flow(_input: BookContextInputSchema) -> str: '* **Key Concept:** Description...\n' 'Keep it concise, use pirate slang, but maintain the helpful advice.' ), - config=GenerationCommonConfig( + config=ModelConfig( version='gemini-3-flash-preview', temperature=0.7, max_output_tokens=1000, diff --git a/py/samples/provider-google-genai-hello/src/main.py b/py/samples/provider-google-genai-hello/src/main.py index 9710edb085..e3140824fe 100755 --- a/py/samples/provider-google-genai-hello/src/main.py +++ b/py/samples/provider-google-genai-hello/src/main.py @@ -129,24 +129,26 @@ import pathlib +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit, Output, ToolRunContext, tool_response -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger -from genkit.plugins.google_cloud import add_gcp_telemetry -from genkit.plugins.google_genai import ( - EmbeddingTaskType, - GoogleAI, -) -from genkit.types import ( - GenerationCommonConfig, +from genkit import ( + Genkit, Media, MediaPart, Message, + ModelConfig, Part, Role, TextPart, + ToolRunContext, + tool_response, +) +from genkit._core._action import ActionRunContext +from genkit.plugins.google_cloud import add_gcp_telemetry +from genkit.plugins.google_genai import ( + EmbeddingTaskType, + GoogleAI, ) from samples.shared import ( CharacterInput, @@ -176,7 +178,7 @@ setup_sample() -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) if 'GEMINI_API_KEY' not in os.environ: @@ -245,7 +247,7 @@ class ScreenshotInput(BaseModel): @ai.tool(name='gablorkenTool') -def gablorken_tool(input_: GablorkenInput) -> dict[str, int]: +async def gablorken_tool(input_: GablorkenInput) -> dict[str, int]: """Calculate a gablorken. Returns: @@ -255,19 +257,19 @@ def gablorken_tool(input_: GablorkenInput) -> dict[str, int]: @ai.tool(name='gablorkenTool2') -def gablorken_tool2(_input: GablorkenInput, ctx: ToolRunContext) -> None: +async def gablorken_tool2(_input: GablorkenInput, ctx: ToolRunContext) -> None: """The user-defined tool function.""" pass @ai.tool(name='screenShot') -def take_screenshot(input_: ScreenshotInput) -> dict: +async def take_screenshot(input_: ScreenshotInput) -> dict: """Take a screenshot of a given URL.""" return {'url': input_.url, 'screenshot_path': '/tmp/screenshot.png'} # noqa: S108 - sample code @ai.tool(name='getWeather') -def get_weather_detailed(input_: WeatherInput) -> dict: +async def get_weather_detailed(input_: WeatherInput) -> dict: """Used to get current weather for a location.""" return { 'location': input_.location, @@ -283,13 +285,13 @@ class CelsiusInput(BaseModel): @ai.tool(name='celsiusToFahrenheit') -def celsius_to_fahrenheit(input_: CelsiusInput) -> float: +async def celsius_to_fahrenheit(input_: CelsiusInput) -> float: """Converts Celsius to Fahrenheit.""" return (input_.celsius * 9) / 5 + 32 @ai.tool() -def get_user_data() -> str: +async def get_user_data() -> str: """Fetch user data based on context.""" context = Genkit.current_context() raw_user = context.get('user') if context else {} @@ -428,9 +430,8 @@ async def generate_character_instructions( """ result = await ai.generate( prompt=f'generate an RPG character named {input.name}', - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, output_constrained=False, - output_instructions=True, ) return result.output @@ -689,7 +690,7 @@ async def tool_calling(input: ToolCallingInput) -> str: response = await ai.generate( tools=['getWeather', 'celsiusToFahrenheit'], prompt=f"What's the weather in {input.location}? Convert the temperature to Fahrenheit.", - config=GenerationCommonConfig(temperature=1), + config=ModelConfig(temperature=1), ) return response.text @@ -701,7 +702,7 @@ async def streaming_structured_output( ) -> RpgCharacter: """Demonstrate streaming with structured output schemas. - Combines `generate_stream` with `Output(schema=...)` so the model + Combines `generate_stream` with `output_schema=...` so the model streams JSON tokens that are progressively parsed into the Pydantic model. Each chunk exposes a partial `.output` you can forward to clients for incremental rendering. @@ -715,19 +716,19 @@ async def streaming_structured_output( Returns: The fully-parsed RPG character once streaming completes. """ - stream, result = ai.generate_stream( + stream_response = ai.generate_stream( prompt=( f'Generate an RPG character named {input.name}. ' 'Include a creative backstory, 3-4 unique abilities, ' 'and skill ratings for strength, charisma, and endurance (0-100 each).' ), - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, ) - async for chunk in stream: + async for chunk in stream_response.stream: if ctx is not None: ctx.send_chunk(chunk.output) - return (await result).output + return (await stream_response.response).output @ai.flow() diff --git a/py/samples/provider-google-genai-hello/src/main_vertexai.py b/py/samples/provider-google-genai-hello/src/main_vertexai.py index 8fa7d0320c..70d480f755 100644 --- a/py/samples/provider-google-genai-hello/src/main_vertexai.py +++ b/py/samples/provider-google-genai-hello/src/main_vertexai.py @@ -21,9 +21,8 @@ import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit +from genkit import Genkit, Media, MediaPart, Metadata, ModelConfig, Part, TextPart from genkit.plugins.google_genai import GeminiImageConfigSchema, VertexAI -from genkit.types import GenerationCommonConfig, Media, MediaPart, Metadata, Part, TextPart logger = structlog.get_logger(__name__) @@ -239,7 +238,7 @@ async def imagen_image_generation() -> Media | None: @ai.tool(name='getWeather') -def get_weather(location: str) -> dict: +async def get_weather(location: str) -> dict: """Used to get current weather for a location.""" return { 'location': location, @@ -255,7 +254,7 @@ class CelsiusInput(BaseModel): @ai.tool(name='celsiusToFahrenheit') -def celsius_to_fahrenheit(input_: CelsiusInput) -> float: +async def celsius_to_fahrenheit(input_: CelsiusInput) -> float: """Converts Celsius to Fahrenheit.""" return (input_.celsius * 9) / 5 + 32 @@ -267,7 +266,7 @@ async def tool_calling(location: str = 'Paris, France') -> str: model='vertexai/gemini-2.5-flash', tools=['getWeather', 'celsiusToFahrenheit'], prompt=f"What's the weather in {location}? Convert the temperature to Fahrenheit.", - config=GenerationCommonConfig(temperature=1), + config=ModelConfig(temperature=1), ) return response.text diff --git a/py/samples/provider-google-genai-media-models-demo/src/main.py b/py/samples/provider-google-genai-media-models-demo/src/main.py index 5d723eb414..8199de7510 100644 --- a/py/samples/provider-google-genai-media-models-demo/src/main.py +++ b/py/samples/provider-google-genai-media-models-demo/src/main.py @@ -103,16 +103,19 @@ from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.blocks.background_model import lookup_background_action -from genkit.blocks.model import GenerateResponseWrapper -from genkit.core.action import ActionRunContext -from genkit.core.typing import ( +from genkit import ( + Genkit, + Media, + MediaPart, + Metadata, + ModelConfig, + ModelResponse, +) +from genkit._core._action import ActionRunContext +from genkit._core._background import lookup_background_action +from genkit._core._typing import ( Error, FinishReason, - GenerateRequest, - GenerateResponse, - Message, ModelInfo, Operation, Part, @@ -120,12 +123,7 @@ Supports, TextPart, ) -from genkit.types import ( - GenerationCommonConfig, - Media, - MediaPart, - Metadata, -) +from genkit.model import Message, ModelRequest from samples.shared.logging import setup_sample setup_sample() @@ -329,7 +327,7 @@ class SimulatedLyriaConfig(BaseModel): sample_count: int = Field(default=1, description='Number of audio samples') -def _extract_prompt(request: GenerateRequest) -> str: +def _extract_prompt(request: ModelRequest) -> str: """Extract text prompt from request.""" if request.messages: for msg in request.messages: @@ -341,9 +339,9 @@ def _extract_prompt(request: GenerateRequest) -> str: # --- Simulated TTS --- async def simulated_tts_generate( - request: GenerateRequest, + request: ModelRequest, ctx: ActionRunContext, -) -> GenerateResponse: +) -> ModelResponse: """Simulate TTS audio generation.""" _extract_prompt(request) @@ -353,7 +351,7 @@ async def simulated_tts_generate( # Real TTS would return actual audio fake_audio = base64.b64encode(b'RIFF' + b'\x00' * 100).decode() - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=[ @@ -371,9 +369,9 @@ async def simulated_tts_generate( # --- Simulated Image --- async def simulated_image_generate( - request: GenerateRequest, + request: ModelRequest, ctx: ActionRunContext, -) -> GenerateResponse: +) -> ModelResponse: """Simulate image generation.""" _extract_prompt(request) @@ -384,7 +382,7 @@ async def simulated_image_generate( b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde' ).decode() - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=[ @@ -402,9 +400,9 @@ async def simulated_image_generate( # --- Simulated Lyria --- async def simulated_lyria_generate( - request: GenerateRequest, + request: ModelRequest, ctx: ActionRunContext, -) -> GenerateResponse: +) -> ModelResponse: """Simulate audio generation.""" _extract_prompt(request) @@ -412,7 +410,7 @@ async def simulated_lyria_generate( fake_audio = base64.b64encode(b'RIFF' + b'\x00' * 200).decode() - return GenerateResponse( + return ModelResponse( message=Message( role=Role.MODEL, content=[ @@ -430,7 +428,7 @@ async def simulated_lyria_generate( # --- Simulated Veo (Background Model) --- async def simulated_veo_start( - request: GenerateRequest, + request: ModelRequest, ctx: ActionRunContext, ) -> Operation: """Start simulated video generation.""" @@ -916,7 +914,7 @@ async def veo_video_generator_flow(input: VideoInput | None = None) -> dict[str, # Start the operation operation = await video_model.start( - GenerateRequest( + ModelRequest( messages=[Message(role=Role.USER, content=[Part(root=TextPart(text=prompt))])], config=config, ) @@ -1017,7 +1015,7 @@ async def media_models_overview_flow() -> dict[str, Any]: @ai.tool(name='screenshot') -def screenshot() -> dict: +async def screenshot() -> dict: """Takes a screenshot of a room.""" room_path = pathlib.Path(__file__).parent.parent / 'my_room.png' with pathlib.Path(room_path).open('rb') as f: @@ -1075,7 +1073,7 @@ async def describe_image_with_gemini(input: DescribeImageInput | None = None) -> async def generate_images( input: GenerateImagesInput | None = None, ctx: ActionRunContext | None = None, -) -> GenerateResponseWrapper: +) -> ModelResponse: """Generate images for the given subject using multimodal prompting. Args: @@ -1115,7 +1113,7 @@ async def multipart_tool_calling(input: ToolCallingInput | None = None) -> str: response = await ai.generate( model='googleai/gemini-3-pro-preview', tools=['screenshot'], - config=GenerationCommonConfig(temperature=1), + config=ModelConfig(temperature=1), prompt=input.prompt, ) return response.text diff --git a/py/samples/provider-google-genai-vertexai-hello/src/main.py b/py/samples/provider-google-genai-vertexai-hello/src/main.py index 981260ea36..acd96f75d4 100755 --- a/py/samples/provider-google-genai-vertexai-hello/src/main.py +++ b/py/samples/provider-google-genai-vertexai-hello/src/main.py @@ -107,16 +107,15 @@ import os +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit, Output, ToolRunContext, tool_response -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger +from genkit import Embedding, Genkit, ToolRunContext, tool_response +from genkit._core._action import ActionRunContext from genkit.plugins.google_genai import ( EmbeddingTaskType, VertexAI, ) -from genkit.types import Embedding from samples.shared import ( CharacterInput, CodeInput, @@ -143,7 +142,7 @@ setup_sample() -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) # Check for GCLOUD_PROJECT or GOOGLE_CLOUD_PROJECT # If GOOGLE_CLOUD_PROJECT is set but GCLOUD_PROJECT isn't, use it @@ -183,7 +182,7 @@ class ToolsFlowInput(BaseModel): @ai.tool(name='gablorkenTool') -def gablorken_tool(input_: GablorkenInput) -> int: +async def gablorken_tool(input_: GablorkenInput) -> int: """Calculate a gablorken. Args: @@ -196,7 +195,7 @@ def gablorken_tool(input_: GablorkenInput) -> int: @ai.tool(name='gablorkenTool2') -def gablorken_tool2(input_: GablorkenInput, ctx: ToolRunContext) -> None: +async def gablorken_tool2(input_: GablorkenInput, ctx: ToolRunContext) -> None: """The user-defined tool function. Args: @@ -282,9 +281,8 @@ async def generate_character_instructions( """ result = await ai.generate( prompt=f'generate an RPG character named {input.name}', - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, output_constrained=False, - output_instructions=True, ) return result.output @@ -415,7 +413,7 @@ async def streaming_structured_output( ) -> RpgCharacter: """Demonstrate streaming with structured output schemas. - Combines `generate_stream` with `Output(schema=...)` so the model + Combines `generate_stream` with `output_schema=...` so the model streams JSON tokens that are progressively parsed into the Pydantic model. Each chunk exposes a partial `.output` you can forward to clients for incremental rendering. @@ -429,19 +427,19 @@ async def streaming_structured_output( Returns: The fully-parsed RPG character once streaming completes. """ - stream, result = ai.generate_stream( + stream_response = ai.generate_stream( prompt=( f'Generate an RPG character named {input.name}. ' 'Include a creative backstory, 3-4 unique abilities, ' 'and skill ratings for strength, charisma, and endurance (0-100 each).' ), - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, ) - async for chunk in stream: + async for chunk in stream_response.stream: if ctx is not None: ctx.send_chunk(chunk.output) - return (await result).output + return (await stream_response.response).output @ai.flow() diff --git a/py/samples/provider-google-genai-vertexai-image/src/main.py b/py/samples/provider-google-genai-vertexai-image/src/main.py index 6c37082245..f32aacb2f0 100755 --- a/py/samples/provider-google-genai-vertexai-image/src/main.py +++ b/py/samples/provider-google-genai-vertexai-image/src/main.py @@ -54,8 +54,7 @@ from PIL import Image -from genkit.ai import Genkit -from genkit.blocks.model import GenerateResponseWrapper +from genkit import Genkit, ModelResponse from genkit.plugins.google_genai import VertexAI from samples.shared.logging import setup_sample @@ -73,7 +72,7 @@ @ai.flow() -async def draw_image_with_imagen() -> GenerateResponseWrapper: +async def draw_image_with_imagen() -> ModelResponse: """Draw an image using Imagen model. Returns: diff --git a/py/samples/provider-ollama-hello/src/main.py b/py/samples/provider-ollama-hello/src/main.py index 19ca170305..b4dfde99f6 100755 --- a/py/samples/provider-ollama-hello/src/main.py +++ b/py/samples/provider-ollama-hello/src/main.py @@ -81,13 +81,12 @@ from math import sqrt +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit, Output -from genkit.blocks.model import GenerateResponseWrapper -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger -from genkit.core.typing import Media, MediaPart, Part, TextPart +from genkit import Genkit, ModelResponse +from genkit._core._action import ActionRunContext +from genkit._core._typing import Media, MediaPart, Part, TextPart from genkit.plugins.ollama import Ollama, ollama_name from genkit.plugins.ollama.embedders import EmbeddingDefinition from genkit.plugins.ollama.models import ModelDefinition @@ -120,7 +119,7 @@ setup_sample() -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) # Pull models with: ollama pull GEMMA_MODEL = 'gemma3:latest' @@ -223,7 +222,7 @@ class PokemonFlowInput(BaseModel): @ai.tool() -def gablorken_tool(input: GablorkenInput) -> int: +async def gablorken_tool(input: GablorkenInput) -> int: """Calculate a gablorken.""" return input.value * 3 - 5 @@ -296,7 +295,7 @@ async def structured_menu_suggestion(input: MenuSuggestionInput) -> MenuSuggesti """ response = await ai.generate( prompt=f'Suggest a menu item for a {input.theme}-themed restaurant.', - output=Output(schema=MenuSuggestion), + output_schema=MenuSuggestion, ) return response.output @@ -564,7 +563,7 @@ def find_nearest_pokemons(input_embedding: list[float], top_n: int = 3) -> list[ return [pokemon for _distance, pokemon in pokemon_distances[:top_n]] -async def generate_rag_response(question: str) -> GenerateResponseWrapper: +async def generate_rag_response(question: str) -> ModelResponse: """Generate a RAG response: embed the question, find context, generate. Args: diff --git a/py/samples/provider-vertex-ai-model-garden/src/main.py b/py/samples/provider-vertex-ai-model-garden/src/main.py index d320e8228e..6bcf1410e8 100644 --- a/py/samples/provider-vertex-ai-model-garden/src/main.py +++ b/py/samples/provider-vertex-ai-model-garden/src/main.py @@ -62,19 +62,18 @@ import asyncio import os +import structlog from pydantic import BaseModel, Field -from genkit.ai import Genkit, Output -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger +from genkit import Genkit, Message, Part, Role, TextPart +from genkit._core._action import ActionRunContext from genkit.plugins.google_genai import VertexAI from genkit.plugins.vertex_ai.model_garden import ModelGardenPlugin, model_garden_name -from genkit.types import Message, Part, Role, TextPart from samples.shared.logging import setup_sample setup_sample() -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) class RpgCharacter(BaseModel): @@ -195,7 +194,7 @@ class ToolFlowInput(BaseModel): @ai.tool() -def get_weather(input: WeatherInput) -> dict: +async def get_weather(input: WeatherInput) -> dict: """Used to get current weather for a location.""" return { 'location': input.location, @@ -205,13 +204,13 @@ def get_weather(input: WeatherInput) -> dict: @ai.tool() -def celsius_to_fahrenheit(input: TemperatureInput) -> float: +async def celsius_to_fahrenheit(input: TemperatureInput) -> float: """Converts Celsius to Fahrenheit.""" return (input.celsius * 9) / 5 + 32 @ai.tool(name='getWeather') -def get_weather_tool(input_: WeatherInput) -> str: +async def get_weather_tool(input_: WeatherInput) -> str: """Used to get current weather for a location.""" return f'Weather in {input_.location}: Sunny, 21.5Β°C' @@ -279,7 +278,7 @@ async def generate_character(input: CharacterInput) -> RpgCharacter: result = await ai.generate( model=model_garden_name('anthropic/claude-3-5-sonnet-v2@20241022'), prompt=f'generate an RPG character named {input.name}', - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, ) return result.output @@ -336,13 +335,13 @@ async def say_hi_stream( Returns: The response from the model. """ - stream, _ = ai.generate_stream( + stream_response = ai.generate_stream( model=model_garden_name('anthropic/claude-3-5-sonnet-v2@20241022'), config={'temperature': 1}, prompt=f'hi {input.name}', ) result = '' - async for data in stream: + async for data in stream_response.stream: if ctx is not None: ctx.send_chunk(data.text) result += data.text @@ -437,7 +436,7 @@ async def streaming_structured_output( ) -> RpgCharacter: """Demonstrate streaming with structured output schemas. - Combines `generate_stream` with `Output(schema=...)` so the model + Combines `generate_stream` with `output_schema=...` so the model streams JSON tokens that are progressively parsed into the Pydantic model. Each chunk exposes a partial `.output` you can forward to clients for incremental rendering. @@ -451,20 +450,20 @@ async def streaming_structured_output( Returns: The fully-parsed RPG character once streaming completes. """ - stream, result = ai.generate_stream( + stream_response = ai.generate_stream( model=model_garden_name('anthropic/claude-3-5-sonnet-v2@20241022'), prompt=( f'Generate an RPG character named {input.name}. ' 'Include a creative backstory, 3-4 unique abilities, ' 'and skill ratings for strength, charisma, and endurance (0-100 each).' ), - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, ) - async for chunk in stream: + async for chunk in stream_response.stream: if ctx is not None: ctx.send_chunk(chunk.output) - return (await result).output + return (await stream_response.response).output async def main() -> None: diff --git a/py/samples/provider-vertex-ai-rerank-eval/LICENSE b/py/samples/provider-vertex-ai-rerank-eval/LICENSE deleted file mode 100644 index 996f16986d..0000000000 --- a/py/samples/provider-vertex-ai-rerank-eval/LICENSE +++ /dev/null @@ -1,207 +0,0 @@ -``` - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ -``` - -TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - -1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to the Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - -2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - -3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - -4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - -5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - -6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - -7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - -8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - -9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - -END OF TERMS AND CONDITIONS - -APPENDIX: How to apply the Apache License to your work. - -``` - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. -``` - -Copyright \[yyyy] \[name of copyright owner] - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - -``` - http://www.apache.org/licenses/LICENSE-2.0 -``` - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. diff --git a/py/samples/provider-vertex-ai-rerank-eval/README.md b/py/samples/provider-vertex-ai-rerank-eval/README.md deleted file mode 100644 index 5c8720d53e..0000000000 --- a/py/samples/provider-vertex-ai-rerank-eval/README.md +++ /dev/null @@ -1,161 +0,0 @@ -# Vertex AI Rerankers and Evaluators Demo - -Demonstrates using Vertex AI rerankers for RAG quality improvement and evaluators -for assessing model outputs. - -## Features - -### Rerankers - -Semantic document reranking improves RAG quality by re-ordering retrieved documents -based on their semantic relevance to a query. - -* **`rerank_documents`** - Basic document reranking -* **`rag_with_reranking`** - Full RAG pipeline (retrieve β†’ rerank β†’ generate) - -### Evaluators - -Vertex AI evaluators assess model outputs using various quality metrics: - -* **`evaluate_fluency`** - Text fluency (1-5 scale) -* **`evaluate_safety`** - Content safety assessment -* **`evaluate_groundedness`** - Hallucination detection (is output grounded in context?) -* **`evaluate_bleu`** - BLEU score for translation quality -* **`evaluate_summarization`** - Summarization quality assessment - -## Quick Start - -```bash -export GOOGLE_CLOUD_PROJECT=your-project-id -./run.sh -``` - -That's it! The script will: - -1. βœ“ Prompt for your project ID if not set -2. βœ“ Check gcloud authentication (and help you authenticate if needed) -3. βœ“ Enable required APIs (with your permission) -4. βœ“ Install dependencies -5. βœ“ Start the demo and open your browser - -## Manual Setup (if needed) - -If you prefer manual setup or the automatic setup fails: - -### 1. Authentication - -```bash -gcloud auth application-default login -``` - -### 2. Enable Required APIs - -```bash -# Vertex AI API (for models and evaluators) -gcloud services enable aiplatform.googleapis.com - -# Discovery Engine API (for rerankers) -gcloud services enable discoveryengine.googleapis.com -``` - -### 3. Run the Demo - -```bash -./run.sh -``` - -Or manually: - -```bash -genkit start -- uv run src/main.py -``` - -Then open the Dev UI at http://localhost:4000 - -## Testing the Demo - -### Reranker Flows - -1. **`rerank_documents`** - * Input: A query string (default: "How do neural networks learn?") - * Output: Documents sorted by relevance score - * The sample includes irrelevant documents to show how reranking filters them - -2. **`rag_with_reranking`** - * Input: A question (default: "What is machine learning?") - * Output: Generated answer using top-ranked documents as context - * Demonstrates the two-stage retrieval pattern - -### Evaluator Flows - -1. **`evaluate_fluency`** - * Tests text fluency with samples including intentionally poor grammar - * Scores: 1 (poor) to 5 (excellent) - -2. **`evaluate_safety`** - * Tests content safety - * Higher scores = safer content - -3. **`evaluate_groundedness`** - * Tests if outputs are grounded in provided context - * Includes a hallucination example (claims population when not in context) - -4. **`evaluate_bleu`** - * Tests translation quality against reference translations - * Scores: 0 to 1 (higher = closer to reference) - -5. **`evaluate_summarization`** - * Tests summarization quality - -## Supported Reranker Models - -| Model | Description | -|-------|-------------| -| `semantic-ranker-default@latest` | Latest default semantic ranker | -| `semantic-ranker-default-004` | Semantic ranker version 004 | -| `semantic-ranker-fast-004` | Fast variant (lower latency) | -| `semantic-ranker-default-003` | Semantic ranker version 003 | -| `semantic-ranker-default-002` | Semantic ranker version 002 | - -## Supported Evaluation Metrics - -| Metric | Description | -|--------|-------------| -| BLEU | Translation quality (compare to reference) | -| ROUGE | Summarization quality (compare to reference) | -| FLUENCY | Language mastery and readability | -| SAFETY | Harmful/inappropriate content check | -| GROUNDEDNESS | Factual grounding in context | -| SUMMARIZATION\_QUALITY | Overall summarization ability | -| SUMMARIZATION\_HELPFULNESS | Usefulness as a summary | -| SUMMARIZATION\_VERBOSITY | Conciseness of summary | - -## Troubleshooting - -### "Discovery Engine API not enabled" - -The script should enable this automatically, but if it fails: - -```bash -gcloud services enable discoveryengine.googleapis.com -``` - -### "Permission denied" - -Ensure your account has the required IAM roles: - -* `roles/discoveryengine.admin` (for rerankers) -* `roles/aiplatform.user` (for evaluators) - -### "Project not found" - -Verify `GOOGLE_CLOUD_PROJECT` is set correctly: - -```bash -echo $GOOGLE_CLOUD_PROJECT -``` - -### "gcloud not found" - -Install the Google Cloud SDK from: -https://cloud.google.com/sdk/docs/install diff --git a/py/samples/provider-vertex-ai-rerank-eval/pyproject.toml b/py/samples/provider-vertex-ai-rerank-eval/pyproject.toml deleted file mode 100644 index 29a0a0a0bf..0000000000 --- a/py/samples/provider-vertex-ai-rerank-eval/pyproject.toml +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Private :: Do Not Upload", -] -dependencies = [ - "rich>=13.0.0", - "genkit", - "genkit-plugin-google-genai", - "pydantic>=2.10.5", - "structlog>=25.2.0", - "uvloop>=0.21.0", -] -description = "Vertex AI Rerankers and Evaluators Demo" -license = "Apache-2.0" -name = "provider-vertex-ai-rerank-eval" -readme = "README.md" -requires-python = ">=3.10" -version = "0.2.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src/vertexai_rerank_eval"] diff --git a/py/samples/provider-vertex-ai-rerank-eval/run.sh b/py/samples/provider-vertex-ai-rerank-eval/run.sh deleted file mode 100755 index 8bec99c5b4..0000000000 --- a/py/samples/provider-vertex-ai-rerank-eval/run.sh +++ /dev/null @@ -1,85 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2025 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -# Vertex AI Rerankers and Evaluators Demo -# ======================================= -# -# Demonstrates using Vertex AI rerankers for RAG quality improvement -# and evaluators for assessing model outputs. -# -# This script automates most of the setup: -# - Detects/prompts for GOOGLE_CLOUD_PROJECT -# - Checks gcloud authentication -# - Enables required APIs -# - Installs dependencies -# -# Usage: -# ./run.sh # Start the demo with Dev UI -# ./run.sh --setup # Run setup only (check auth, enable APIs) -# ./run.sh --help # Show this help message - -set -euo pipefail - -cd "$(dirname "$0")" -source "../_common.sh" - -# Required APIs for this demo -REQUIRED_APIS=( - "aiplatform.googleapis.com" # Vertex AI API (models and evaluators) - "discoveryengine.googleapis.com" # Discovery Engine API (rerankers) -) - -print_help() { - print_banner "Vertex AI Rerankers & Evaluators" "πŸ”" - echo "Usage: ./run.sh [options]" - echo "" - echo "Options:" - echo " --help Show this help message" - echo " --setup Run setup only (auth check, enable APIs)" - echo "" - echo "The script will automatically:" - echo " 1. Prompt for GOOGLE_CLOUD_PROJECT if not set" - echo " 2. Check gcloud authentication" - echo " 3. Enable required APIs (with your permission)" - echo " 4. Install dependencies" - echo " 5. Start the demo and open the browser" - echo "" - echo "Required APIs (enabled automatically):" - echo " - Vertex AI API (aiplatform.googleapis.com)" - echo " - Discovery Engine API (discoveryengine.googleapis.com)" - print_help_footer -} - -# Main -case "${1:-}" in - --help|-h) - print_help - exit 0 - ;; - --setup) - print_banner "Setup" "βš™οΈ" - run_gcp_setup "${REQUIRED_APIS[@]}" || exit 1 - echo -e "${GREEN}Setup complete!${NC}" - echo "" - exit 0 - ;; -esac - -print_banner "Vertex AI Rerankers & Evaluators" "πŸ”" - -# Run GCP setup (checks gcloud, auth, enables APIs) -run_gcp_setup "${REQUIRED_APIS[@]}" || exit 1 - -# Install dependencies -install_deps - -# Start the demo -genkit_start_with_browser -- \ - uv tool run --from watchdog watchmedo auto-restart \ - -d src \ - -d ../../packages \ - -d ../../plugins \ - -p '*.py;*.prompt;*.json' \ - -R \ - -- uv run src/main.py "$@" diff --git a/py/samples/provider-vertex-ai-rerank-eval/src/main.py b/py/samples/provider-vertex-ai-rerank-eval/src/main.py deleted file mode 100644 index 94d7601fb6..0000000000 --- a/py/samples/provider-vertex-ai-rerank-eval/src/main.py +++ /dev/null @@ -1,389 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Vertex AI Rerankers and Evaluators Demo. - -This sample demonstrates: -- Semantic document reranking for RAG quality improvement -- Model output evaluation using Vertex AI metrics (BLEU, ROUGE, fluency, safety, etc.) - -Prerequisites: -- GOOGLE_CLOUD_PROJECT environment variable set -- gcloud auth application-default login -- Discovery Engine API enabled (for rerankers) -- Vertex AI API enabled (for evaluators) -""" - -from typing import Any, cast - -import structlog -from pydantic import BaseModel - -from genkit.ai import Genkit -from genkit.blocks.document import Document -from genkit.core.typing import BaseDataPoint, DocumentData, Score -from genkit.plugins.google_genai import VertexAI -from samples.shared.logging import setup_sample - -setup_sample() - -logger = structlog.get_logger(__name__) - - -ai = Genkit( - plugins=[ - VertexAI(location='us-central1'), - ], - model='vertexai/gemini-2.5-flash', -) - - -class RerankResult(BaseModel): - """Result of a rerank operation.""" - - query: str - ranked_documents: list[dict[str, Any]] - - -@ai.flow() -async def rerank_documents(query: str = 'How do neural networks learn?') -> RerankResult: - """Rerank documents based on relevance to query. - - This demonstrates using Vertex AI's semantic reranker to re-order - documents by their semantic relevance to a query. Useful for improving - RAG (Retrieval-Augmented Generation) quality. - """ - # Sample documents to rerank (in a real app, these would come from a retriever) - documents: list[Document] = [ - Document.from_text('Neural networks learn through backpropagation, adjusting weights based on errors.'), - Document.from_text('Python is a popular programming language for machine learning.'), - Document.from_text('The gradient descent algorithm minimizes the loss function during training.'), - Document.from_text('Cats are popular pets known for their independence.'), - Document.from_text('Deep learning models use multiple layers to extract hierarchical features.'), - Document.from_text('The weather today is sunny with a high of 75 degrees.'), - Document.from_text('Transformers use attention mechanisms to process sequential data efficiently.'), - ] - - # Rerank documents using Vertex AI semantic reranker - # Document extends DocumentData, so we can cast and pass documents directly - ranked_docs = await ai.rerank( - reranker='vertexai/semantic-ranker-default@latest', - query=query, - documents=cast(list[DocumentData], documents), - options={'top_n': 5}, - ) - - # Format results - results: list[dict[str, Any]] = [] - for doc in ranked_docs: - results.append({ - 'text': doc.text(), - 'score': doc.score, - }) - - return RerankResult(query=query, ranked_documents=results) - - -@ai.flow() -async def rag_with_reranking(question: str = 'What is machine learning?') -> str: - """Full RAG pipeline with reranking. - - Demonstrates a two-stage retrieval pattern: - 1. Initial retrieval (simulated with sample docs) - 2. Reranking for quality - 3. Generation using top-k results - """ - # Simulated retrieval results (in production, use a real retriever) - retrieved_docs: list[Document] = [ - Document.from_text('Machine learning is a subset of artificial intelligence.'), - Document.from_text('Supervised learning uses labeled data to train models.'), - Document.from_text('The stock market closed higher today.'), - Document.from_text('ML algorithms can identify patterns in large datasets.'), - Document.from_text('Unsupervised learning finds hidden patterns without labels.'), - Document.from_text('Pizza is a popular Italian dish.'), - Document.from_text('Deep learning uses neural networks with many layers.'), - Document.from_text('Reinforcement learning learns from rewards and penalties.'), - ] - - # Stage 2: Rerank for quality - # Document extends DocumentData, so we can cast and pass documents directly - ranked_docs = await ai.rerank( - reranker='vertexai/semantic-ranker-default@latest', - query=question, - documents=cast(list[DocumentData], retrieved_docs), - options={'top_n': 3}, - ) - - # Build context from top-ranked documents - context = '\n'.join([f'- {doc.text()}' for doc in ranked_docs]) - - # Stage 3: Generate answer using reranked context - response = await ai.generate( - model='vertexai/gemini-2.5-flash', - prompt=f"""Answer the following question based on the provided context. - -Context: -{context} - -Question: {question} - -Answer:""", - ) - - return response.text - - -class EvalResult(BaseModel): - """Result of an evaluation.""" - - metric: str - scores: list[dict[str, Any]] - - -def _extract_score(evaluation: Score | list[Score]) -> float | str | bool | None: - """Extract score from evaluation result.""" - if isinstance(evaluation, list): - return evaluation[0].score if evaluation else None - return evaluation.score - - -def _extract_reasoning(evaluation: Score | list[Score]) -> str | None: - """Extract reasoning from evaluation result.""" - if isinstance(evaluation, list): - if evaluation and evaluation[0].details: - return evaluation[0].details.reasoning - return None - if evaluation.details: - return evaluation.details.reasoning - return None - - -@ai.flow() -async def evaluate_fluency() -> EvalResult: - """Evaluate text fluency using Vertex AI. - - Fluency measures how natural and readable the text is. - Scores range from 1 (poor) to 5 (excellent). - """ - dataset: list[BaseDataPoint] = [ - BaseDataPoint( - input='Write a sentence about AI.', - output=( - 'Artificial intelligence is transforming industries by enabling ' - 'machines to learn from data and make intelligent decisions.' - ), - ), - BaseDataPoint( - input='Describe machine learning.', - # Intentionally poor grammar for demonstration - output='Machine learning algorithms data patterns the find automatically.', - ), - BaseDataPoint( - input='Explain neural networks.', - output=( - 'Neural networks are computational models inspired by the human brain, ' - 'consisting of interconnected nodes that process information in layers.' - ), - ), - ] - - results = await ai.evaluate( - evaluator='vertexai/fluency', - dataset=dataset, - ) - - scores: list[dict[str, Any]] = [] - for i, result in enumerate(results.root): - scores.append({ - 'sample_index': i, - 'test_case_id': result.test_case_id, - 'score': _extract_score(result.evaluation), - 'reasoning': _extract_reasoning(result.evaluation), - }) - - return EvalResult(metric='fluency', scores=scores) - - -@ai.flow() -async def evaluate_safety() -> EvalResult: - """Evaluate text safety using Vertex AI. - - Safety checks for harmful, biased, or inappropriate content. - Higher scores indicate safer content. - """ - dataset: list[BaseDataPoint] = [ - BaseDataPoint( - input='Write about technology.', - output=( - 'Technology has revolutionized communication, enabling people to connect instantly across the globe.' - ), - ), - BaseDataPoint( - input='Describe a helpful assistant.', - output=( - 'A helpful AI assistant provides accurate information, ' - 'respects user privacy, and declines harmful requests.' - ), - ), - ] - - results = await ai.evaluate( - evaluator='vertexai/safety', - dataset=dataset, - ) - - scores: list[dict[str, Any]] = [] - for i, result in enumerate(results.root): - scores.append({ - 'sample_index': i, - 'test_case_id': result.test_case_id, - 'score': _extract_score(result.evaluation), - }) - - return EvalResult(metric='safety', scores=scores) - - -@ai.flow() -async def evaluate_groundedness() -> EvalResult: - """Evaluate groundedness using Vertex AI. - - Groundedness checks if the output is factually grounded in the provided context. - This helps detect hallucinations in RAG applications. - """ - dataset: list[BaseDataPoint] = [ - BaseDataPoint( - input='What is the capital of France?', - output='The capital of France is Paris.', - context=[ - 'France is a country in Western Europe. Its capital city is Paris, which is known for the Eiffel Tower.' - ], - ), - BaseDataPoint( - input='What is the population of Paris?', - # Hallucinated - context doesn't mention population - output='Paris has a population of about 12 million people.', - context=['Paris is the capital of France. It is known for art, fashion, and culture.'], - ), - BaseDataPoint( - input='What is France known for?', - output='France is known for wine, cheese, and the Eiffel Tower.', - context=[ - 'France is famous for its cuisine, especially wine and cheese. ' - 'The Eiffel Tower in Paris is a major landmark.' - ], - ), - ] - - results = await ai.evaluate( - evaluator='vertexai/groundedness', - dataset=dataset, - ) - - scores: list[dict[str, Any]] = [] - for i, result in enumerate(results.root): - scores.append({ - 'sample_index': i, - 'test_case_id': result.test_case_id, - 'score': _extract_score(result.evaluation), - 'reasoning': _extract_reasoning(result.evaluation), - }) - - return EvalResult(metric='groundedness', scores=scores) - - -@ai.flow() -async def evaluate_bleu() -> EvalResult: - """Evaluate using BLEU score. - - BLEU (Bilingual Evaluation Understudy) compares output to a reference. - Commonly used for translation and text generation quality. - Scores range from 0 to 1, with higher being better. - """ - dataset: list[BaseDataPoint] = [ - BaseDataPoint( - input='Translate to French: Hello, how are you?', - output='Bonjour, comment allez-vous?', - reference='Bonjour, comment allez-vous?', # Perfect match - ), - BaseDataPoint( - input='Translate to French: Good morning', - output='Bon matin', - reference='Bonjour', # Different but valid translation - ), - ] - - results = await ai.evaluate( - evaluator='vertexai/bleu', - dataset=dataset, - ) - - scores: list[dict[str, Any]] = [] - for i, result in enumerate(results.root): - scores.append({ - 'sample_index': i, - 'test_case_id': result.test_case_id, - 'score': _extract_score(result.evaluation), - }) - - return EvalResult(metric='bleu', scores=scores) - - -@ai.flow() -async def evaluate_summarization() -> EvalResult: - """Evaluate summarization quality using Vertex AI. - - Summarization quality assesses how well a summary captures the key points - of the original text. - """ - dataset: list[BaseDataPoint] = [ - BaseDataPoint( - input='Summarize this article about climate change.', - output='Climate change is causing rising temperatures and extreme weather events globally.', - context=[ - 'Climate change refers to long-term shifts in temperatures and weather patterns. ' - 'Human activities have been the main driver since the 1800s, primarily due to ' - 'burning fossil fuels. This has led to rising global temperatures, melting ice ' - 'caps, rising sea levels, and more frequent extreme weather events like ' - 'hurricanes, droughts, and floods.' - ], - ), - ] - - results = await ai.evaluate( - evaluator='vertexai/summarization_quality', - dataset=dataset, - ) - - scores: list[dict[str, Any]] = [] - for i, result in enumerate(results.root): - scores.append({ - 'sample_index': i, - 'test_case_id': result.test_case_id, - 'score': _extract_score(result.evaluation), - 'reasoning': _extract_reasoning(result.evaluation), - }) - - return EvalResult(metric='summarization_quality', scores=scores) - - -async def main() -> None: - """Main function.""" - # Example run logic can go here or be empty for pure flow server - pass - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/samples/provider-vertex-ai-vector-search-bigquery/LICENSE b/py/samples/provider-vertex-ai-vector-search-bigquery/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-bigquery/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/samples/provider-vertex-ai-vector-search-bigquery/README.md b/py/samples/provider-vertex-ai-vector-search-bigquery/README.md deleted file mode 100644 index 2b492ae48a..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-bigquery/README.md +++ /dev/null @@ -1,101 +0,0 @@ -# Vertex AI - Vector Search BigQuery - -An example demonstrating the use Vector Search API with BigQuery retriever for Vertex AI - -### Monitoring and Running - -For an enhanced development experience, use the provided `run.sh` script to start the sample with automatic reloading: - -```bash -./run.sh -``` - -This script uses `watchmedo` to monitor changes in: -- `src/` (Python logic) -- `../../packages` (Genkit core) -- `../../plugins` (Genkit plugins) -- File patterns: `*.py`, `*.prompt`, `*.json` - -Changes will automatically trigger a restart of the sample. You can also pass command-line arguments directly to the script, e.g., `./run.sh --some-flag`. - -## Setup environment - -1. Install [GCP CLI](https://cloud.google.com/sdk/docs/install). -2. Run the following code to connect to VertexAI. -```bash -gcloud auth application-default login -``` -3. Set the following env vars to run the sample -``` -export LOCATION='' -export PROJECT_ID='' -export BIGQUERY_DATASET_NAME='' -export BIGQUERY_TABLE_NAME='' -export VECTOR_SEARCH_DEPLOYED_INDEX_ID='' -export VECTOR_SEARCH_INDEX_ENDPOINT_PATH='' -export VECTOR_SEARCH_API_ENDPOINT='' -``` -4. Run the sample. - -## Env vars definition -| Variable | Definition | -| ----------------------------------- | --------------------------------------------------------------------------------------------------------- | -| `LOCATION` | The Google Cloud region or multi-region where your resources (e.g., Vertex AI Index, BigQuery dataset) are located. Example: `us-central1`. | -| `PROJECT_ID` | The name or unique identifier for your Google Cloud Project. | -| `BIGQUERY_DATASET_NAME` | The name of the BigQuery dataset that contains your source data or will store results. | -| `BIGQUERY_TABLE_NAME` | The name of the specific table within the BigQuery dataset. | -| `VECTOR_SEARCH_DEPLOYED_INDEX_ID` | The ID of your deployed Vertex AI Vector Search index that you want to query. Numeric identifier. | -| `VECTOR_SEARCH_INDEX_ENDPOINT_PATH` | The full storage path of the Vertex AI Vector Search Index Endpoint. Example: `projects/YOUR_PROJECT_ID/locations/YOUR_LOCATION/indexEndpoints/YOUR_INDEX_ENDPOINT_ID`. | -| `VECTOR_SEARCH_API_ENDPOINT` | The regional API endpoint for making calls to the Vertex AI Vector Search service. Example: `YOUR_LOCATION-aiplatform.googleapis.com`. | - -## Run the sample - -```bash -genkit start -- uv run src/main.py -``` - -## Set up env for sample -In the file `setup_env.py` you will find some code that will help you to create the bigquery dataset, table with the expected schema, encode the content of the table and push this to the VertexAI Vector Search index. -This index must be created with update method set as `stream`. VertexAI Index is expected to be already created. - -## Testing This Demo - -1. **Prerequisites** - Set up GCP resources: - ```bash - # Required environment variables - export LOCATION=us-central1 - export PROJECT_ID=your_project_id - export BIGQUERY_DATASET_NAME=your_dataset - export BIGQUERY_TABLE_NAME=your_table - export VECTOR_SEARCH_DEPLOYED_INDEX_ID=your_deployed_index_id - export VECTOR_SEARCH_INDEX_ENDPOINT_PATH=your_endpoint_path - export VECTOR_SEARCH_API_ENDPOINT=your_api_endpoint - - # Authenticate with GCP - gcloud auth application-default login - ``` - -2. **GCP Setup Required**: - - Create Vertex AI Vector Search index - - Deploy index to an endpoint - - Create BigQuery dataset and table with embeddings - - Ensure table schema matches expected format - -3. **Run the demo**: - ```bash - cd py/samples/provider-vertex-ai-vector-search-bigquery - ./run.sh - ``` - -4. **Open DevUI** at http://localhost:4000 - -5. **Test the flows**: - - [ ] `retrieve_documents` - Vector similarity search - - [ ] Test with limit options - - [ ] Check performance metrics in output - -6. **Expected behavior**: - - Query is embedded and sent to Vector Search - - Similar vectors are found and IDs returned - - BigQuery is queried for full document content - - Duration metrics show performance diff --git a/py/samples/provider-vertex-ai-vector-search-bigquery/pyproject.toml b/py/samples/provider-vertex-ai-vector-search-bigquery/pyproject.toml deleted file mode 100644 index 2bb3bcf5e3..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-bigquery/pyproject.toml +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Private :: Do Not Upload", -] -dependencies = [ - "rich>=13.0.0", - "genkit", - "genkit-plugin-google-genai", - "genkit-plugin-vertex-ai", - "google-cloud-aiplatform", - "google-cloud-bigquery", - "pydantic>=2.10.5", - "strenum>=0.4.15; python_version < '3.11'", - "structlog>=25.2.0", - "uvloop>=0.21.0", -] -description = "An example demonstrating the use Vector Search API with BigQuery retriever for Vertex AI" -license = "Apache-2.0" -name = "provider-vertex-ai-vector-search-bigquery" -readme = "README.md" -requires-python = ">=3.10" -version = "0.2.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src/sample"] diff --git a/py/samples/provider-vertex-ai-vector-search-bigquery/run.sh b/py/samples/provider-vertex-ai-vector-search-bigquery/run.sh deleted file mode 100755 index f74499d7b6..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-bigquery/run.sh +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -# Vertex AI Vector Search with BigQuery Demo -# ========================================== -# -# Demonstrates enterprise-grade vector search using GCP services. -# -# Prerequisites: -# - GOOGLE_CLOUD_PROJECT environment variable set -# - gcloud CLI authenticated -# - BigQuery dataset and Vertex AI index configured -# -# Usage: -# ./run.sh # Start the demo with Dev UI -# ./run.sh --help # Show this help message - -set -euo pipefail - -cd "$(dirname "$0")" -source "../_common.sh" - -print_help() { - print_banner "Vertex AI Vector Search (BigQuery)" "πŸ“Š" - echo "Usage: ./run.sh [options]" - echo "" - echo "Options:" - echo " --help Show this help message" - echo "" - echo "Environment Variables:" - echo " GOOGLE_CLOUD_PROJECT Required. Your GCP project ID" - echo "" - echo "Getting Started:" - echo " 1. Authenticate: gcloud auth application-default login" - echo " 2. Set project: export GOOGLE_CLOUD_PROJECT=your-project" - echo " 3. Configure BigQuery dataset and Vertex AI index" - print_help_footer -} - -case "${1:-}" in - --help|-h) - print_help - exit 0 - ;; -esac - -print_banner "Vertex AI Vector Search (BigQuery)" "πŸ“Š" - -check_env_var "PROJECT_ID" "" || exit 1 -check_env_var "LOCATION" "" || exit 1 -check_env_var "BIGQUERY_DATASET_NAME" "" || exit 1 -check_env_var "BIGQUERY_TABLE_NAME" "" || exit 1 -check_env_var "VECTOR_SEARCH_DEPLOYED_INDEX_ID" "" || exit 1 - -install_deps - -genkit_start_with_browser -- \ - uv tool run --from watchdog watchmedo auto-restart \ - -d src \ - -d ../../packages \ - -d ../../plugins \ - -p '*.py;*.prompt;*.json' \ - -R \ - -- uv run src/main.py "$@" diff --git a/py/samples/provider-vertex-ai-vector-search-bigquery/src/main.py b/py/samples/provider-vertex-ai-vector-search-bigquery/src/main.py deleted file mode 100755 index 19c11baccf..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-bigquery/src/main.py +++ /dev/null @@ -1,243 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Vertex AI Vector Search with BigQuery sample. - -This sample demonstrates how to use Vertex AI Vector Search with BigQuery -as the document store for large-scale analytics-friendly vector search. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ BigQuery β”‚ Google's data warehouse. Store and query β”‚ - β”‚ β”‚ massive amounts of data (petabytes!) fast. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Vertex AI Vector β”‚ Google's enterprise vector search. Handles β”‚ - β”‚ Search β”‚ billions of vectors with sub-second latency. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Dataset/Table β”‚ BigQuery organization. Dataset = folder, β”‚ - β”‚ β”‚ Table = spreadsheet with your data. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Deployed Index β”‚ Your vector index running in the cloud. β”‚ - β”‚ β”‚ Ready 24/7 to answer similarity queries. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Analytics + AI β”‚ Combine SQL analytics with AI search. β”‚ - β”‚ β”‚ The best of both worlds! β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow (BigQuery + Vector Search):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW BIGQUERY + VERTEX VECTOR SEARCH WORK TOGETHER β”‚ - β”‚ β”‚ - β”‚ Query: "Find similar products to this one" β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) Convert to embedding β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ Query β†’ [0.3, -0.2, 0.7, ...] β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Search Vector Index β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Vertex AI β”‚ Returns IDs: ["prod_1", "prod_2", ...] β”‚ - β”‚ β”‚ Vector Search β”‚ (ranked by similarity) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Fetch from BigQuery (with extra analytics) β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ BigQuery β”‚ JOIN with sales data, filter by region, β”‚ - β”‚ β”‚ (SQL query) β”‚ add business logic, return full records β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (4) Rich results with analytics context β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Your App β”‚ Products + sales data + similarity scores β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Features -============ -| Feature Description | Example Function / Code Snippet | -|-----------------------------------------|-------------------------------------| -| BigQuery Vector Search Definition | `define_vertex_vector_search_big_query`| -| BigQuery Client Integration | `bigquery.Client()` | -| Document Retrieval with Filters | `ai.retrieve(..., options={'limit': ...})`| -| Performance Metrics | Duration tracking | - -Testing This Demo -================= -1. **Prerequisites** - Set up GCP resources: - ```bash - # Required environment variables - export LOCATION=us-central1 - export PROJECT_ID=your_project_id - export BIGQUERY_DATASET_NAME=your_dataset - export BIGQUERY_TABLE_NAME=your_table - export VECTOR_SEARCH_DEPLOYED_INDEX_ID=your_deployed_index_id - export VECTOR_SEARCH_INDEX_ENDPOINT_PATH=your_endpoint_path - export VECTOR_SEARCH_API_ENDPOINT=your_api_endpoint - - # Authenticate with GCP - gcloud auth application-default login - ``` - -2. **GCP Setup Required**: - - Create Vertex AI Vector Search index - - Deploy index to an endpoint - - Create BigQuery dataset and table with embeddings - - Ensure table schema matches expected format - -3. **Run the demo**: - ```bash - cd py/samples/provider-vertex-ai-vector-search-bigquery - ./run.sh - ``` - -4. **Open DevUI** at http://localhost:4000 - -5. **Test the flows**: - - [ ] `retrieve_documents` - Vector similarity search - - [ ] Test with limit options - - [ ] Check performance metrics in output - -6. **Expected behavior**: - - Query is embedded and sent to Vector Search - - Similar vectors are found and IDs returned - - BigQuery is queried for full document content - - Duration metrics show performance -""" - -import os -import time - -from google.cloud import aiplatform, bigquery -from pydantic import BaseModel, Field - -from genkit.ai import Genkit -from genkit.blocks.document import Document -from genkit.core.logging import get_logger -from genkit.plugins.google_genai import VertexAI -from genkit.plugins.vertex_ai import define_vertex_vector_search_big_query -from samples.shared.logging import setup_sample - -setup_sample() - -LOCATION = os.getenv('LOCATION') -PROJECT_ID = os.getenv('PROJECT_ID') - -BIGQUERY_DATASET_NAME = os.getenv('BIGQUERY_DATASET_NAME') -BIGQUERY_TABLE_NAME = os.getenv('BIGQUERY_TABLE_NAME') - -VECTOR_SEARCH_DEPLOYED_INDEX_ID = os.getenv('VECTOR_SEARCH_DEPLOYED_INDEX_ID') -VECTOR_SEARCH_INDEX_ENDPOINT_PATH = os.getenv('VECTOR_SEARCH_INDEX_ENDPOINT_PATH') -VECTOR_SEARCH_API_ENDPOINT = os.getenv('VECTOR_SEARCH_API_ENDPOINT') - -bq_client = bigquery.Client(project=PROJECT_ID) -aiplatform.init(project=PROJECT_ID, location=LOCATION) - -logger = get_logger(__name__) - -ai = Genkit(plugins=[VertexAI()]) - -# Define Vertex AI Vector Search with BigQuery -define_vertex_vector_search_big_query( - ai, - name='my-vector-search', - embedder='vertexai/gemini-embedding-001', - embedder_options={ - 'task': 'RETRIEVAL_DOCUMENT', - 'output_dimensionality': 128, - }, - bq_client=bq_client, - dataset_id=BIGQUERY_DATASET_NAME or 'default_dataset', - table_id=BIGQUERY_TABLE_NAME or 'default_table', -) - - -class QueryFlowInputSchema(BaseModel): - """Input schema.""" - - query: str = Field(default='document 1', description='Search query text') - k: int = Field(default=5, description='Number of results to return') - - -class QueryFlowOutputSchema(BaseModel): - """Output schema.""" - - result: list[dict[str, object]] - length: int - time: int - - -@ai.flow(name='queryFlow') -async def query_flow(_input: QueryFlowInputSchema) -> QueryFlowOutputSchema: - """Executes a vector search with VertexAI Vector Search.""" - start_time = time.time() - - query_document = Document.from_text(text=_input.query) - query_document.metadata = { - 'api_endpoint': VECTOR_SEARCH_API_ENDPOINT, - 'index_endpoint_path': VECTOR_SEARCH_INDEX_ENDPOINT_PATH, - 'deployed_index_id': VECTOR_SEARCH_DEPLOYED_INDEX_ID, - } - - response = await ai.retrieve( - retriever='my-vector-search', - query=query_document, - options={'limit': 10}, - ) - - end_time = time.time() - - duration = int(end_time - start_time) - - result_data = [] - for doc in response.documents: - metadata = doc.metadata or {} - result_data.append({ - 'id': metadata.get('id'), - 'text': doc.content[0].root.text if doc.content and doc.content[0].root.text else '', - 'distance': metadata.get('distance'), - }) - - result_data = sorted(result_data, key=lambda x: x['distance']) - - return QueryFlowOutputSchema( - result=result_data, - length=len(result_data), - time=duration, - ) - - -async def main() -> None: - """Main function.""" - query_input = QueryFlowInputSchema( - query='Content for doc', - k=3, - ) - - result = await query_flow(query_input) - await logger.ainfo(str(result)) - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/samples/provider-vertex-ai-vector-search-bigquery/src/setup_env.py b/py/samples/provider-vertex-ai-vector-search-bigquery/src/setup_env.py deleted file mode 100644 index 9ea4405056..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-bigquery/src/setup_env.py +++ /dev/null @@ -1,251 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Example of using Genkit to fill VertexAI Index for Vector Search with BigQuery.""" - -import json -import os -from typing import Any - -import structlog -from google.cloud import aiplatform, aiplatform_v1, bigquery - -from genkit.ai import Genkit -from genkit.core.typing import DocumentPart -from genkit.plugins.google_genai import VertexAI -from genkit.plugins.vertex_ai import define_vertex_vector_search_big_query -from genkit.types import Document, TextPart - -# Environment Variables -LOCATION = os.environ['LOCATION'] -PROJECT_ID = os.environ['PROJECT_ID'] -EMBEDDING_MODEL = 'gemini-embedding-001' - -BIGQUERY_DATASET_NAME = os.environ['BIGQUERY_DATASET_NAME'] -BIGQUERY_TABLE_NAME = os.environ['BIGQUERY_TABLE_NAME'] - -VECTOR_SEARCH_INDEX_ID = os.environ['VECTOR_SEARCH_INDEX_ID'] - -bq_client = bigquery.Client(project=PROJECT_ID) -aiplatform.init(project=PROJECT_ID, location=LOCATION) - -logger = structlog.get_logger(__name__) - -ai = Genkit(plugins=[VertexAI()]) - -define_vertex_vector_search_big_query( - ai, - name='bigquery-vector-search', - embedder=f'vertexai/{EMBEDDING_MODEL}', - bq_client=bq_client, - dataset_id=BIGQUERY_DATASET_NAME, - table_id=BIGQUERY_TABLE_NAME, -) - - -@ai.flow(name='generateEmbeddings') -async def generate_embeddings() -> None: - """Generates document embeddings and upserts them to the Vertex AI Vector Search index. - - This flow retrieves data from BigQuery, generates embeddings for the documents, - and then upserts these embeddings to the specified Vector Search index. - """ - toy_documents = [ - { - 'id': 'doc1', - 'content': {'title': 'Document 1', 'body': 'This is the content of document 1.'}, - 'metadata': {'author': 'Alice', 'date': '2024-01-15'}, - }, - { - 'id': 'doc2', - 'content': {'title': 'Document 2', 'body': 'This is the content of document 2.'}, - 'metadata': {'author': 'Bob', 'date': '2024-02-20'}, - }, - { - 'id': 'doc3', - 'content': {'title': 'Document 3', 'body': 'Content for doc 3'}, - 'metadata': {'author': 'Charlie', 'date': '2024-03-01'}, - }, - ] - - create_bigquery_dataset_and_table( - PROJECT_ID, - LOCATION, - BIGQUERY_DATASET_NAME, - BIGQUERY_TABLE_NAME, - toy_documents, - ) - - results_dict = get_data_from_bigquery( - bq_client=bq_client, - project_id=PROJECT_ID, - dataset_id=BIGQUERY_DATASET_NAME, - table_id=BIGQUERY_TABLE_NAME, - ) - - genkit_documents = [Document(content=[DocumentPart(root=TextPart(text=text))]) for text in results_dict.values()] - - embed_response = await ai.embed_many( - embedder=f'vertexai/{EMBEDDING_MODEL}', - content=genkit_documents, - options={'task': 'RETRIEVAL_DOCUMENT', 'output_dimensionality': 128}, - ) - - embeddings = [emb.embedding for emb in embed_response] - - ids = list(results_dict.keys())[: len(embeddings)] - data_embeddings = list(zip(ids, embeddings, strict=True)) - - upsert_data = [(str(id), embedding) for id, embedding in data_embeddings] - upsert_index(PROJECT_ID, LOCATION, VECTOR_SEARCH_INDEX_ID, upsert_data) - - -def create_bigquery_dataset_and_table( - project_id: str, - location: str, - dataset_id: str, - table_id: str, - documents: list[dict[str, Any]], -) -> None: - """Creates a BigQuery dataset and table, and inserts documents. - - Args: - project_id: The ID of the Google Cloud project. - location: The location for the BigQuery resources. - dataset_id: The ID of the BigQuery dataset. - table_id: The ID of the BigQuery table. - documents: A list of dictionaries, where each dictionary represents a document - with 'id', 'content', and 'metadata' keys. 'content' and 'metadata' - are expected to be JSON serializable. - """ - client = bigquery.Client(project=project_id) - dataset_ref = bigquery.DatasetReference(project_id, dataset_id) - dataset = bigquery.Dataset(dataset_ref) - dataset.location = location - - try: - dataset = client.create_dataset(dataset, exists_ok=True) - logger.debug('Dataset %s.%s created.', client.project, dataset.dataset_id) - except Exception as e: - logger.exception('Error creating dataset: %s', e) - raise e - - schema = [ - bigquery.SchemaField('id', 'STRING', mode='REQUIRED'), - bigquery.SchemaField('content', 'JSON'), - bigquery.SchemaField('metadata', 'JSON'), - ] - - table_ref = dataset_ref.table(table_id) - table = bigquery.Table(table_ref, schema=schema) - try: - table = client.create_table(table, exists_ok=True) - logger.debug( - 'Table %s.%s.%s created.', - table.project, - table.dataset_id, - table.table_id, - ) - except Exception as e: - logger.exception('Error creating table: %s', e) - raise e - - rows_to_insert = [ - { - 'id': doc['id'], - 'content': json.dumps(doc['content']), - 'metadata': json.dumps(doc['metadata']), - } - for doc in documents - ] - - errors = client.insert_rows_json(table, rows_to_insert) - if errors: - logger.error('Errors inserting rows: %s', errors) - raise Exception(f'Failed to insert rows: {errors}') - else: - logger.debug('Inserted %s rows into BigQuery.', len(rows_to_insert)) - - -def get_data_from_bigquery( - bq_client: bigquery.Client, - project_id: str, - dataset_id: str, - table_id: str, -) -> dict[str, str]: - """Retrieves data from a BigQuery table. - - Args: - bq_client: The BigQuery client. - project_id: The ID of the Google Cloud project. - dataset_id: The ID of the BigQuery dataset. - table_id: The ID of the BigQuery table. - - Returns: - A dictionary where keys are document IDs and values are JSON strings - representing the document content. - """ - table_ref = bigquery.TableReference.from_string(f'{project_id}.{dataset_id}.{table_id}') - query = f'SELECT id, content FROM `{table_ref}`' # noqa: S608 - table ref from trusted config - query_job = bq_client.query(query) - rows = query_job.result() - - results = {row['id']: json.dumps(row['content']) for row in rows} - logger.debug('Found %s rows with different ids into BigQuery.', len(results)) - - return results - - -def upsert_index( - project_id: str, - region: str, - index_name: str, - data: list[tuple[str, list[float]]], -) -> None: - """Upserts data points to a Vertex AI Index using batch processing. - - Args: - project_id: The ID of your Google Cloud project. - region: The region where the Index is located. - index_name: The name of the Vertex AI Index. - data: A list of tuples, where each tuple contains (id, embedding). - id should be a string, and embedding should be a list of floats. - """ - aiplatform.init(project=project_id, location=region) - - index_client = aiplatform_v1.IndexServiceClient( - client_options={'api_endpoint': f'{region}-aiplatform.googleapis.com'} - ) - - index_path = index_client.index_path(project=project_id, location=region, index=index_name) - - datapoints = [aiplatform_v1.IndexDatapoint(datapoint_id=id, feature_vector=embedding) for id, embedding in data] - - logger.debug('Attempting to insert %s rows into Index %s', len(datapoints), index_path) - - upsert_request = aiplatform_v1.UpsertDatapointsRequest(index=index_path, datapoints=datapoints) - - index_client.upsert_datapoints(request=upsert_request) - logger.info('Upserted %s datapoints.', len(datapoints)) - - -async def main() -> None: - """Main function.""" - await logger.ainfo(await generate_embeddings()) - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/samples/provider-vertex-ai-vector-search-firestore/LICENSE b/py/samples/provider-vertex-ai-vector-search-firestore/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-firestore/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/samples/provider-vertex-ai-vector-search-firestore/README.md b/py/samples/provider-vertex-ai-vector-search-firestore/README.md deleted file mode 100644 index 43a3a80d8c..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-firestore/README.md +++ /dev/null @@ -1,94 +0,0 @@ -# Vertex AI - Vector Search Firestore - -An example demonstrating the use Vector Search API with Firestore retriever for Vertex AI - -### Monitoring and Running - -For an enhanced development experience, use the provided `run.sh` script to start the sample with automatic reloading: - -```bash -./run.sh -``` - -This script uses `watchmedo` to monitor changes in: -- `src/` (Python logic) -- `../../packages` (Genkit core) -- `../../plugins` (Genkit plugins) -- File patterns: `*.py`, `*.prompt`, `*.json` - -Changes will automatically trigger a restart of the sample. You can also pass command-line arguments directly to the script, e.g., `./run.sh --some-flag`. - -## Setup environment - -1. Install [GCP CLI](https://cloud.google.com/sdk/docs/install). -2. Run the following code to connect to VertexAI. -```bash -gcloud auth application-default login -``` -3. Set the following env vars to run the sample -``` -export LOCATION='' -export PROJECT_ID='' -export FIRESTORE_COLLECTION='' -export VECTOR_SEARCH_DEPLOYED_INDEX_ID='' -export VECTOR_SEARCH_INDEX_ENDPOINT_PATH='' -export VECTOR_SEARCH_API_ENDPOINT='' -``` -4. Run the sample. - -## Env vars definition -| Variable | Definition | -| ----------------------------------- | --------------------------------------------------------------------------------------------------------- | -| `LOCATION` | The Google Cloud region or multi-region where your resources (e.g., Vertex AI Index, Firestore database) are located. Example: `us-central1`. | -| `PROJECT_ID` | The name or unique identifier for your Google Cloud Project. | -| `FIRESTORE_COLLECTION` | The name of the Firestore collection used, for example, to store metadata associated with your vectors or the source documents. | -| `VECTOR_SEARCH_DEPLOYED_INDEX_ID` | The ID of your deployed Vertex AI Vector Search index that you want to query. | -| `VECTOR_SEARCH_INDEX_ENDPOINT_PATH` | The full storage path of the Vertex AI Vector Search Index Endpoint. Example: `projects/YOUR_PROJECT_ID/locations/YOUR_LOCATION/indexEndpoints/YOUR_INDEX_ENDPOINT_ID`. | -| `VECTOR_SEARCH_API_ENDPOINT` | The regional API endpoint for making calls to the Vertex AI Vector Search service. Example: `YOUR_LOCATION-aiplatform.googleapis.com`. | - -## Run the sample - -```bash -genkit start -- uv run src/main.py -``` - -## Testing This Demo - -1. **Prerequisites** - Set up GCP resources: - ```bash - # Required environment variables - export LOCATION=us-central1 - export PROJECT_ID=your_project_id - export FIRESTORE_COLLECTION=your_collection_name - export VECTOR_SEARCH_DEPLOYED_INDEX_ID=your_deployed_index_id - export VECTOR_SEARCH_INDEX_ENDPOINT_PATH=your_endpoint_path - export VECTOR_SEARCH_API_ENDPOINT=your_api_endpoint - - # Authenticate with GCP - gcloud auth application-default login - ``` - -2. **GCP Setup Required**: - - Create Vertex AI Vector Search index - - Deploy index to an endpoint - - Create Firestore collection with documents - - Ensure documents have matching IDs in both services - -3. **Run the demo**: - ```bash - cd py/samples/provider-vertex-ai-vector-search-firestore - ./run.sh - ``` - -4. **Open DevUI** at http://localhost:4000 - -5. **Test the flows**: - - [ ] `retrieve_documents` - Vector similarity search - - [ ] Check results are ranked by distance - - [ ] Verify Firestore document metadata is returned - -6. **Expected behavior**: - - Query is embedded and sent to Vector Search - - Similar vectors are found and IDs returned - - Firestore is queried for full document content - - Results sorted by similarity distance diff --git a/py/samples/provider-vertex-ai-vector-search-firestore/pyproject.toml b/py/samples/provider-vertex-ai-vector-search-firestore/pyproject.toml deleted file mode 100644 index 1a1144f344..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-firestore/pyproject.toml +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [ - { name = "Google" }, - { name = "Yesudeep Mangalapilly", email = "yesudeep@google.com" }, - { name = "Elisa Shen", email = "mengqin@google.com" }, - { name = "Niraj Nepal", email = "nnepal@google.com" }, -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Environment :: Web Environment", - "Intended Audience :: Developers", - "Operating System :: OS Independent", - "Programming Language :: Python", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries", - "Private :: Do Not Upload", -] -dependencies = [ - "rich>=13.0.0", - "genkit", - "genkit-plugin-google-genai", - "genkit-plugin-vertex-ai", - "google-cloud-aiplatform", - "google-cloud-firestore", - "pydantic>=2.10.5", - "strenum>=0.4.15; python_version < '3.11'", - "structlog>=25.2.0", - "uvloop>=0.21.0", -] -description = "An example demonstrating the use Vector Search API with Firestore retriever for Vertex AI" -license = "Apache-2.0" -name = "provider-vertex-ai-vector-search-firestore" -readme = "README.md" -requires-python = ">=3.10" -version = "0.2.0" - -[project.optional-dependencies] -dev = ["watchdog>=6.0.0"] - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["src/sample"] diff --git a/py/samples/provider-vertex-ai-vector-search-firestore/run.sh b/py/samples/provider-vertex-ai-vector-search-firestore/run.sh deleted file mode 100755 index ec8aac48a4..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-firestore/run.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2026 Google LLC -# SPDX-License-Identifier: Apache-2.0 - -# Vertex AI Vector Search with Firestore Demo -# ============================================ -# -# Demonstrates enterprise-grade vector search using GCP services. -# -# Prerequisites: -# - GOOGLE_CLOUD_PROJECT environment variable set -# - gcloud CLI authenticated -# - Firestore database and Vertex AI index configured -# -# Usage: -# ./run.sh # Start the demo with Dev UI -# ./run.sh --help # Show this help message - -set -euo pipefail - -cd "$(dirname "$0")" -source "../_common.sh" - -print_help() { - print_banner "Vertex AI Vector Search (Firestore)" "πŸ”" - echo "Usage: ./run.sh [options]" - echo "" - echo "Options:" - echo " --help Show this help message" - echo "" - echo "Environment Variables:" - echo " GOOGLE_CLOUD_PROJECT Required. Your GCP project ID" - echo "" - echo "Getting Started:" - echo " 1. Authenticate: gcloud auth application-default login" - echo " 2. Set project: export GOOGLE_CLOUD_PROJECT=your-project" - echo " 3. Configure Firestore and Vertex AI index" - print_help_footer -} - -case "${1:-}" in - --help|-h) - print_help - exit 0 - ;; -esac - -print_banner "Vertex AI Vector Search (Firestore)" "πŸ”" - -check_env_var "PROJECT_ID" "" || exit 1 -check_env_var "LOCATION" "" || exit 1 -check_env_var "FIRESTORE_COLLECTION" "" || exit 1 -check_env_var "VECTOR_SEARCH_DEPLOYED_INDEX_ID" "" || exit 1 - -install_deps - -genkit_start_with_browser -- \ - uv tool run --from watchdog watchmedo auto-restart \ - -d src \ - -d ../../packages \ - -d ../../plugins \ - -p '*.py;*.prompt;*.json' \ - -R \ - -- uv run src/main.py "$@" diff --git a/py/samples/provider-vertex-ai-vector-search-firestore/src/main.py b/py/samples/provider-vertex-ai-vector-search-firestore/src/main.py deleted file mode 100755 index 796025ada2..0000000000 --- a/py/samples/provider-vertex-ai-vector-search-firestore/src/main.py +++ /dev/null @@ -1,241 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Vertex AI Vector Search with Firestore sample. - -This sample demonstrates how to use Vertex AI Vector Search with Firestore -as the document store for enterprise-scale vector similarity search. - -Key Concepts (ELI5):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ Concept β”‚ ELI5 Explanation β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Vertex AI Vector β”‚ Google's enterprise vector search. Handles β”‚ - β”‚ Search β”‚ billions of items with fast nearest-neighbor. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Firestore β”‚ Stores the actual document content. Vector Search β”‚ - β”‚ β”‚ returns IDs, Firestore returns full text. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Deployed Index β”‚ Your vector index running on Google's servers. β”‚ - β”‚ β”‚ Ready to answer similarity queries 24/7. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Endpoint β”‚ The address where your index listens for queries. β”‚ - β”‚ β”‚ Like a phone number for your search service. β”‚ - β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ - β”‚ Distance β”‚ How "far apart" two vectors are. Lower = more β”‚ - β”‚ β”‚ similar. 0 = identical match. β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Data Flow (Enterprise RAG):: - - β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” - β”‚ HOW VERTEX VECTOR SEARCH + FIRESTORE WORK TOGETHER β”‚ - β”‚ β”‚ - β”‚ Query: "What movies are about time travel?" β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (1) Convert to embedding β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Embedder β”‚ Query β†’ [0.3, -0.2, 0.7, ...] β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (2) Send to Vector Search endpoint β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Vertex AI β”‚ Returns: ["doc_123", "doc_456", ...] β”‚ - β”‚ β”‚ Vector Search β”‚ (IDs only, ranked by similarity) β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (3) Fetch full documents from Firestore β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Firestore β”‚ Returns: full document content β”‚ - β”‚ β”‚ (by ID) β”‚ "Back to the Future is a 1985 film..." β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β”‚ β”‚ β”‚ - β”‚ β”‚ (4) Sorted results returned β”‚ - β”‚ β–Ό β”‚ - β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ - β”‚ β”‚ Your App β”‚ Complete docs, ranked by relevance β”‚ - β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - -Key Features -============ -| Feature Description | Example Function / Code Snippet | -|-----------------------------------------|-------------------------------------| -| Firestore Vector Search Definition | `define_vertex_vector_search_firestore`| -| Firestore Async Client Integration | `firestore.AsyncClient()` | -| Document Retrieval | `ai.retrieve()` | -| Result Ranking | Custom sorting by distance | - -Testing This Demo -================= -1. **Prerequisites** - Set up GCP resources: - ```bash - # Required environment variables - export LOCATION=us-central1 - export PROJECT_ID=your_project_id - export FIRESTORE_COLLECTION=your_collection_name - export VECTOR_SEARCH_DEPLOYED_INDEX_ID=your_deployed_index_id - export VECTOR_SEARCH_INDEX_ENDPOINT_PATH=your_endpoint_path - export VECTOR_SEARCH_API_ENDPOINT=your_api_endpoint - - # Authenticate with GCP - gcloud auth application-default login - ``` - -2. **GCP Setup Required**: - - Create Vertex AI Vector Search index - - Deploy index to an endpoint - - Create Firestore collection with documents - - Ensure documents have matching IDs in both services - -3. **Run the demo**: - ```bash - cd py/samples/provider-vertex-ai-vector-search-firestore - ./run.sh - ``` - -4. **Open DevUI** at http://localhost:4000 - -5. **Test the flows**: - - [ ] `retrieve_documents` - Vector similarity search - - [ ] Check results are ranked by distance - - [ ] Verify Firestore document metadata is returned - -6. **Expected behavior**: - - Query is embedded and sent to Vector Search - - Similar vectors are found and IDs returned - - Firestore is queried for full document content - - Results sorted by similarity distance -""" - -import os -import time - -from google.cloud import aiplatform, firestore -from pydantic import BaseModel, Field - -from genkit.ai import Genkit -from genkit.blocks.document import Document -from genkit.core.logging import get_logger -from genkit.core.typing import RetrieverResponse -from genkit.plugins.google_genai import VertexAI -from genkit.plugins.vertex_ai import define_vertex_vector_search_firestore -from samples.shared.logging import setup_sample - -setup_sample() - -LOCATION = os.environ['LOCATION'] -PROJECT_ID = os.environ['PROJECT_ID'] - -FIRESTORE_COLLECTION = os.environ['FIRESTORE_COLLECTION'] - -VECTOR_SEARCH_DEPLOYED_INDEX_ID = os.environ['VECTOR_SEARCH_DEPLOYED_INDEX_ID'] -VECTOR_SEARCH_INDEX_ENDPOINT_PATH = os.environ['VECTOR_SEARCH_INDEX_ENDPOINT_PATH'] -VECTOR_SEARCH_API_ENDPOINT = os.environ['VECTOR_SEARCH_API_ENDPOINT'] - -firestore_client = firestore.AsyncClient(project=PROJECT_ID) -aiplatform.init(project=PROJECT_ID, location=LOCATION) - -logger = get_logger(__name__) - -ai = Genkit(plugins=[VertexAI()]) - -# Define Vertex AI Vector Search with Firestore -define_vertex_vector_search_firestore( - ai, - name='my-vector-search', - embedder='vertexai/gemini-embedding-001', - embedder_options={ - 'task': 'RETRIEVAL_DOCUMENT', - 'output_dimensionality': 128, - }, - firestore_client=firestore_client, - collection_name=FIRESTORE_COLLECTION, -) - - -class QueryFlowInputSchema(BaseModel): - """Input schema.""" - - query: str = Field(default='document 1', description='Search query text') - k: int = Field(default=5, description='Number of results to return') - - -class QueryFlowOutputSchema(BaseModel): - """Output schema.""" - - result: list[dict[str, object]] - length: int - time: int - - -@ai.flow(name='queryFlow') -async def query_flow(_input: QueryFlowInputSchema) -> QueryFlowOutputSchema: - """Executes a vector search with VertexAI Vector Search.""" - start_time = time.time() - - query_document = Document.from_text(text=_input.query) - query_document.metadata = { - 'api_endpoint': VECTOR_SEARCH_API_ENDPOINT, - 'index_endpoint_path': VECTOR_SEARCH_INDEX_ENDPOINT_PATH, - 'deployed_index_id': VECTOR_SEARCH_DEPLOYED_INDEX_ID, - } - - result: RetrieverResponse = await ai.retrieve( - retriever='my-vector-search', - query=query_document, - options={'limit': 10}, - ) - - end_time = time.time() - - duration = int(end_time - start_time) - - result_data = [] - for doc in result.documents: - metadata = doc.metadata or {} - result_data.append({ - 'id': metadata.get('id'), - 'text': doc.content[0].root.text if doc.content and doc.content[0].root.text else '', - 'distance': metadata.get('distance', 0.0), - }) - - result_data = sorted(result_data, key=lambda x: x['distance']) - - return QueryFlowOutputSchema( - result=result_data, - length=len(result_data), - time=duration, - ) - - -async def main() -> None: - """Main function.""" - query_input = QueryFlowInputSchema( - query='Content for doc', - k=3, - ) - - result = await query_flow(query_input) - await logger.ainfo(str(result)) - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/samples/setup.sh b/py/samples/setup.sh index 09ed1116f0..0157b76f00 100755 --- a/py/samples/setup.sh +++ b/py/samples/setup.sh @@ -813,7 +813,7 @@ _configure_aws_auth() { echo "" echo -e " ${BOLD}── AWS Authentication ──${NC}" - echo -e " ${DIM}Used by: provider-amazon-bedrock-hello${NC}" + echo -e " ${DIM}(No sample currently uses AWS)${NC}" # Check if already authenticated. if aws sts get-caller-identity &>/dev/null 2>&1; then @@ -1225,7 +1225,7 @@ _configure_azure_keys() { _configure_aws_keys() { echo "" echo -e " ${BOLD}── AWS Bedrock ──${NC}" - echo -e " ${DIM}Used by: provider-amazon-bedrock-hello${NC}" + echo -e " ${DIM}(No sample currently uses AWS)${NC}" echo -e " ${DIM}Tip: If you ran 'aws configure' above, these are already set.${NC}" _prompt_key "AWS_REGION" \ @@ -1359,7 +1359,6 @@ if ! $CHECK_ONLY; then echo -e "${DIM}This saves credentials to the standard CLI config locations.${NC}" _configure_gcloud_auth - _configure_aws_auth _configure_az_auth fi @@ -1387,7 +1386,6 @@ else _configure_huggingface_keys _configure_cohere_keys _configure_cloudflare_keys - _configure_aws_keys _configure_azure_keys _configure_gcp_keys _configure_observability_keys diff --git a/py/samples/shared/flows.py b/py/samples/shared/flows.py index 7518ec5881..7e6a161d41 100644 --- a/py/samples/shared/flows.py +++ b/py/samples/shared/flows.py @@ -21,14 +21,14 @@ Provider-specific flow logic stays in each sample's main.py. """ -from genkit.ai import Genkit, Output -from genkit.core.action import ActionRunContext -from genkit.core.logging import get_logger -from genkit.types import Media, MediaPart, Message, Part, Role, TextPart +import structlog + +from genkit import Genkit, Media, MediaPart, Message, Part, Role, TextPart +from genkit._core._action import ActionRunContext from .types import CalculatorInput, CurrencyExchangeInput, RpgCharacter, WeatherInput -logger = get_logger(__name__) +logger = structlog.get_logger(__name__) async def calculation_logic(ai: Genkit, input: CalculatorInput, model: str | None = None) -> str: @@ -123,7 +123,7 @@ async def generate_character_logic(ai: Genkit, name: str) -> RpgCharacter: """ result = await ai.generate( prompt=f'Generate a RPG character named {name}.\n{schema_hint}', - output=Output(schema=RpgCharacter), + output_schema=RpgCharacter, ) return result.output @@ -207,14 +207,12 @@ async def generate_streaming_story_logic(ai: Genkit, name: str, ctx: ActionRunCo Returns: Complete story text. """ - stream, response = ai.generate_stream( - prompt=f'Tell me a short story about {name}', - ) - async for chunk in stream: + stream_response = ai.generate_stream(prompt=f'Tell me a short story about {name}') + async for chunk in stream_response.stream: if chunk.text: if ctx is not None: ctx.send_chunk(chunk.text) - return (await response).text + return (await stream_response.response).text async def generate_streaming_with_tools_logic( @@ -234,16 +232,16 @@ async def generate_streaming_with_tools_logic( Returns: The complete generated text. """ - stream, response = ai.generate_stream( + stream_response = ai.generate_stream( model=model, prompt=f'What is the weather in {location}? Describe it poetically.', tools=['get_weather'], ) - async for chunk in stream: + async for chunk in stream_response.stream: if chunk.text: if ctx is not None: ctx.send_chunk(chunk.text) - return (await response).text + return (await stream_response.response).text async def generate_weather_logic(ai: Genkit, input: WeatherInput, model: str | None = None) -> str: diff --git a/py/samples/web-fastapi-bugbot/src/main.py b/py/samples/web-fastapi-bugbot/src/main.py index 4acfdc1877..e27fa1c71a 100644 --- a/py/samples/web-fastapi-bugbot/src/main.py +++ b/py/samples/web-fastapi-bugbot/src/main.py @@ -10,7 +10,6 @@ """ import asyncio -from collections.abc import Awaitable from pathlib import Path from typing import Literal @@ -20,8 +19,7 @@ from pydantic import BaseModel, Field from typing_extensions import Never -from genkit import Genkit, Input, Output -from genkit.ai import FlowWrapper +from genkit import Flow, Genkit from genkit.plugins.fastapi import genkit_fastapi_handler from genkit.plugins.google_genai import GoogleAI @@ -71,10 +69,10 @@ class DiffInput(BaseModel): context: str = '' -security_prompt = ai.prompt('analyze_security', input=Input(schema=CodeInput), output=Output(schema=Analysis)) -bugs_prompt = ai.prompt('analyze_bugs', input=Input(schema=CodeInput), output=Output(schema=Analysis)) -style_prompt = ai.prompt('analyze_style', input=Input(schema=CodeInput), output=Output(schema=Analysis)) -diff_prompt = ai.prompt('analyze_diff', input=Input(schema=DiffInput), output=Output(schema=Analysis)) +security_prompt = ai.prompt('analyze_security', input_schema=CodeInput, output_schema=Analysis) +bugs_prompt = ai.prompt('analyze_bugs', input_schema=CodeInput, output_schema=Analysis) +style_prompt = ai.prompt('analyze_style', input_schema=CodeInput, output_schema=Analysis) +diff_prompt = ai.prompt('analyze_diff', input_schema=DiffInput, output_schema=Analysis) @ai.flow() @@ -139,14 +137,14 @@ async def review_diff_endpoint(diff: str, context: str = '') -> Analysis: @app.post('/flow/review', response_model=None) @genkit_fastapi_handler(ai) -def flow_review() -> FlowWrapper[..., Awaitable[Analysis], Analysis, Never]: +def flow_review() -> Flow[CodeInput, Analysis, Never]: """Expose review_code flow directly via {"data": {"code": "...", "language": "..."}}.""" return review_code @app.post('/flow/security', response_model=None) @genkit_fastapi_handler(ai) -def flow_security() -> FlowWrapper[..., Awaitable[Analysis], Analysis, Never]: +def flow_security() -> Flow[CodeInput, Analysis, Never]: """Expose analyze_security flow directly.""" return analyze_security diff --git a/py/samples/web-flask-hello/src/main.py b/py/samples/web-flask-hello/src/main.py index 7a0e26dc18..a23977c831 100755 --- a/py/samples/web-flask-hello/src/main.py +++ b/py/samples/web-flask-hello/src/main.py @@ -58,10 +58,9 @@ from flask import Flask from pydantic import BaseModel, Field -from genkit.ai import Genkit -from genkit.blocks.model import GenerateResponseWrapper -from genkit.core.action import ActionRunContext -from genkit.core.context import RequestData +from genkit import Genkit, ModelResponse +from genkit._core._action import ActionRunContext +from genkit._core._context import RequestData from genkit.plugins.flask import genkit_flask_handler from genkit.plugins.google_genai import GoogleAI from genkit.plugins.google_genai.models.gemini import GoogleAIGeminiVersion @@ -100,10 +99,13 @@ async def my_context_provider(request: RequestData[dict[str, object]]) -> dict[s async def say_hi( input: SayHiInput, ctx: ActionRunContext | None = None, -) -> GenerateResponseWrapper: +) -> ModelResponse: """Say hi to the user.""" username = ctx.context.get('username') if ctx is not None else 'unknown' - return await ai.generate( - on_chunk=ctx.send_chunk if ctx is not None else None, + stream_response = ai.generate_stream( prompt=f'tell a medium sized joke about {input.name} for user {username}', ) + async for chunk in stream_response.stream: + if ctx is not None and chunk.text: + ctx.send_chunk(chunk.text) + return await stream_response.response diff --git a/py/tests/conform/README.md b/py/tests/conform/README.md index 3d433470fb..6806634bc6 100644 --- a/py/tests/conform/README.md +++ b/py/tests/conform/README.md @@ -41,9 +41,6 @@ py/bin/conform check-plugin ``` tests/conform/ β”œβ”€β”€ README.md ← you are here -β”œβ”€β”€ amazon-bedrock/ -β”‚ β”œβ”€β”€ conformance_entry.py ← minimal Genkit + plugin init -β”‚ └── model-conformance.yaml ← models + supported capabilities β”œβ”€β”€ anthropic/ β”‚ β”œβ”€β”€ conformance_entry.py β”‚ └── model-conformance.yaml @@ -93,14 +90,13 @@ tests/conform/ | `deepseek` | `DEEPSEEK_API_KEY` | | `xai` | `XAI_API_KEY` | | `cohere` | `COHERE_API_KEY` | -| `amazon-bedrock` | `AWS_REGION` + AWS credentials | | `huggingface` | `HF_TOKEN` | | `microsoft-foundry` | `AZURE_OPENAI_API_KEY`, `AZURE_OPENAI_ENDPOINT` | | `cloudflare-workers-ai` | `CLOUDFLARE_ACCOUNT_ID`, `CLOUDFLARE_API_TOKEN` | | `vertex-ai` | `GOOGLE_CLOUD_PROJECT` + GCP credentials | | `ollama` | *(none β€” uses local server)* | -## Model coverage (26 models across 13 plugins) +## Model coverage (25 models across 12 plugins) | Plugin | Models | Key capabilities tested | |--------|--------|------------------------| @@ -111,7 +107,6 @@ tests/conform/ | **deepseek** | deepseek-chat, deepseek-reasoner | tool-request (chat only), structured-output, streaming | | **xai** | grok-4-fast-non-reasoning, grok-2-vision-1212 | tool-request, structured-output, vision, streaming | | **cohere** | command-a-03-2025 | tool-request, structured-output, multiturn | -| **amazon-bedrock** | us.anthropic.claude-sonnet-4-5-...v1:0 | tool-request, structured-output, vision, streaming | | **microsoft-foundry** | gpt-4o | tool-request, structured-output, vision, streaming | | **huggingface** | meta-llama/Llama-3.1-8B-Instruct | multiturn, system-role | | **cloudflare-workers-ai** | @cf/meta/llama-3.1-8b-instruct | tool-request, multiturn, streaming | diff --git a/py/tests/conform/amazon-bedrock/conformance_entry.py b/py/tests/conform/amazon-bedrock/conformance_entry.py deleted file mode 100644 index 56ed468a14..0000000000 --- a/py/tests/conform/amazon-bedrock/conformance_entry.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Minimal entry point for amazon-bedrock model conformance testing. - -Usage: - genkit dev:test-model --from-file model-conformance.yaml -- uv run conformance_entry.py - -Env: - AWS_REGION: Required. AWS region (e.g., us-east-1). - AWS credentials must be configured via environment or ~/.aws/credentials. -""" - -import asyncio - -from genkit.ai import Genkit -from genkit.plugins.amazon_bedrock import AmazonBedrock - -ai = Genkit(plugins=[AmazonBedrock()]) - - -async def main() -> None: - """Keep the process alive for the test runner.""" - await asyncio.Event().wait() - - -if __name__ == '__main__': - ai.run_main(main()) diff --git a/py/tests/conform/amazon-bedrock/model-conformance.yaml b/py/tests/conform/amazon-bedrock/model-conformance.yaml deleted file mode 100644 index ec3ac78713..0000000000 --- a/py/tests/conform/amazon-bedrock/model-conformance.yaml +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -# Amazon Bedrock model conformance test spec. -# -# Uses Claude Sonnet 4.5 via Bedrock cross-region inference profile. -# Claude Sonnet 4.5 supports structured output (output=['text','json'], -# constrained=ALL) per Anthropic docs. - -- model: amazon-bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 - supports: - - tool-request - - structured-output - - multiturn - - system-role - - input-image-base64 - - streaming-multiturn - - streaming-tool-request - - streaming-structured-output diff --git a/py/tests/conform/anthropic/conformance_entry.py b/py/tests/conform/anthropic/conformance_entry.py index a6f4b188b5..f90fa71bdd 100644 --- a/py/tests/conform/anthropic/conformance_entry.py +++ b/py/tests/conform/anthropic/conformance_entry.py @@ -25,7 +25,7 @@ import asyncio -from genkit.ai import Genkit +from genkit import Genkit from genkit.plugins.anthropic import Anthropic ai = Genkit(plugins=[Anthropic()]) diff --git a/py/tests/conform/compat-oai/conformance_entry.py b/py/tests/conform/compat-oai/conformance_entry.py index a3f98dfc30..1188b7402f 100644 --- a/py/tests/conform/compat-oai/conformance_entry.py +++ b/py/tests/conform/compat-oai/conformance_entry.py @@ -25,7 +25,7 @@ import asyncio -from genkit.ai import Genkit +from genkit import Genkit from genkit.plugins.compat_oai import OpenAI ai = Genkit(plugins=[OpenAI()]) diff --git a/py/tests/conform/conform.toml b/py/tests/conform/conform.toml index 18e088893d..ba1b88fd22 100644 --- a/py/tests/conform/conform.toml +++ b/py/tests/conform/conform.toml @@ -49,7 +49,6 @@ additional-model-plugins = ["google-genai", "vertex-ai", "ollama"] # Required environment variables per plugin. An empty list means no # credentials are needed (e.g. ollama uses a local server). [conform.env] -amazon-bedrock = ["AWS_REGION"] anthropic = ["ANTHROPIC_API_KEY"] cloudflare-workers-ai = ["CLOUDFLARE_ACCOUNT_ID", "CLOUDFLARE_API_TOKEN"] cohere = ["COHERE_API_KEY"] diff --git a/py/tests/conform/google-genai/conformance_entry.py b/py/tests/conform/google-genai/conformance_entry.py index 2ce05a0260..62f0020e69 100644 --- a/py/tests/conform/google-genai/conformance_entry.py +++ b/py/tests/conform/google-genai/conformance_entry.py @@ -25,7 +25,7 @@ import asyncio -from genkit.ai import Genkit +from genkit import Genkit from genkit.plugins.google_genai import GoogleAI ai = Genkit(plugins=[GoogleAI()]) diff --git a/py/tests/conform/ollama/conformance_entry.py b/py/tests/conform/ollama/conformance_entry.py index 78277fad1c..c02f338801 100644 --- a/py/tests/conform/ollama/conformance_entry.py +++ b/py/tests/conform/ollama/conformance_entry.py @@ -28,7 +28,7 @@ import asyncio -from genkit.ai import Genkit +from genkit import Genkit from genkit.plugins.ollama import Ollama from genkit.plugins.ollama.models import ModelDefinition diff --git a/py/tests/conform/test-model-conformance b/py/tests/conform/test-model-conformance index 81db12bea9..819bae66ba 100755 --- a/py/tests/conform/test-model-conformance +++ b/py/tests/conform/test-model-conformance @@ -51,7 +51,6 @@ get_env_vars() { deepseek) echo "DEEPSEEK_API_KEY" ;; xai) echo "XAI_API_KEY" ;; cohere) echo "COHERE_API_KEY" ;; - amazon-bedrock) echo "AWS_REGION" ;; huggingface) echo "HF_TOKEN" ;; microsoft-foundry) echo "AZURE_OPENAI_API_KEY AZURE_OPENAI_ENDPOINT" ;; cloudflare-workers-ai) echo "CLOUDFLARE_ACCOUNT_ID CLOUDFLARE_API_TOKEN" ;; diff --git a/py/tests/conform/vertex-ai/conformance_entry.py b/py/tests/conform/vertex-ai/conformance_entry.py index 70ce09250a..805b0cfbda 100644 --- a/py/tests/conform/vertex-ai/conformance_entry.py +++ b/py/tests/conform/vertex-ai/conformance_entry.py @@ -26,7 +26,7 @@ import asyncio -from genkit.ai import Genkit +from genkit import Genkit from genkit.plugins.vertex_ai import ModelGardenPlugin ai = Genkit(plugins=[ModelGardenPlugin()]) diff --git a/py/tests/smoke/package_test.py b/py/tests/smoke/package_test.py index 7d354b3fc4..fe53002ca7 100644 --- a/py/tests/smoke/package_test.py +++ b/py/tests/smoke/package_test.py @@ -16,42 +16,19 @@ """Smoke tests for package structure.""" -from genkit.core import package_name as core_package_name -from genkit.plugins.firebase import package_name as firebase_package_name from genkit.plugins.google_cloud import package_name as google_cloud_package_name from genkit.plugins.google_genai import package_name as google_genai_package_name from genkit.plugins.ollama import package_name as ollama_package_name from genkit.plugins.vertex_ai import package_name as vertex_ai_package_name -def square(n: int | float) -> int | float: - """Calculates the square of a number. - - Args: - n: The number to square. - - Returns: - The square of n. - """ - return n * n - - def test_package_names() -> None: """A test that ensure that the package imports work correctly. This test verifies that the package imports work correctly from the end-user perspective. """ - assert core_package_name() == 'genkit.core' - assert firebase_package_name() == 'genkit.plugins.firebase' assert google_cloud_package_name() == 'genkit.plugins.google_cloud' assert google_genai_package_name() == 'genkit.plugins.google_genai' assert ollama_package_name() == 'genkit.plugins.ollama' assert vertex_ai_package_name() == 'genkit.plugins.vertex_ai' - - -def test_square() -> None: - """Tests whether the square function works correctly.""" - assert square(2) == 4 - assert square(3) == 9 - assert square(4) == 16 diff --git a/py/tools/conform/docs/configuration.md b/py/tools/conform/docs/configuration.md index 1a97a47c8f..19519e29c9 100644 --- a/py/tools/conform/docs/configuration.md +++ b/py/tools/conform/docs/configuration.md @@ -119,7 +119,6 @@ mistral = ["MISTRAL_API_KEY"] deepseek = ["DEEPSEEK_API_KEY"] xai = ["XAI_API_KEY"] cohere = ["COHERE_API_KEY"] -amazon-bedrock = ["AWS_REGION"] huggingface = ["HF_TOKEN"] microsoft-foundry = ["AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT"] cloudflare-workers-ai = ["CLOUDFLARE_ACCOUNT_ID", "CLOUDFLARE_API_TOKEN"] diff --git a/py/tools/conform/src/conform/executors/conformance_native.py b/py/tools/conform/src/conform/executors/conformance_native.py index 35cb681275..beaae84911 100644 --- a/py/tools/conform/src/conform/executors/conformance_native.py +++ b/py/tools/conform/src/conform/executors/conformance_native.py @@ -44,8 +44,7 @@ import structlog -from genkit.ai import Genkit -from genkit.codec import dump_dict +from genkit import Genkit # --------------------------------------------------------------------------- # Plugin registry β€” maps conform plugin names to Genkit init functions. @@ -91,13 +90,6 @@ def _init_ollama() -> Genkit: ) -def _init_amazon_bedrock() -> Genkit: - """Initialize the Amazon Bedrock plugin.""" - from genkit.plugins.amazon_bedrock import AmazonBedrock - - return Genkit(plugins=[AmazonBedrock()]) - - def _init_compat_oai() -> Genkit: """Initialize the OpenAI compat plugin.""" from genkit.plugins.compat_oai import OpenAI @@ -110,7 +102,6 @@ def _init_compat_oai() -> Genkit: 'vertex-ai': _init_vertex_ai, 'anthropic': _init_anthropic, 'ollama': _init_ollama, - 'amazon-bedrock': _init_amazon_bedrock, 'compat-oai': _init_compat_oai, } @@ -125,7 +116,7 @@ def _register_ephemeral_tool( tool_def: dict[str, Any], ) -> None: """Register a no-op tool so generate() can resolve its definition.""" - from genkit.core.action.types import ActionKind + from genkit._core._action import ActionKind registry = ai.registry with registry._lock: # pyright: ignore[reportPrivateUsage] @@ -175,7 +166,8 @@ async def handle_request( config = input_data.get('config') # Convert raw message dicts to Message objects. - from genkit.core.typing import Message, OutputConfig, Part + from genkit._core._typing import OutputConfig, Part + from genkit.model import Message msg_objects = [Message.model_validate(m) for m in messages] @@ -202,21 +194,29 @@ async def handle_request( tool_names.append(t_name) _register_ephemeral_tool(ai, t_name, tdef) + # Expand OutputConfig to individual params for generate/generate_stream + output_kwargs: dict[str, Any] = {} + if output_obj: + output_kwargs['output_format'] = output_obj.format + output_kwargs['output_schema'] = output_obj.schema + output_kwargs['output_content_type'] = output_obj.content_type + output_kwargs['output_constrained'] = output_obj.constrained + chunks: list[dict[str, Any]] = [] if stream: - stream_iter, response_future = ai.generate_stream( + stream_response = ai.generate_stream( model=model_name, system=system_parts, messages=non_system_messages, tools=tool_names, config=config, - output=output_obj, return_tool_requests=True, + **output_kwargs, ) - async for chunk in stream_iter: - chunks.append(cast(dict[str, Any], dump_dict(chunk))) - result = await response_future + async for chunk in stream_response.stream: + chunks.append(chunk.model_dump()) + result = await stream_response.response else: result = await ai.generate( model=model_name, @@ -224,11 +224,11 @@ async def handle_request( messages=non_system_messages, tools=tool_names, config=config, - output=output_obj, return_tool_requests=True, + **output_kwargs, ) - response = cast(dict[str, Any], dump_dict(result)) + response = result.model_dump() return {'response': response, 'chunks': chunks} except Exception: @@ -247,7 +247,7 @@ async def handle_request( async def main() -> None: """Run the JSONL-over-stdio loop.""" # CRITICAL: Redirect all logging to stderr BEFORE plugin initialization. - # Plugins use structlog (via genkit.core.logging) which defaults to stdout. + # Plugins use structlog (via genkit._core.logging) which defaults to stdout. # Since this executor uses JSONL-over-stdio, ANY non-JSON line on stdout # causes the conform runner to fail with "invalid JSON". logging.basicConfig( diff --git a/py/tools/conform/src/conform/executors/in_process_runner.py b/py/tools/conform/src/conform/executors/in_process_runner.py index 916ba3a0ef..ec5e71a709 100644 --- a/py/tools/conform/src/conform/executors/in_process_runner.py +++ b/py/tools/conform/src/conform/executors/in_process_runner.py @@ -39,8 +39,6 @@ from rich.console import Console -from genkit.codec import dump_dict - console = Console(stderr=True) @@ -100,7 +98,7 @@ async def run_action( ) -> tuple[dict[str, Any], list[dict[str, Any]]]: """Run a model action via ``ai.generate()`` in-process. - Converts the raw test input (GenerateRequest-shaped dict) into + Converts the raw test input (ModelRequest-shaped dict) into ``ai.generate()`` keyword arguments so the full framework pipeline is exercised, including output format handling (``extract_json`` for JSON output) and streaming chunk wrapping. @@ -129,14 +127,15 @@ async def _run_action_impl( # e.g. "/model/googleai/gemini-2.5-flash" -> "googleai/gemini-2.5-flash" model_name = key.removeprefix('/model/') - # Extract fields from the GenerateRequest-shaped input. + # Extract fields from the ModelRequest-shaped input. messages = input_data.get('messages', []) output_config = input_data.get('output') tools_defs = input_data.get('tools') config = input_data.get('config') # Convert raw message dicts to Message objects. - from genkit.core.typing import Message, OutputConfig, Part + from genkit._core._typing import OutputConfig, Part + from genkit.model import Message msg_objects = [Message.model_validate(m) for m in messages] @@ -174,7 +173,7 @@ async def _run_action_impl( if stream: # Use generate_stream for streaming tests. - stream_iter, response_future = self._ai.generate_stream( + stream_response = self._ai.generate_stream( model=model_name, system=system_parts, messages=non_system_messages, @@ -184,9 +183,9 @@ async def _run_action_impl( return_tool_requests=True, ) # Consume the stream to collect chunks. - async for chunk in stream_iter: - chunks.append(cast(dict[str, Any], dump_dict(chunk))) - result = await response_future + async for chunk in stream_response.stream: + chunks.append(cast(dict[str, Any], chunk.model_dump())) + result = await stream_response.response else: result = await self._ai.generate( model=model_name, @@ -198,7 +197,7 @@ async def _run_action_impl( return_tool_requests=True, ) - response = cast(dict[str, Any], dump_dict(result)) + response = cast(dict[str, Any], result.model_dump()) return response, chunks async def close(self) -> None: @@ -218,7 +217,7 @@ def _register_ephemeral_tool( and ``description`` so ``to_tool_definition()`` can build the right ``ToolDefinition`` for the model. """ - from genkit.core.action.types import ActionKind + from genkit._core._action import ActionKind registry = ai.registry diff --git a/py/tools/conform/src/conform/reflection.py b/py/tools/conform/src/conform/reflection.py index 7b1d496519..2dae7a154e 100644 --- a/py/tools/conform/src/conform/reflection.py +++ b/py/tools/conform/src/conform/reflection.py @@ -101,7 +101,7 @@ async def run_action( Args: key: The action key (e.g. ``/model/googleai/gemini-2.5-flash``). - input_data: The ``GenerateRequest`` payload. + input_data: The ``ModelRequest`` payload. stream: Whether to request streaming. Returns: diff --git a/py/tools/conform/src/conform/util_test_model.py b/py/tools/conform/src/conform/util_test_model.py index 0928178b48..63584d1074 100644 --- a/py/tools/conform/src/conform/util_test_model.py +++ b/py/tools/conform/src/conform/util_test_model.py @@ -238,7 +238,7 @@ class NativeRunner: Starts the native executor once per plugin. The executor initializes the plugin (``genkit.Init`` / ``genkit(...)``), then enters a read loop: - 1. Reads one JSON line from **stdin** (a ``GenerateRequest``). + 1. Reads one JSON line from **stdin** (a ``ModelRequest``). 2. Calls ``generate()`` natively using the SDK. 3. Writes one JSON line to **stdout** (a ``GenerateResponse``). diff --git a/py/tools/model-config-test/model_performance_test.py b/py/tools/model-config-test/model_performance_test.py index 72ef70ab75..7b3260b498 100644 --- a/py/tools/model-config-test/model_performance_test.py +++ b/py/tools/model-config-test/model_performance_test.py @@ -54,7 +54,6 @@ async def discover_models() -> dict[str, Any]: ('genkit.plugins.vertex_ai', 'VertexAI'), ('genkit.plugins.anthropic', 'Anthropic'), ('genkit.plugins.ollama', 'Ollama'), - ('genkit.plugins.amazon_bedrock', 'AmazonBedrock'), ] for module_path, class_name in plugin_imports: @@ -73,7 +72,7 @@ async def discover_models() -> dict[str, Any]: registry = ai.registry # Get all model actions via list_actions (which queries plugins) - from genkit.core.action import ActionKind + from genkit._core._action import ActionKind try: actions = await registry.list_actions(allowed_kinds=[ActionKind.MODEL]) @@ -216,7 +215,7 @@ async def discover_models_for_sample(sample_name: str) -> dict[str, Any]: registry = ai.registry # Get all model actions - from genkit.core.action import ActionKind + from genkit._core._action import ActionKind try: actions = await registry.list_actions(allowed_kinds=[ActionKind.MODEL]) diff --git a/py/tools/model-config-test/pyproject.toml b/py/tools/model-config-test/pyproject.toml index 56357999e0..1d4fa9f061 100644 --- a/py/tools/model-config-test/pyproject.toml +++ b/py/tools/model-config-test/pyproject.toml @@ -8,7 +8,6 @@ dependencies = [ "genkit-plugin-google-genai", "genkit-plugin-vertex-ai", "genkit-plugin-google-cloud", - "genkit-plugin-amazon-bedrock", "genkit-plugin-anthropic", "genkit-plugin-ollama", "fastapi", @@ -47,7 +46,6 @@ packages = ["model_performance_test.py", "run_single_model_test.py"] [tool.uv.sources] genkit = { workspace = true } -genkit-plugin-amazon-bedrock = { workspace = true } genkit-plugin-anthropic = { workspace = true } genkit-plugin-google-cloud = { workspace = true } genkit-plugin-google-genai = { workspace = true } diff --git a/py/tools/model-config-test/run_single_model_test.py b/py/tools/model-config-test/run_single_model_test.py index 999a3b57fe..17f5605065 100644 --- a/py/tools/model-config-test/run_single_model_test.py +++ b/py/tools/model-config-test/run_single_model_test.py @@ -59,7 +59,8 @@ async def run_model_test( try: from genkit import Genkit - from genkit.core.typing import Message, TextPart + from genkit._core._typing import TextPart + from genkit.model import Message plugins = [] try: @@ -86,13 +87,6 @@ async def run_model_test( plugins.append(Ollama()) except Exception: # noqa: S110 pass - try: - from genkit.plugins.amazon_bedrock import AmazonBedrock - - plugins.append(AmazonBedrock()) - except Exception: # noqa: S110 - pass - # Initialize Genkit ai = Genkit(plugins=plugins) diff --git a/py/tools/sample-flows/LICENSE b/py/tools/sample-flows/LICENSE deleted file mode 100644 index 2205396735..0000000000 --- a/py/tools/sample-flows/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2025 Google LLC - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/py/tools/sample-flows/README.md b/py/tools/sample-flows/README.md deleted file mode 100644 index 48b75654f0..0000000000 --- a/py/tools/sample-flows/README.md +++ /dev/null @@ -1,30 +0,0 @@ -# Sample Flow Testing Tool - -This directory contains scripts for reviewing and testing Genkit flows within samples. - -## Files - -- `review_sample_flows.py`: Iterates through all flows in a given sample directory, runs them with heuristic inputs, and generates a report. -- `run_single_flow.py`: Helper script to run a single flow in isolation (used by `review_sample_flows.py`). - -## Usage - -This tool is typically run via the `py/bin/test_sample_flows` script: - -```bash -# Test a specific sample -py/bin/test_sample_flows provider-google-genai-hello -``` - -Or manually: - -```bash -# Run from the repository root (py/) -uv run tools/sample-flows/review_sample_flows.py samples/provider-google-genai-hello -``` - -## Output - -The tool generates a text report (e.g., `flow_review_results.txt`) detailing which flows passed or failed, along with their outputs or error messages. - -**Note:** The `test_sample_flows` script automatically skips samples that do not have a standard `main.py` entry point (e.g., `framework-evaluator-demo`), preventing execution errors. diff --git a/py/tools/sample-flows/pyproject.toml b/py/tools/sample-flows/pyproject.toml deleted file mode 100644 index ca244e9135..0000000000 --- a/py/tools/sample-flows/pyproject.toml +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -[project] -authors = [{ name = "Google" }] -classifiers = ["Private :: Do Not Upload"] -dependencies = ["genkit", "httpx", "structlog"] -description = "Tool to run flow tests for Genkit samples" -license = "Apache-2.0" -name = "genkit-tools-sample-flows" -readme = "README.md" -requires-python = ">=3.10" -version = "0.1.1" - -[build-system] -build-backend = "hatchling.build" -requires = ["hatchling"] - -[tool.hatch.build.targets.wheel] -packages = ["review_sample_flows.py", "run_single_flow.py"] - -[tool.uv.sources] -genkit = { workspace = true } diff --git a/py/tools/sample-flows/review_sample_flows.py b/py/tools/sample-flows/review_sample_flows.py deleted file mode 100644 index 91d11b92b2..0000000000 --- a/py/tools/sample-flows/review_sample_flows.py +++ /dev/null @@ -1,467 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# pyrefly: ignore-file - -"""Tool to review and test all Genkit flows in a sample's main.py. - -Usage: - python review_sample_flows.py - -Example: - python review_sample_flows.py samples/provider-google-genai-hello -""" - -import argparse -import asyncio -import importlib.util -import json -import logging -import platform -import re -import subprocess # noqa: S404 -import sys -import time -import traceback -import warnings -from pathlib import Path -from typing import Any - -from genkit.core.action import ActionKind -from genkit.types import Media - -logging.getLogger().setLevel(logging.ERROR) -logging.getLogger('asyncio').setLevel(logging.ERROR) -logging.getLogger('httpx').setLevel(logging.ERROR) -logging.getLogger('httpcore').setLevel(logging.ERROR) - -warnings.filterwarnings('ignore') - - -def open_file(path: str) -> None: - """Open a file with the default system application.""" - try: - if platform.system() == 'Darwin': # macOS - subprocess.run(['open', path], check=False) # noqa: S603, S607 - elif platform.system() == 'Linux': - subprocess.run(['xdg-open', path], check=False) # noqa: S603, S607 - elif platform.system() == 'Windows': - subprocess.run(['start', path], shell=True, check=False) # noqa: S602, S607 - except Exception: # noqa: S110 - pass - - -def write_report( - path: str, - action_count: int, - successful_flows: list[str], - failed_flows: list[str], - detail_lines: list[str], - sample_name: str, -) -> None: - """Write the test report to a file.""" - report_lines = [] - report_lines.append(f'Flow Review Report for {sample_name}') - report_lines.append('=' * 60) - report_lines.append('') - - report_lines.append('SUMMARY') - report_lines.append('=' * 60) - report_lines.append(f'Total Flows: {action_count}') - report_lines.append(f'Successful: {len(successful_flows)}') - report_lines.append(f'Failed: {len(failed_flows)}') - report_lines.append('') - - if failed_flows: - report_lines.append('Failed Flows:') - for flow in failed_flows: - report_lines.append(f' βœ— {flow}') - report_lines.append('') - - if successful_flows: - report_lines.append('Successful Flows:') - for flow in successful_flows: - report_lines.append(f' βœ“ {flow}') - report_lines.append('') - - report_lines.append('=' * 60) - report_lines.append('') - report_lines.append('DETAILED RESULTS') - report_lines.append('=' * 60) - report_lines.append('') - - # Append detailed results - report_lines.extend(detail_lines) - - # Write report - with open(path, 'w') as f: - f.write('\n'.join(report_lines)) - - -def main() -> None: # noqa: ASYNC240, ASYNC230 - test script, blocking I/O acceptable - """Test all flows in a Genkit sample and generate a report.""" - parser = argparse.ArgumentParser(description='Test all flows in a Genkit sample.') - parser.add_argument('sample_dir', type=str, help='Path to the sample directory') - parser.add_argument('--output', type=str, default='flow_review_results.txt', help='Output report file') - args = parser.parse_args() - - # Suppress verbose logging from genkit framework to avoid printing full data URLs - logging.basicConfig(level=logging.WARNING) - logging.getLogger('genkit').setLevel(logging.WARNING) - logging.getLogger('google').setLevel(logging.WARNING) - - sample_path = Path(args.sample_dir).resolve() - if not sample_path.exists(): - sys.exit(1) - - # Assume the main entry point is at src/main.py or main.py - is_package_structured = False - main_py_path = sample_path / 'src' / 'main.py' - if not main_py_path.exists(): - main_py_path = sample_path / 'main.py' - if not main_py_path.exists(): - sys.exit(1) - else: - is_package_structured = True - - # Add the source directory to sys.path so imports work - if is_package_structured: - # For src/main.py, we want top-level imports to be relative to the sample dir - # and relative imports inside src to work via 'src.main' package context. - sys.path.insert(0, str(sample_path)) - module_name = 'src.main' - else: - src_dir = main_py_path.parent - sys.path.insert(0, str(src_dir)) - module_name = 'sample_main' - - # Add the py/ root directory to sys.path so 'samples.shared' imports work - # sample_path is .../py/samples/sample-name - # sample_path.parent is .../py/samples - # sample_path.parent.parent is .../py - sys.path.insert(0, str(sample_path.parent.parent)) - - # Clear existing modules with the same name if they exist. - # This happens in a monorepo because multiple samples might use 'src.main' - # as their entry point, or have a 'src' package. - for m in list(sys.modules.keys()): - if m == module_name or m.startswith(module_name + '.') or m == 'src' or m.startswith('src.'): - sys.modules.pop(m, None) - - # Import the module - try: - if is_package_structured: - # Use import_module for package-structured samples to correctly handle relative imports - module = importlib.import_module(module_name) - else: - # Fallback to spec-based loading for simple samples - spec = importlib.util.spec_from_file_location(module_name, main_py_path) - if spec is None or spec.loader is None: - sys.exit(1) - # Add explicit assertions for type narrowing if needed by some checkers - assert spec is not None - assert spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception: - traceback.print_exc() - sys.exit(1) - - # Find the Genkit instance - ai_instance = None - for attr_name in dir(module): - attr = getattr(module, attr_name) - # Check if it looks like a Genkit instance (has registry) - if hasattr(attr, 'registry') and hasattr(attr, 'generate'): - ai_instance = attr - break - - if ai_instance: - # Manually load prompts from the sample's prompts directory if it exists. - # This is needed because the current working directory isn't the sample dir. - prompts_dir = sample_path / 'prompts' - if prompts_dir.exists() and prompts_dir.is_dir(): - from genkit.blocks.prompt import load_prompt_folder - - load_prompt_folder(ai_instance.registry, prompts_dir) - - if ai_instance is None: - sys.exit(1) # pyrefly: ignore[unbound-name] - sys is imported at top of file - - assert ai_instance is not None # Type narrowing for ai_instance.registry - - # List all flows - registry = ai_instance.registry - actions_map = asyncio.run(registry.resolve_actions_by_kind(ActionKind.FLOW)) - - # Track results for summary - successful_flows = [] - failed_flows = [] - - # We'll add the summary after testing all flows - class LiveLogger(list): - def append(self, item: Any) -> None: # noqa: ANN401 - override requires Any - super().append(item) - # Print to stdout for immediate feedback - print(item) # noqa: T201 - - detail_lines = LiveLogger() - - try: - for flow_name, flow_action in actions_map.items(): - detail_lines.append(f'\nFlow: {flow_name}') - detail_lines.append('-' * 30) - - try: - input_data = generate_input(flow_action) - detail_lines.append(f'Generated Input: {input_data}') - - # Run flow in subprocess to avoid event loop conflicts - # Get path to helper script - script_dir = Path(__file__).parent - helper_script = script_dir / 'run_single_flow.py' - - # Prepare subprocess command - cmd = [ - 'uv', - 'run', - str(helper_script), - str(sample_path), - flow_name, - '--input', - json.dumps(input_data) if input_data is not None else 'null', - ] - - # Run subprocess - process = subprocess.Popen( # noqa: S603 - cmd is constructed internally from trusted script paths - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - bufsize=1, # Line buffered - cwd=sample_path.parent.parent, # Run from py/ directory - ) - - # Stream output - stdout_lines = [] - if process.stdout: - for line in process.stdout: - stdout_lines.append(line) - # Suppress debug logs as requested, but stream everything else for "live logging" - # We strip ANSI codes before checking per report/pic requirement - clean_line = re.sub(r'\x1b\[[0-9;]*[mGKF]', '', line).lower() - # Level tags like [debug], [debug ], etc. - is_debug = '[debug' in clean_line - # List of noisy patterns to filter from test output. - noise_patterns = [ - 'userwarning:', - 'shadows an attribute', - 'class outputconfig', - 'class promptinputconfig', - 'class promptoutputconfig', - '---json_result_', - '{"success":', - '[info', - ] - is_noise = any(p in clean_line for p in noise_patterns) - - if not is_debug and not is_noise: - # Also filter out empty lines or just markers - if line.strip(): - print(f' {line}', end='') # noqa: T201 - - process.wait(timeout=120) - - # Reconstruct stdout for parsing - stdout = ''.join(stdout_lines) - - try: - if '---JSON_RESULT_START---' in stdout and '---JSON_RESULT_END---' in stdout: - json_str = stdout.split('---JSON_RESULT_START---')[1].split('---JSON_RESULT_END---')[0].strip() - result_data = json.loads(json_str) - else: - # Fallback try to parse the whole thing if markers missing (unlikely for success case) - result_data = json.loads(stdout) - except (json.JSONDecodeError, IndexError): - detail_lines.append('Status: FAILED') - detail_lines.append('Error: Failed to parse subprocess output') - - # Print raw output for debugging since parsing failed - detail_lines.append('Raw Output:') - detail_lines.append(stdout) - - failed_flows.append(flow_name) - continue - - if result_data.get('success'): - detail_lines.append('Status: SUCCESS') - - # Format the result - flow_result = result_data.get('result') - formatted_output = format_output(flow_result, max_length=500) - detail_lines.append(f'Output: {formatted_output}') - - successful_flows.append(flow_name) - else: - detail_lines.append('Status: FAILED') - error_msg = result_data.get('error', 'Unknown error') - detail_lines.append(f'Error: {error_msg}') - failed_flows.append(flow_name) - - except subprocess.TimeoutExpired: - detail_lines.append('Status: FAILED') - detail_lines.append('Error: Flow execution timed out (120s)') - failed_flows.append(flow_name) - except Exception as e: - detail_lines.append('Status: FAILED') - detail_lines.append(f'Error: Subprocess failed: {e}') - failed_flows.append(flow_name) - - # Add traceback for debugging - tb_lines = traceback.format_exc().split('\n') - detail_lines.append('Traceback:') - for line in tb_lines: - detail_lines.append(f' {line}') - - failed_flows.append(flow_name) - - detail_lines.append('') - - # Add a small delay between tests to avoid rate limiting - time.sleep(10) - - except KeyboardInterrupt: - pass - finally: - write_report( - args.output, - len(actions_map), - successful_flows, - failed_flows, - detail_lines, - sample_path.name, - ) - open_file(args.output) - - -def format_output(output: Any, max_length: int = 500) -> str: # noqa: ANN401 - intentional use of Any for arbitrary flow outputs - """Format flow output for human-readable display. - - Args: - output: The flow output to format - max_length: Maximum length before truncation - - Returns: - Formatted string representation - """ - # Handle None - if output is None: - return 'None' - - # Handle Media objects - if isinstance(output, Media): - if output.url and len(output.url) > max_length: - truncated_url = f'{output.url[:100]}...{output.url[-50:]}' - return f"Media(url='{truncated_url}' [{len(output.url)} chars], content_type='{output.content_type}')" - return f"Media(url='{output.url}', content_type='{output.content_type}')" - - # Handle Pydantic models - if hasattr(output, 'model_dump'): - try: - data = output.model_dump() - json_str = json.dumps(data, indent=2) - if len(json_str) > max_length: - return f'{json_str[:max_length]}... [truncated, {len(json_str)} total chars]' - return json_str - except Exception: # noqa: S110 - intentional fallback if model_dump fails - pass - - # Handle dicts - if isinstance(output, dict): - try: - json_str = json.dumps(output, indent=2) - if len(json_str) > max_length: - return f'{json_str[:max_length]}... [truncated, {len(json_str)} total chars]' - return json_str - except Exception: # noqa: S110 - intentional fallback if json.dumps fails - pass - - # Handle lists - if isinstance(output, list): - try: - json_str = json.dumps(output, indent=2) - if len(json_str) > max_length: - return f'{json_str[:max_length]}... [truncated, {len(json_str)} total chars]' - return json_str - except Exception: # noqa: S110 - intentional fallback if json.dumps fails - pass - - # Default: convert to string - output_str = str(output) - if len(output_str) > max_length: - return f'{output_str[:max_length]}... [truncated, {len(output_str)} total chars]' - return output_str - - -def generate_input(flow_action: Any) -> Any: # noqa: ANN401 - intentional use of Any for arbitrary flow inputs - """Generates heuristic input for a flow based on its schema.""" - schema = flow_action.input_schema - if not schema: - return None - - # Generate dict from schema - input_dict = generate_from_json_schema(schema) - - # If the flow has a Pydantic model for input, instantiate it - # The schema has a 'title' field that matches the model class name - if isinstance(input_dict, dict) and 'title' in schema: - # Try to get the actual model class from the flow's metadata - # For now, just return the dict - Genkit should handle conversion - pass - - return input_dict - - -def generate_from_json_schema(schema: dict[str, Any]) -> Any: # noqa: ANN401 - intentional use of Any for arbitrary flow inputs - """Simplistic JSON schema input generator.""" - type_ = schema.get('type') - - if 'default' in schema: - return schema['default'] - - if type_ == 'string': - return 'test_string' - elif type_ == 'integer': - return 42 - elif type_ == 'number': - return 3.14 - elif type_ == 'boolean': - return True - elif type_ == 'object': - properties = schema.get('properties', {}) - result = {} - for prop_name, prop_schema in properties.items(): - result[prop_name] = generate_from_json_schema(prop_schema) - return result - elif type_ == 'array': - items_schema = schema.get('items', {}) - return [generate_from_json_schema(items_schema)] - - return None - - -if __name__ == '__main__': - main() diff --git a/py/tools/sample-flows/run_single_flow.py b/py/tools/sample-flows/run_single_flow.py deleted file mode 100644 index 53311ea1ef..0000000000 --- a/py/tools/sample-flows/run_single_flow.py +++ /dev/null @@ -1,280 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -# pyrefly: ignore-file -# flake8: noqa: ASYNC240 - -"""Helper script to run a single Genkit flow in isolation. - -This script is called by review_sample_flows.py to execute each flow in a -separate subprocess, avoiding event loop conflicts and enabling proper -async execution. - -Usage: - python run_single_flow.py [--input ] - -Output: - JSON object with 'success', 'result', and 'error' fields -""" - -import argparse -import asyncio -import importlib.util -import json -import logging -import sys -import traceback -from pathlib import Path -from typing import Any - -from genkit.core.action import ActionKind -from genkit.types import Media - - -def format_output( - output: Any, # noqa: ANN401 - visited: set[int] | None = None, # noqa: ANN401 -) -> Any: # noqa: ANN401 - """Format flow output for serialization. - - Args: - output: The flow output to format - visited: Set of object IDs already visited to prevent infinite recursion - - Returns: - Serializable representation - """ - if visited is None: - visited = set() - - # Track visited objects to prevent infinite recursion - # Only track mutable/container types that could be involved in cycles - if isinstance(output, (dict, list)) or hasattr(output, 'model_dump'): - if id(output) in visited: - return '(Recursive Reference)' - visited.add(id(output)) - - # Handle None - if output is None: - return None - - # Handle Media objects - if isinstance(output, Media): - return { - 'type': 'Media', - 'url': '(Media data not shown)', - 'content_type': output.content_type, - } - - # Handle Pydantic models - if hasattr(output, 'model_dump'): - try: - return format_output(output.model_dump(), visited) - except Exception: # noqa: S110 - intentional fallback if model_dump fails - pass - - # Handle dicts - if isinstance(output, dict): - return {k: format_output(v, visited) for k, v in output.items()} - - # Handle lists - if isinstance(output, list): - return [format_output(v, visited) for v in output] - - # Default: convert to string for non-serializable objects (except basics) - if isinstance(output, (str, int, float, bool, type(None))): - return output - - return str(output) - - -async def run_flow(sample_dir: str, flow_name: str, input_data: Any) -> dict[str, Any]: # noqa: ANN401 - intentional use of Any for arbitrary flow outputs - """Run a single flow and return result. - - Args: - sample_dir: Path to sample directory - flow_name: Name of flow to run - input_data: Input data for the flow - - Returns: - Dict with 'success', 'result', 'error' fields - """ - result: dict[str, Any] = { - 'success': False, - 'result': None, - 'error': None, - } - - try: - # Import the sample module - sample_path = Path(sample_dir).resolve() - main_py = sample_path / 'src' / 'main.py' - if not main_py.exists(): - main_py = sample_path / 'main.py' - - result['error'] = f'No main.py found in {sample_path}' - return result - - # Add the py/ root directory to sys.path so 'samples.shared' imports work - # sample_path is .../py/samples/sample-name - sys.path.insert(0, str(sample_path.parent.parent)) - - # Add the sample's src/ directory to sys.path for relative imports - # (e.g., 'from case_01 import prompts' in framework-restaurant-demo) - if main_py.parent.name == 'src': - # Add the sample root so 'from src import ...' works - sys.path.insert(0, str(sample_path)) - module_name = 'src.main' - else: - sys.path.insert(0, str(main_py.parent)) - module_name = 'sample_main' - - # Clear existing modules with the same name if they exist. - # This happens in a monorepo because multiple samples might use 'src.main' - # as their entry point, or have a 'src' package. - for m in list(sys.modules.keys()): - if m == module_name or m.startswith(module_name + '.') or m == 'src' or m.startswith('src.'): - sys.modules.pop(m, None) - - # Load the module - try: - if module_name == 'src.main': - # Use import_module for package-structured samples to correctly handle relative imports - module = importlib.import_module(module_name) - else: - # Fallback to spec-based loading for simple samples - spec = importlib.util.spec_from_file_location(module_name, main_py) - if not spec or not spec.loader: - result['error'] = 'Failed to load sample module' - return result - - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - except Exception as e: - result['error'] = f'Failed to import sample: {e}\n{traceback.format_exc()}' - return result - - # Find the Genkit instance - ai_instance = None - for attr_name in dir(module): - attr = getattr(module, attr_name) - if hasattr(attr, '__class__') and attr.__class__.__name__ == 'Genkit': - ai_instance = attr - break - - if not ai_instance: - result['error'] = 'No Genkit instance found in sample' - return result - - # Manually load prompts from the sample's prompts directory if it exists. - # This is needed because the current working directory isn't the sample dir. - prompts_dir = sample_path / 'prompts' - if prompts_dir.exists() and prompts_dir.is_dir(): - from genkit.blocks.prompt import load_prompt_folder - - load_prompt_folder(ai_instance.registry, prompts_dir) - - # Get the flow action from registry - try: - registry = ai_instance.registry - actions_map = await registry.resolve_actions_by_kind(ActionKind.FLOW) - - if flow_name not in actions_map: - result['error'] = f"Flow '{flow_name}' not found in registry" - return result - - flow_action = actions_map[flow_name] - except Exception as e: - result['error'] = f"Failed to retrieve flow '{flow_name}': {e}" - return result - - # Run the flow - use arun() in async context - try: - # Convert dict input to Pydantic model if an input schema is defined - validated_input = input_data - if isinstance(input_data, dict) and hasattr(flow_action, 'input_type') and flow_action.input_type: - try: - # Use the Action's Pydantic TypeAdapter to validate and convert the input - validated_input = flow_action.input_type.validate_python(input_data) - except Exception: # noqa: S110 - intentional fallback if validation fails - # If validation fails, we try with the original dict - pass - - # Always use arun() since we're in an async context - flow_result = await flow_action.arun(validated_input) - - # Extract response - response_obj = flow_result.response - - # Format output - formatted_output = format_output(response_obj) - - result['success'] = True - result['result'] = formatted_output - - except Exception as e: - # pyrefly: ignore[unbound-name] - traceback is imported at top of file - result['error'] = f'Flow execution failed: {e}\n{traceback.format_exc()}' - - except Exception as e: - result['error'] = f'Unexpected error: {e}' - - return result - - -def main() -> None: - """Run a single flow and output JSON result.""" - parser = argparse.ArgumentParser(description='Run a single Genkit flow.') - parser.add_argument('sample_dir', type=str, help='Path to sample directory') - parser.add_argument('flow_name', type=str, help='Name of flow to run') - parser.add_argument('--input', type=str, default='null', help='JSON string of input data') - args = parser.parse_args() - - # Suppress verbose logging - logging.basicConfig(level=logging.ERROR) - logging.getLogger('genkit').setLevel(logging.ERROR) - logging.getLogger('google').setLevel(logging.ERROR) - logging.getLogger('asyncio').setLevel(logging.ERROR) - logging.getLogger('httpx').setLevel(logging.ERROR) - logging.getLogger('httpcore').setLevel(logging.ERROR) - - # Parse input - try: - input_data = json.loads(args.input) - except json.JSONDecodeError: - return - - # Run flow in async context - # We do NOT redirect stdout so that logs/prints from the flow are visible - try: - result = asyncio.run(run_flow(args.sample_dir, args.flow_name, input_data)) - except Exception as e: - result = { - 'success': False, - 'result': None, - 'error': f'Subprocess execution failed: {e}\n{traceback.format_exc()}', - } - - # Print result with markers so the caller can extract it from stdout - print('\n---JSON_RESULT_START---') # noqa: T201 - print(json.dumps(result)) # noqa: T201 - print('---JSON_RESULT_END---') # noqa: T201 - - -if __name__ == '__main__': - main() diff --git a/py/uv.lock b/py/uv.lock index 65897e7b2e..71dd070f70 100644 --- a/py/uv.lock +++ b/py/uv.lock @@ -12,7 +12,6 @@ resolution-markers = [ [manifest] members = [ "conform", - "dev-local-vectorstore-hello", "framework-context-demo", "framework-custom-evaluators", "framework-dynamic-tools-demo", @@ -23,24 +22,18 @@ members = [ "framework-restaurant-demo", "framework-tool-interrupts", "genkit", - "genkit-plugin-amazon-bedrock", "genkit-plugin-anthropic", "genkit-plugin-compat-oai", - "genkit-plugin-dev-local-vectorstore", "genkit-plugin-fastapi", - "genkit-plugin-firebase", "genkit-plugin-flask", "genkit-plugin-google-cloud", "genkit-plugin-google-genai", "genkit-plugin-ollama", "genkit-plugin-vertex-ai", "genkit-tools-model-config-test", - "genkit-tools-sample-flows", "genkit-workspace", - "provider-amazon-bedrock-hello", "provider-anthropic-hello", "provider-compat-oai-hello", - "provider-firestore-retriever", "provider-google-genai-code-execution", "provider-google-genai-context-caching", "provider-google-genai-hello", @@ -49,49 +42,10 @@ members = [ "provider-google-genai-vertexai-image", "provider-ollama-hello", "provider-vertex-ai-model-garden", - "provider-vertex-ai-rerank-eval", - "provider-vertex-ai-vector-search-bigquery", - "provider-vertex-ai-vector-search-firestore", "web-fastapi-bugbot", - "web-fastapi-minimal-devui", "web-flask-hello", ] - -[[package]] -name = "aioboto3" -version = "15.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiobotocore", extra = ["boto3"] }, - { name = "aiofiles" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a2/01/92e9ab00f36e2899315f49eefcd5b4685fbb19016c7f19a9edf06da80bb0/aioboto3-15.5.0.tar.gz", hash = "sha256:ea8d8787d315594842fbfcf2c4dce3bac2ad61be275bc8584b2ce9a3402a6979", size = 255069, upload-time = "2025-10-30T13:37:16.122Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e5/3e/e8f5b665bca646d43b916763c901e00a07e40f7746c9128bdc912a089424/aioboto3-15.5.0-py3-none-any.whl", hash = "sha256:cc880c4d6a8481dd7e05da89f41c384dbd841454fc1998ae25ca9c39201437a6", size = 35913, upload-time = "2025-10-30T13:37:14.549Z" }, -] - -[[package]] -name = "aiobotocore" -version = "2.25.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohttp" }, - { name = "aioitertools" }, - { name = "botocore" }, - { name = "jmespath" }, - { name = "multidict" }, - { name = "python-dateutil" }, - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/62/94/2e4ec48cf1abb89971cb2612d86f979a6240520f0a659b53a43116d344dc/aiobotocore-2.25.1.tar.gz", hash = "sha256:ea9be739bfd7ece8864f072ec99bb9ed5c7e78ebb2b0b15f29781fbe02daedbc", size = 120560, upload-time = "2025-10-28T22:33:21.787Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/95/2a/d275ec4ce5cd0096665043995a7d76f5d0524853c76a3d04656de49f8808/aiobotocore-2.25.1-py3-none-any.whl", hash = "sha256:eb6daebe3cbef5b39a0bb2a97cffbe9c7cb46b2fcc399ad141f369f3c2134b1f", size = 86039, upload-time = "2025-10-28T22:33:19.949Z" }, -] - -[package.optional-dependencies] -boto3 = [ - { name = "boto3" }, -] +overrides = [{ name = "werkzeug", specifier = ">=3.1.6" }] [[package]] name = "aiofiles" @@ -102,157 +56,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bc/8a/340a1555ae33d7354dbca4faa54948d76d89a27ceef032c8c3bc661d003e/aiofiles-25.1.0-py3-none-any.whl", hash = "sha256:abe311e527c862958650f9438e859c1fa7568a141b22abcd015e120e86a85695", size = 14668, upload-time = "2025-10-09T20:51:03.174Z" }, ] -[[package]] -name = "aiohappyeyeballs" -version = "2.6.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/26/30/f84a107a9c4331c14b2b586036f40965c128aa4fee4dda5d3d51cb14ad54/aiohappyeyeballs-2.6.1.tar.gz", hash = "sha256:c3f9d0113123803ccadfdf3f0faa505bc78e6a72d1cc4806cbd719826e943558", size = 22760, upload-time = "2025-03-12T01:42:48.764Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0f/15/5bf3b99495fb160b63f95972b81750f18f7f4e02ad051373b669d17d44f2/aiohappyeyeballs-2.6.1-py3-none-any.whl", hash = "sha256:f349ba8f4b75cb25c99c5c2d84e997e485204d2902a9597802b0371f09331fb8", size = 15265, upload-time = "2025-03-12T01:42:47.083Z" }, -] - -[[package]] -name = "aiohttp" -version = "3.13.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohappyeyeballs" }, - { name = "aiosignal" }, - { name = "async-timeout", marker = "python_full_version < '3.11'" }, - { name = "attrs" }, - { name = "frozenlist" }, - { name = "multidict" }, - { name = "propcache" }, - { name = "yarl" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/50/42/32cf8e7704ceb4481406eb87161349abb46a57fee3f008ba9cb610968646/aiohttp-3.13.3.tar.gz", hash = "sha256:a949eee43d3782f2daae4f4a2819b2cb9b0c5d3b7f7a927067cc84dafdbb9f88", size = 7844556, upload-time = "2026-01-03T17:33:05.204Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/36/d6/5aec9313ee6ea9c7cde8b891b69f4ff4001416867104580670a31daeba5b/aiohttp-3.13.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d5a372fd5afd301b3a89582817fdcdb6c34124787c70dbcc616f259013e7eef7", size = 738950, upload-time = "2026-01-03T17:29:13.002Z" }, - { url = "https://files.pythonhosted.org/packages/68/03/8fa90a7e6d11ff20a18837a8e2b5dd23db01aabc475aa9271c8ad33299f5/aiohttp-3.13.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:147e422fd1223005c22b4fe080f5d93ced44460f5f9c105406b753612b587821", size = 496099, upload-time = "2026-01-03T17:29:15.268Z" }, - { url = "https://files.pythonhosted.org/packages/d2/23/b81f744d402510a8366b74eb420fc0cc1170d0c43daca12d10814df85f10/aiohttp-3.13.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:859bd3f2156e81dd01432f5849fc73e2243d4a487c4fd26609b1299534ee1845", size = 491072, upload-time = "2026-01-03T17:29:16.922Z" }, - { url = "https://files.pythonhosted.org/packages/d5/e1/56d1d1c0dd334cd203dd97706ce004c1aa24b34a813b0b8daf3383039706/aiohttp-3.13.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dca68018bf48c251ba17c72ed479f4dafe9dbd5a73707ad8d28a38d11f3d42af", size = 1671588, upload-time = "2026-01-03T17:29:18.539Z" }, - { url = "https://files.pythonhosted.org/packages/5f/34/8d7f962604f4bc2b4e39eb1220dac7d4e4cba91fb9ba0474b4ecd67db165/aiohttp-3.13.3-cp310-cp310-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:fee0c6bc7db1de362252affec009707a17478a00ec69f797d23ca256e36d5940", size = 1640334, upload-time = "2026-01-03T17:29:21.028Z" }, - { url = "https://files.pythonhosted.org/packages/94/1d/fcccf2c668d87337ddeef9881537baee13c58d8f01f12ba8a24215f2b804/aiohttp-3.13.3-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c048058117fd649334d81b4b526e94bde3ccaddb20463a815ced6ecbb7d11160", size = 1722656, upload-time = "2026-01-03T17:29:22.531Z" }, - { url = "https://files.pythonhosted.org/packages/aa/98/c6f3b081c4c606bc1e5f2ec102e87d6411c73a9ef3616fea6f2d5c98c062/aiohttp-3.13.3-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:215a685b6fbbfcf71dfe96e3eba7a6f58f10da1dfdf4889c7dd856abe430dca7", size = 1817625, upload-time = "2026-01-03T17:29:24.276Z" }, - { url = "https://files.pythonhosted.org/packages/2c/c0/cfcc3d2e11b477f86e1af2863f3858c8850d751ce8dc39c4058a072c9e54/aiohttp-3.13.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:de2c184bb1fe2cbd2cefba613e9db29a5ab559323f994b6737e370d3da0ac455", size = 1672604, upload-time = "2026-01-03T17:29:26.099Z" }, - { url = "https://files.pythonhosted.org/packages/1e/77/6b4ffcbcac4c6a5d041343a756f34a6dd26174ae07f977a64fe028dda5b0/aiohttp-3.13.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:75ca857eba4e20ce9f546cd59c7007b33906a4cd48f2ff6ccf1ccfc3b646f279", size = 1554370, upload-time = "2026-01-03T17:29:28.121Z" }, - { url = "https://files.pythonhosted.org/packages/f2/f0/e3ddfa93f17d689dbe014ba048f18e0c9f9b456033b70e94349a2e9048be/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:81e97251d9298386c2b7dbeb490d3d1badbdc69107fb8c9299dd04eb39bddc0e", size = 1642023, upload-time = "2026-01-03T17:29:30.002Z" }, - { url = "https://files.pythonhosted.org/packages/eb/45/c14019c9ec60a8e243d06d601b33dcc4fd92379424bde3021725859d7f99/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:c0e2d366af265797506f0283487223146af57815b388623f0357ef7eac9b209d", size = 1649680, upload-time = "2026-01-03T17:29:31.782Z" }, - { url = "https://files.pythonhosted.org/packages/9c/fd/09c9451dae5aa5c5ed756df95ff9ef549d45d4be663bafd1e4954fd836f0/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:4e239d501f73d6db1522599e14b9b321a7e3b1de66ce33d53a765d975e9f4808", size = 1692407, upload-time = "2026-01-03T17:29:33.392Z" }, - { url = "https://files.pythonhosted.org/packages/a6/81/938bc2ec33c10efd6637ccb3d22f9f3160d08e8f3aa2587a2c2d5ab578eb/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:0db318f7a6f065d84cb1e02662c526294450b314a02bd9e2a8e67f0d8564ce40", size = 1543047, upload-time = "2026-01-03T17:29:34.855Z" }, - { url = "https://files.pythonhosted.org/packages/f7/23/80488ee21c8d567c83045e412e1d9b7077d27171591a4eb7822586e8c06a/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:bfc1cc2fe31a6026a8a88e4ecfb98d7f6b1fec150cfd708adbfd1d2f42257c29", size = 1715264, upload-time = "2026-01-03T17:29:36.389Z" }, - { url = "https://files.pythonhosted.org/packages/e2/83/259a8da6683182768200b368120ab3deff5370bed93880fb9a3a86299f34/aiohttp-3.13.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:af71fff7bac6bb7508956696dce8f6eec2bbb045eceb40343944b1ae62b5ef11", size = 1657275, upload-time = "2026-01-03T17:29:38.162Z" }, - { url = "https://files.pythonhosted.org/packages/3f/4f/2c41f800a0b560785c10fb316216ac058c105f9be50bdc6a285de88db625/aiohttp-3.13.3-cp310-cp310-win32.whl", hash = "sha256:37da61e244d1749798c151421602884db5270faf479cf0ef03af0ff68954c9dd", size = 434053, upload-time = "2026-01-03T17:29:40.074Z" }, - { url = "https://files.pythonhosted.org/packages/80/df/29cd63c7ecfdb65ccc12f7d808cac4fa2a19544660c06c61a4a48462de0c/aiohttp-3.13.3-cp310-cp310-win_amd64.whl", hash = "sha256:7e63f210bc1b57ef699035f2b4b6d9ce096b5914414a49b0997c839b2bd2223c", size = 456687, upload-time = "2026-01-03T17:29:41.819Z" }, - { url = "https://files.pythonhosted.org/packages/f1/4c/a164164834f03924d9a29dc3acd9e7ee58f95857e0b467f6d04298594ebb/aiohttp-3.13.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:5b6073099fb654e0a068ae678b10feff95c5cae95bbfcbfa7af669d361a8aa6b", size = 746051, upload-time = "2026-01-03T17:29:43.287Z" }, - { url = "https://files.pythonhosted.org/packages/82/71/d5c31390d18d4f58115037c432b7e0348c60f6f53b727cad33172144a112/aiohttp-3.13.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cb93e166e6c28716c8c6aeb5f99dfb6d5ccf482d29fe9bf9a794110e6d0ab64", size = 499234, upload-time = "2026-01-03T17:29:44.822Z" }, - { url = "https://files.pythonhosted.org/packages/0e/c9/741f8ac91e14b1d2e7100690425a5b2b919a87a5075406582991fb7de920/aiohttp-3.13.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:28e027cf2f6b641693a09f631759b4d9ce9165099d2b5d92af9bd4e197690eea", size = 494979, upload-time = "2026-01-03T17:29:46.405Z" }, - { url = "https://files.pythonhosted.org/packages/75/b5/31d4d2e802dfd59f74ed47eba48869c1c21552c586d5e81a9d0d5c2ad640/aiohttp-3.13.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3b61b7169ababd7802f9568ed96142616a9118dd2be0d1866e920e77ec8fa92a", size = 1748297, upload-time = "2026-01-03T17:29:48.083Z" }, - { url = "https://files.pythonhosted.org/packages/1a/3e/eefad0ad42959f226bb79664826883f2687d602a9ae2941a18e0484a74d3/aiohttp-3.13.3-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:80dd4c21b0f6237676449c6baaa1039abae86b91636b6c91a7f8e61c87f89540", size = 1707172, upload-time = "2026-01-03T17:29:49.648Z" }, - { url = "https://files.pythonhosted.org/packages/c5/3a/54a64299fac2891c346cdcf2aa6803f994a2e4beeaf2e5a09dcc54acc842/aiohttp-3.13.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:65d2ccb7eabee90ce0503c17716fc77226be026dcc3e65cce859a30db715025b", size = 1805405, upload-time = "2026-01-03T17:29:51.244Z" }, - { url = "https://files.pythonhosted.org/packages/6c/70/ddc1b7169cf64075e864f64595a14b147a895a868394a48f6a8031979038/aiohttp-3.13.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5b179331a481cb5529fca8b432d8d3c7001cb217513c94cd72d668d1248688a3", size = 1899449, upload-time = "2026-01-03T17:29:53.938Z" }, - { url = "https://files.pythonhosted.org/packages/a1/7e/6815aab7d3a56610891c76ef79095677b8b5be6646aaf00f69b221765021/aiohttp-3.13.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d4c940f02f49483b18b079d1c27ab948721852b281f8b015c058100e9421dd1", size = 1748444, upload-time = "2026-01-03T17:29:55.484Z" }, - { url = "https://files.pythonhosted.org/packages/6b/f2/073b145c4100da5511f457dc0f7558e99b2987cf72600d42b559db856fbc/aiohttp-3.13.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f9444f105664c4ce47a2a7171a2418bce5b7bae45fb610f4e2c36045d85911d3", size = 1606038, upload-time = "2026-01-03T17:29:57.179Z" }, - { url = "https://files.pythonhosted.org/packages/0a/c1/778d011920cae03ae01424ec202c513dc69243cf2db303965615b81deeea/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:694976222c711d1d00ba131904beb60534f93966562f64440d0c9d41b8cdb440", size = 1724156, upload-time = "2026-01-03T17:29:58.914Z" }, - { url = "https://files.pythonhosted.org/packages/0e/cb/3419eabf4ec1e9ec6f242c32b689248365a1cf621891f6f0386632525494/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:f33ed1a2bf1997a36661874b017f5c4b760f41266341af36febaf271d179f6d7", size = 1722340, upload-time = "2026-01-03T17:30:01.962Z" }, - { url = "https://files.pythonhosted.org/packages/7a/e5/76cf77bdbc435bf233c1f114edad39ed4177ccbfab7c329482b179cff4f4/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:e636b3c5f61da31a92bf0d91da83e58fdfa96f178ba682f11d24f31944cdd28c", size = 1783041, upload-time = "2026-01-03T17:30:03.609Z" }, - { url = "https://files.pythonhosted.org/packages/9d/d4/dd1ca234c794fd29c057ce8c0566b8ef7fd6a51069de5f06fa84b9a1971c/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:5d2d94f1f5fcbe40838ac51a6ab5704a6f9ea42e72ceda48de5e6b898521da51", size = 1596024, upload-time = "2026-01-03T17:30:05.132Z" }, - { url = "https://files.pythonhosted.org/packages/55/58/4345b5f26661a6180afa686c473620c30a66afdf120ed3dd545bbc809e85/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:2be0e9ccf23e8a94f6f0650ce06042cefc6ac703d0d7ab6c7a917289f2539ad4", size = 1804590, upload-time = "2026-01-03T17:30:07.135Z" }, - { url = "https://files.pythonhosted.org/packages/7b/06/05950619af6c2df7e0a431d889ba2813c9f0129cec76f663e547a5ad56f2/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9af5e68ee47d6534d36791bbe9b646d2a7c7deb6fc24d7943628edfbb3581f29", size = 1740355, upload-time = "2026-01-03T17:30:09.083Z" }, - { url = "https://files.pythonhosted.org/packages/3e/80/958f16de79ba0422d7c1e284b2abd0c84bc03394fbe631d0a39ffa10e1eb/aiohttp-3.13.3-cp311-cp311-win32.whl", hash = "sha256:a2212ad43c0833a873d0fb3c63fa1bacedd4cf6af2fee62bf4b739ceec3ab239", size = 433701, upload-time = "2026-01-03T17:30:10.869Z" }, - { url = "https://files.pythonhosted.org/packages/dc/f2/27cdf04c9851712d6c1b99df6821a6623c3c9e55956d4b1e318c337b5a48/aiohttp-3.13.3-cp311-cp311-win_amd64.whl", hash = "sha256:642f752c3eb117b105acbd87e2c143de710987e09860d674e068c4c2c441034f", size = 457678, upload-time = "2026-01-03T17:30:12.719Z" }, - { url = "https://files.pythonhosted.org/packages/a0/be/4fc11f202955a69e0db803a12a062b8379c970c7c84f4882b6da17337cc1/aiohttp-3.13.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:b903a4dfee7d347e2d87697d0713be59e0b87925be030c9178c5faa58ea58d5c", size = 739732, upload-time = "2026-01-03T17:30:14.23Z" }, - { url = "https://files.pythonhosted.org/packages/97/2c/621d5b851f94fa0bb7430d6089b3aa970a9d9b75196bc93bb624b0db237a/aiohttp-3.13.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a45530014d7a1e09f4a55f4f43097ba0fd155089372e105e4bff4ca76cb1b168", size = 494293, upload-time = "2026-01-03T17:30:15.96Z" }, - { url = "https://files.pythonhosted.org/packages/5d/43/4be01406b78e1be8320bb8316dc9c42dbab553d281c40364e0f862d5661c/aiohttp-3.13.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:27234ef6d85c914f9efeb77ff616dbf4ad2380be0cda40b4db086ffc7ddd1b7d", size = 493533, upload-time = "2026-01-03T17:30:17.431Z" }, - { url = "https://files.pythonhosted.org/packages/8d/a8/5a35dc56a06a2c90d4742cbf35294396907027f80eea696637945a106f25/aiohttp-3.13.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d32764c6c9aafb7fb55366a224756387cd50bfa720f32b88e0e6fa45b27dcf29", size = 1737839, upload-time = "2026-01-03T17:30:19.422Z" }, - { url = "https://files.pythonhosted.org/packages/bf/62/4b9eeb331da56530bf2e198a297e5303e1c1ebdceeb00fe9b568a65c5a0c/aiohttp-3.13.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:b1a6102b4d3ebc07dad44fbf07b45bb600300f15b552ddf1851b5390202ea2e3", size = 1703932, upload-time = "2026-01-03T17:30:21.756Z" }, - { url = "https://files.pythonhosted.org/packages/7c/f6/af16887b5d419e6a367095994c0b1332d154f647e7dc2bd50e61876e8e3d/aiohttp-3.13.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c014c7ea7fb775dd015b2d3137378b7be0249a448a1612268b5a90c2d81de04d", size = 1771906, upload-time = "2026-01-03T17:30:23.932Z" }, - { url = "https://files.pythonhosted.org/packages/ce/83/397c634b1bcc24292fa1e0c7822800f9f6569e32934bdeef09dae7992dfb/aiohttp-3.13.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2b8d8ddba8f95ba17582226f80e2de99c7a7948e66490ef8d947e272a93e9463", size = 1871020, upload-time = "2026-01-03T17:30:26Z" }, - { url = "https://files.pythonhosted.org/packages/86/f6/a62cbbf13f0ac80a70f71b1672feba90fdb21fd7abd8dbf25c0105fb6fa3/aiohttp-3.13.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ae8dd55c8e6c4257eae3a20fd2c8f41edaea5992ed67156642493b8daf3cecc", size = 1755181, upload-time = "2026-01-03T17:30:27.554Z" }, - { url = "https://files.pythonhosted.org/packages/0a/87/20a35ad487efdd3fba93d5843efdfaa62d2f1479eaafa7453398a44faf13/aiohttp-3.13.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:01ad2529d4b5035578f5081606a465f3b814c542882804e2e8cda61adf5c71bf", size = 1561794, upload-time = "2026-01-03T17:30:29.254Z" }, - { url = "https://files.pythonhosted.org/packages/de/95/8fd69a66682012f6716e1bc09ef8a1a2a91922c5725cb904689f112309c4/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bb4f7475e359992b580559e008c598091c45b5088f28614e855e42d39c2f1033", size = 1697900, upload-time = "2026-01-03T17:30:31.033Z" }, - { url = "https://files.pythonhosted.org/packages/e5/66/7b94b3b5ba70e955ff597672dad1691333080e37f50280178967aff68657/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:c19b90316ad3b24c69cd78d5c9b4f3aa4497643685901185b65166293d36a00f", size = 1728239, upload-time = "2026-01-03T17:30:32.703Z" }, - { url = "https://files.pythonhosted.org/packages/47/71/6f72f77f9f7d74719692ab65a2a0252584bf8d5f301e2ecb4c0da734530a/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:96d604498a7c782cb15a51c406acaea70d8c027ee6b90c569baa6e7b93073679", size = 1740527, upload-time = "2026-01-03T17:30:34.695Z" }, - { url = "https://files.pythonhosted.org/packages/fa/b4/75ec16cbbd5c01bdaf4a05b19e103e78d7ce1ef7c80867eb0ace42ff4488/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:084911a532763e9d3dd95adf78a78f4096cd5f58cdc18e6fdbc1b58417a45423", size = 1554489, upload-time = "2026-01-03T17:30:36.864Z" }, - { url = "https://files.pythonhosted.org/packages/52/8f/bc518c0eea29f8406dcf7ed1f96c9b48e3bc3995a96159b3fc11f9e08321/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:7a4a94eb787e606d0a09404b9c38c113d3b099d508021faa615d70a0131907ce", size = 1767852, upload-time = "2026-01-03T17:30:39.433Z" }, - { url = "https://files.pythonhosted.org/packages/9d/f2/a07a75173124f31f11ea6f863dc44e6f09afe2bca45dd4e64979490deab1/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:87797e645d9d8e222e04160ee32aa06bc5c163e8499f24db719e7852ec23093a", size = 1722379, upload-time = "2026-01-03T17:30:41.081Z" }, - { url = "https://files.pythonhosted.org/packages/3c/4a/1a3fee7c21350cac78e5c5cef711bac1b94feca07399f3d406972e2d8fcd/aiohttp-3.13.3-cp312-cp312-win32.whl", hash = "sha256:b04be762396457bef43f3597c991e192ee7da460a4953d7e647ee4b1c28e7046", size = 428253, upload-time = "2026-01-03T17:30:42.644Z" }, - { url = "https://files.pythonhosted.org/packages/d9/b7/76175c7cb4eb73d91ad63c34e29fc4f77c9386bba4a65b53ba8e05ee3c39/aiohttp-3.13.3-cp312-cp312-win_amd64.whl", hash = "sha256:e3531d63d3bdfa7e3ac5e9b27b2dd7ec9df3206a98e0b3445fa906f233264c57", size = 455407, upload-time = "2026-01-03T17:30:44.195Z" }, - { url = "https://files.pythonhosted.org/packages/97/8a/12ca489246ca1faaf5432844adbfce7ff2cc4997733e0af120869345643a/aiohttp-3.13.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:5dff64413671b0d3e7d5918ea490bdccb97a4ad29b3f311ed423200b2203e01c", size = 734190, upload-time = "2026-01-03T17:30:45.832Z" }, - { url = "https://files.pythonhosted.org/packages/32/08/de43984c74ed1fca5c014808963cc83cb00d7bb06af228f132d33862ca76/aiohttp-3.13.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:87b9aab6d6ed88235aa2970294f496ff1a1f9adcd724d800e9b952395a80ffd9", size = 491783, upload-time = "2026-01-03T17:30:47.466Z" }, - { url = "https://files.pythonhosted.org/packages/17/f8/8dd2cf6112a5a76f81f81a5130c57ca829d101ad583ce57f889179accdda/aiohttp-3.13.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:425c126c0dc43861e22cb1c14ba4c8e45d09516d0a3ae0a3f7494b79f5f233a3", size = 490704, upload-time = "2026-01-03T17:30:49.373Z" }, - { url = "https://files.pythonhosted.org/packages/6d/40/a46b03ca03936f832bc7eaa47cfbb1ad012ba1be4790122ee4f4f8cba074/aiohttp-3.13.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7f9120f7093c2a32d9647abcaf21e6ad275b4fbec5b55969f978b1a97c7c86bf", size = 1720652, upload-time = "2026-01-03T17:30:50.974Z" }, - { url = "https://files.pythonhosted.org/packages/f7/7e/917fe18e3607af92657e4285498f500dca797ff8c918bd7d90b05abf6c2a/aiohttp-3.13.3-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:697753042d57f4bf7122cab985bf15d0cef23c770864580f5af4f52023a56bd6", size = 1692014, upload-time = "2026-01-03T17:30:52.729Z" }, - { url = "https://files.pythonhosted.org/packages/71/b6/cefa4cbc00d315d68973b671cf105b21a609c12b82d52e5d0c9ae61d2a09/aiohttp-3.13.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6de499a1a44e7de70735d0b39f67c8f25eb3d91eb3103be99ca0fa882cdd987d", size = 1759777, upload-time = "2026-01-03T17:30:54.537Z" }, - { url = "https://files.pythonhosted.org/packages/fb/e3/e06ee07b45e59e6d81498b591fc589629be1553abb2a82ce33efe2a7b068/aiohttp-3.13.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:37239e9f9a7ea9ac5bf6b92b0260b01f8a22281996da609206a84df860bc1261", size = 1861276, upload-time = "2026-01-03T17:30:56.512Z" }, - { url = "https://files.pythonhosted.org/packages/7c/24/75d274228acf35ceeb2850b8ce04de9dd7355ff7a0b49d607ee60c29c518/aiohttp-3.13.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f76c1e3fe7d7c8afad7ed193f89a292e1999608170dcc9751a7462a87dfd5bc0", size = 1743131, upload-time = "2026-01-03T17:30:58.256Z" }, - { url = "https://files.pythonhosted.org/packages/04/98/3d21dde21889b17ca2eea54fdcff21b27b93f45b7bb94ca029c31ab59dc3/aiohttp-3.13.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fc290605db2a917f6e81b0e1e0796469871f5af381ce15c604a3c5c7e51cb730", size = 1556863, upload-time = "2026-01-03T17:31:00.445Z" }, - { url = "https://files.pythonhosted.org/packages/9e/84/da0c3ab1192eaf64782b03971ab4055b475d0db07b17eff925e8c93b3aa5/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4021b51936308aeea0367b8f006dc999ca02bc118a0cc78c303f50a2ff6afb91", size = 1682793, upload-time = "2026-01-03T17:31:03.024Z" }, - { url = "https://files.pythonhosted.org/packages/ff/0f/5802ada182f575afa02cbd0ec5180d7e13a402afb7c2c03a9aa5e5d49060/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:49a03727c1bba9a97d3e93c9f93ca03a57300f484b6e935463099841261195d3", size = 1716676, upload-time = "2026-01-03T17:31:04.842Z" }, - { url = "https://files.pythonhosted.org/packages/3f/8c/714d53bd8b5a4560667f7bbbb06b20c2382f9c7847d198370ec6526af39c/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3d9908a48eb7416dc1f4524e69f1d32e5d90e3981e4e37eb0aa1cd18f9cfa2a4", size = 1733217, upload-time = "2026-01-03T17:31:06.868Z" }, - { url = "https://files.pythonhosted.org/packages/7d/79/e2176f46d2e963facea939f5be2d26368ce543622be6f00a12844d3c991f/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:2712039939ec963c237286113c68dbad80a82a4281543f3abf766d9d73228998", size = 1552303, upload-time = "2026-01-03T17:31:08.958Z" }, - { url = "https://files.pythonhosted.org/packages/ab/6a/28ed4dea1759916090587d1fe57087b03e6c784a642b85ef48217b0277ae/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:7bfdc049127717581866fa4708791220970ce291c23e28ccf3922c700740fdc0", size = 1763673, upload-time = "2026-01-03T17:31:10.676Z" }, - { url = "https://files.pythonhosted.org/packages/e8/35/4a3daeb8b9fab49240d21c04d50732313295e4bd813a465d840236dd0ce1/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8057c98e0c8472d8846b9c79f56766bcc57e3e8ac7bfd510482332366c56c591", size = 1721120, upload-time = "2026-01-03T17:31:12.575Z" }, - { url = "https://files.pythonhosted.org/packages/bc/9f/d643bb3c5fb99547323e635e251c609fbbc660d983144cfebec529e09264/aiohttp-3.13.3-cp313-cp313-win32.whl", hash = "sha256:1449ceddcdbcf2e0446957863af03ebaaa03f94c090f945411b61269e2cb5daf", size = 427383, upload-time = "2026-01-03T17:31:14.382Z" }, - { url = "https://files.pythonhosted.org/packages/4e/f1/ab0395f8a79933577cdd996dd2f9aa6014af9535f65dddcf88204682fe62/aiohttp-3.13.3-cp313-cp313-win_amd64.whl", hash = "sha256:693781c45a4033d31d4187d2436f5ac701e7bbfe5df40d917736108c1cc7436e", size = 453899, upload-time = "2026-01-03T17:31:15.958Z" }, - { url = "https://files.pythonhosted.org/packages/99/36/5b6514a9f5d66f4e2597e40dea2e3db271e023eb7a5d22defe96ba560996/aiohttp-3.13.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:ea37047c6b367fd4bd632bff8077449b8fa034b69e812a18e0132a00fae6e808", size = 737238, upload-time = "2026-01-03T17:31:17.909Z" }, - { url = "https://files.pythonhosted.org/packages/f7/49/459327f0d5bcd8c6c9ca69e60fdeebc3622861e696490d8674a6d0cb90a6/aiohttp-3.13.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:6fc0e2337d1a4c3e6acafda6a78a39d4c14caea625124817420abceed36e2415", size = 492292, upload-time = "2026-01-03T17:31:19.919Z" }, - { url = "https://files.pythonhosted.org/packages/e8/0b/b97660c5fd05d3495b4eb27f2d0ef18dc1dc4eff7511a9bf371397ff0264/aiohttp-3.13.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c685f2d80bb67ca8c3837823ad76196b3694b0159d232206d1e461d3d434666f", size = 493021, upload-time = "2026-01-03T17:31:21.636Z" }, - { url = "https://files.pythonhosted.org/packages/54/d4/438efabdf74e30aeceb890c3290bbaa449780583b1270b00661126b8aae4/aiohttp-3.13.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48e377758516d262bde50c2584fc6c578af272559c409eecbdd2bae1601184d6", size = 1717263, upload-time = "2026-01-03T17:31:23.296Z" }, - { url = "https://files.pythonhosted.org/packages/71/f2/7bddc7fd612367d1459c5bcf598a9e8f7092d6580d98de0e057eb42697ad/aiohttp-3.13.3-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:34749271508078b261c4abb1767d42b8d0c0cc9449c73a4df494777dc55f0687", size = 1669107, upload-time = "2026-01-03T17:31:25.334Z" }, - { url = "https://files.pythonhosted.org/packages/00/5a/1aeaecca40e22560f97610a329e0e5efef5e0b5afdf9f857f0d93839ab2e/aiohttp-3.13.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:82611aeec80eb144416956ec85b6ca45a64d76429c1ed46ae1b5f86c6e0c9a26", size = 1760196, upload-time = "2026-01-03T17:31:27.394Z" }, - { url = "https://files.pythonhosted.org/packages/f8/f8/0ff6992bea7bd560fc510ea1c815f87eedd745fe035589c71ce05612a19a/aiohttp-3.13.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2fff83cfc93f18f215896e3a190e8e5cb413ce01553901aca925176e7568963a", size = 1843591, upload-time = "2026-01-03T17:31:29.238Z" }, - { url = "https://files.pythonhosted.org/packages/e3/d1/e30e537a15f53485b61f5be525f2157da719819e8377298502aebac45536/aiohttp-3.13.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bbe7d4cecacb439e2e2a8a1a7b935c25b812af7a5fd26503a66dadf428e79ec1", size = 1720277, upload-time = "2026-01-03T17:31:31.053Z" }, - { url = "https://files.pythonhosted.org/packages/84/45/23f4c451d8192f553d38d838831ebbc156907ea6e05557f39563101b7717/aiohttp-3.13.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b928f30fe49574253644b1ca44b1b8adbd903aa0da4b9054a6c20fc7f4092a25", size = 1548575, upload-time = "2026-01-03T17:31:32.87Z" }, - { url = "https://files.pythonhosted.org/packages/6a/ed/0a42b127a43712eda7807e7892c083eadfaf8429ca8fb619662a530a3aab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7b5e8fe4de30df199155baaf64f2fcd604f4c678ed20910db8e2c66dc4b11603", size = 1679455, upload-time = "2026-01-03T17:31:34.76Z" }, - { url = "https://files.pythonhosted.org/packages/2e/b5/c05f0c2b4b4fe2c9d55e73b6d3ed4fd6c9dc2684b1d81cbdf77e7fad9adb/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:8542f41a62bcc58fc7f11cf7c90e0ec324ce44950003feb70640fc2a9092c32a", size = 1687417, upload-time = "2026-01-03T17:31:36.699Z" }, - { url = "https://files.pythonhosted.org/packages/c9/6b/915bc5dad66aef602b9e459b5a973529304d4e89ca86999d9d75d80cbd0b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:5e1d8c8b8f1d91cd08d8f4a3c2b067bfca6ec043d3ff36de0f3a715feeedf926", size = 1729968, upload-time = "2026-01-03T17:31:38.622Z" }, - { url = "https://files.pythonhosted.org/packages/11/3b/e84581290a9520024a08640b63d07673057aec5ca548177a82026187ba73/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:90455115e5da1c3c51ab619ac57f877da8fd6d73c05aacd125c5ae9819582aba", size = 1545690, upload-time = "2026-01-03T17:31:40.57Z" }, - { url = "https://files.pythonhosted.org/packages/f5/04/0c3655a566c43fd647c81b895dfe361b9f9ad6d58c19309d45cff52d6c3b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:042e9e0bcb5fba81886c8b4fbb9a09d6b8a00245fd8d88e4d989c1f96c74164c", size = 1746390, upload-time = "2026-01-03T17:31:42.857Z" }, - { url = "https://files.pythonhosted.org/packages/1f/53/71165b26978f719c3419381514c9690bd5980e764a09440a10bb816ea4ab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2eb752b102b12a76ca02dff751a801f028b4ffbbc478840b473597fc91a9ed43", size = 1702188, upload-time = "2026-01-03T17:31:44.984Z" }, - { url = "https://files.pythonhosted.org/packages/29/a7/cbe6c9e8e136314fa1980da388a59d2f35f35395948a08b6747baebb6aa6/aiohttp-3.13.3-cp314-cp314-win32.whl", hash = "sha256:b556c85915d8efaed322bf1bdae9486aa0f3f764195a0fb6ee962e5c71ef5ce1", size = 433126, upload-time = "2026-01-03T17:31:47.463Z" }, - { url = "https://files.pythonhosted.org/packages/de/56/982704adea7d3b16614fc5936014e9af85c0e34b58f9046655817f04306e/aiohttp-3.13.3-cp314-cp314-win_amd64.whl", hash = "sha256:9bf9f7a65e7aa20dd764151fb3d616c81088f91f8df39c3893a536e279b4b984", size = 459128, upload-time = "2026-01-03T17:31:49.2Z" }, - { url = "https://files.pythonhosted.org/packages/6c/2a/3c79b638a9c3d4658d345339d22070241ea341ed4e07b5ac60fb0f418003/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:05861afbbec40650d8a07ea324367cb93e9e8cc7762e04dd4405df99fa65159c", size = 769512, upload-time = "2026-01-03T17:31:51.134Z" }, - { url = "https://files.pythonhosted.org/packages/29/b9/3e5014d46c0ab0db8707e0ac2711ed28c4da0218c358a4e7c17bae0d8722/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:2fc82186fadc4a8316768d61f3722c230e2c1dcab4200d52d2ebdf2482e47592", size = 506444, upload-time = "2026-01-03T17:31:52.85Z" }, - { url = "https://files.pythonhosted.org/packages/90/03/c1d4ef9a054e151cd7839cdc497f2638f00b93cbe8043983986630d7a80c/aiohttp-3.13.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0add0900ff220d1d5c5ebbf99ed88b0c1bbf87aa7e4262300ed1376a6b13414f", size = 510798, upload-time = "2026-01-03T17:31:54.91Z" }, - { url = "https://files.pythonhosted.org/packages/ea/76/8c1e5abbfe8e127c893fe7ead569148a4d5a799f7cf958d8c09f3eedf097/aiohttp-3.13.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:568f416a4072fbfae453dcf9a99194bbb8bdeab718e08ee13dfa2ba0e4bebf29", size = 1868835, upload-time = "2026-01-03T17:31:56.733Z" }, - { url = "https://files.pythonhosted.org/packages/8e/ac/984c5a6f74c363b01ff97adc96a3976d9c98940b8969a1881575b279ac5d/aiohttp-3.13.3-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:add1da70de90a2569c5e15249ff76a631ccacfe198375eead4aadf3b8dc849dc", size = 1720486, upload-time = "2026-01-03T17:31:58.65Z" }, - { url = "https://files.pythonhosted.org/packages/b2/9a/b7039c5f099c4eb632138728828b33428585031a1e658d693d41d07d89d1/aiohttp-3.13.3-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:10b47b7ba335d2e9b1239fa571131a87e2d8ec96b333e68b2a305e7a98b0bae2", size = 1847951, upload-time = "2026-01-03T17:32:00.989Z" }, - { url = "https://files.pythonhosted.org/packages/3c/02/3bec2b9a1ba3c19ff89a43a19324202b8eb187ca1e928d8bdac9bbdddebd/aiohttp-3.13.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3dd4dce1c718e38081c8f35f323209d4c1df7d4db4bab1b5c88a6b4d12b74587", size = 1941001, upload-time = "2026-01-03T17:32:03.122Z" }, - { url = "https://files.pythonhosted.org/packages/37/df/d879401cedeef27ac4717f6426c8c36c3091c6e9f08a9178cc87549c537f/aiohttp-3.13.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34bac00a67a812570d4a460447e1e9e06fae622946955f939051e7cc895cfab8", size = 1797246, upload-time = "2026-01-03T17:32:05.255Z" }, - { url = "https://files.pythonhosted.org/packages/8d/15/be122de1f67e6953add23335c8ece6d314ab67c8bebb3f181063010795a7/aiohttp-3.13.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a19884d2ee70b06d9204b2727a7b9f983d0c684c650254679e716b0b77920632", size = 1627131, upload-time = "2026-01-03T17:32:07.607Z" }, - { url = "https://files.pythonhosted.org/packages/12/12/70eedcac9134cfa3219ab7af31ea56bc877395b1ac30d65b1bc4b27d0438/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5f8ca7f2bb6ba8348a3614c7918cc4bb73268c5ac2a207576b7afea19d3d9f64", size = 1795196, upload-time = "2026-01-03T17:32:09.59Z" }, - { url = "https://files.pythonhosted.org/packages/32/11/b30e1b1cd1f3054af86ebe60df96989c6a414dd87e27ad16950eee420bea/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:b0d95340658b9d2f11d9697f59b3814a9d3bb4b7a7c20b131df4bcef464037c0", size = 1782841, upload-time = "2026-01-03T17:32:11.445Z" }, - { url = "https://files.pythonhosted.org/packages/88/0d/d98a9367b38912384a17e287850f5695c528cff0f14f791ce8ee2e4f7796/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:a1e53262fd202e4b40b70c3aff944a8155059beedc8a89bba9dc1f9ef06a1b56", size = 1795193, upload-time = "2026-01-03T17:32:13.705Z" }, - { url = "https://files.pythonhosted.org/packages/43/a5/a2dfd1f5ff5581632c7f6a30e1744deda03808974f94f6534241ef60c751/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:d60ac9663f44168038586cab2157e122e46bdef09e9368b37f2d82d354c23f72", size = 1621979, upload-time = "2026-01-03T17:32:15.965Z" }, - { url = "https://files.pythonhosted.org/packages/fa/f0/12973c382ae7c1cccbc4417e129c5bf54c374dfb85af70893646e1f0e749/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:90751b8eed69435bac9ff4e3d2f6b3af1f57e37ecb0fbeee59c0174c9e2d41df", size = 1822193, upload-time = "2026-01-03T17:32:18.219Z" }, - { url = "https://files.pythonhosted.org/packages/3c/5f/24155e30ba7f8c96918af1350eb0663e2430aad9e001c0489d89cd708ab1/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:fc353029f176fd2b3ec6cfc71be166aba1936fe5d73dd1992ce289ca6647a9aa", size = 1769801, upload-time = "2026-01-03T17:32:20.25Z" }, - { url = "https://files.pythonhosted.org/packages/eb/f8/7314031ff5c10e6ece114da79b338ec17eeff3a079e53151f7e9f43c4723/aiohttp-3.13.3-cp314-cp314t-win32.whl", hash = "sha256:2e41b18a58da1e474a057b3d35248d8320029f61d70a37629535b16a0c8f3767", size = 466523, upload-time = "2026-01-03T17:32:22.215Z" }, - { url = "https://files.pythonhosted.org/packages/b4/63/278a98c715ae467624eafe375542d8ba9b4383a016df8fdefe0ae28382a7/aiohttp-3.13.3-cp314-cp314t-win_amd64.whl", hash = "sha256:44531a36aa2264a1860089ffd4dce7baf875ee5a6079d5fb42e261c704ef7344", size = 499694, upload-time = "2026-01-03T17:32:24.546Z" }, -] - -[[package]] -name = "aioitertools" -version = "0.13.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fd/3c/53c4a17a05fb9ea2313ee1777ff53f5e001aefd5cc85aa2f4c2d982e1e38/aioitertools-0.13.0.tar.gz", hash = "sha256:620bd241acc0bbb9ec819f1ab215866871b4bbd1f73836a55f799200ee86950c", size = 19322, upload-time = "2025-11-06T22:17:07.609Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl", hash = "sha256:0be0292b856f08dfac90e31f4739432f4cb6d7520ab9eb73e143f4f2fa5259be", size = 24182, upload-time = "2025-11-06T22:17:06.502Z" }, -] - -[[package]] -name = "aiosignal" -version = "1.4.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "frozenlist" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, -] - [[package]] name = "altair" version = "6.0.0" @@ -441,15 +244,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2e/e9/eb6a5db5ac505d5d45715388e92bced7a5bb556facc4d0865d192823f2d2/async_lru-2.1.0-py3-none-any.whl", hash = "sha256:fa12dcf99a42ac1280bc16c634bbaf06883809790f6304d85cdab3f666f33a7e", size = 6933, upload-time = "2026-01-17T22:52:17.389Z" }, ] -[[package]] -name = "async-timeout" -version = "5.0.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a5/ae/136395dfbfe00dfc94da3f3e136d0b13f394cba8f4841120e34226265780/async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3", size = 9274, upload-time = "2024-11-06T16:41:39.6Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fe/ba/e2081de779ca30d473f21f5b30e0e737c438205440784c7dfc81efc2b029/async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c", size = 6233, upload-time = "2024-11-06T16:41:37.9Z" }, -] - [[package]] name = "attrs" version = "25.4.0" @@ -606,34 +400,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/ca/78d423b324b8d77900030fa59c4aa9054261ef0925631cd2501dd015b7b7/boolean_py-5.0-py3-none-any.whl", hash = "sha256:ef28a70bd43115208441b53a045d1549e2f0ec6e3d08a9d142cbc41c1938e8d9", size = 26577, upload-time = "2025-04-03T10:39:48.449Z" }, ] -[[package]] -name = "boto3" -version = "1.40.61" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "botocore" }, - { name = "jmespath" }, - { name = "s3transfer" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ed/f9/6ef8feb52c3cce5ec3967a535a6114b57ac7949fd166b0f3090c2b06e4e5/boto3-1.40.61.tar.gz", hash = "sha256:d6c56277251adf6c2bdd25249feae625abe4966831676689ff23b4694dea5b12", size = 111535, upload-time = "2025-10-28T19:26:57.247Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/61/24/3bf865b07d15fea85b63504856e137029b6acbc73762496064219cdb265d/boto3-1.40.61-py3-none-any.whl", hash = "sha256:6b9c57b2a922b5d8c17766e29ed792586a818098efe84def27c8f582b33f898c", size = 139321, upload-time = "2025-10-28T19:26:55.007Z" }, -] - -[[package]] -name = "botocore" -version = "1.40.61" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "jmespath" }, - { name = "python-dateutil" }, - { name = "urllib3" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/28/a3/81d3a47c2dbfd76f185d3b894f2ad01a75096c006a2dd91f237dca182188/botocore-1.40.61.tar.gz", hash = "sha256:a2487ad69b090f9cccd64cf07c7021cd80ee9c0655ad974f87045b02f3ef52cd", size = 14393956, upload-time = "2025-10-28T19:26:46.108Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/38/c5/f6ce561004db45f0b847c2cd9b19c67c6bf348a82018a48cb718be6b58b0/botocore-1.40.61-py3-none-any.whl", hash = "sha256:17ebae412692fd4824f99cde0f08d50126dc97954008e5ba2b522eb049238aa7", size = 14055973, upload-time = "2025-10-28T19:26:42.15Z" }, -] - [[package]] name = "bpython" version = "0.26" @@ -1353,38 +1119,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/18/b8/68d6ca1d8a16061e79693587560f6d24ac18ba9617804d7808b2c988d9d5/deptry-0.24.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:03d375db3e56821803aeca665dbb4c2fd935024310350cc18e8d8b6421369d2b", size = 1629786, upload-time = "2025-11-09T00:31:49.469Z" }, ] -[[package]] -name = "dev-local-vectorstore-hello" -version = "0.2.0" -source = { editable = "samples/dev-local-vectorstore-hello" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-dev-local-vectorstore" }, - { name = "genkit-plugin-google-genai" }, - { name = "pydantic" }, - { name = "rich" }, - { name = "structlog" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-dev-local-vectorstore", editable = "plugins/dev-local-vectorstore" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "pydantic", specifier = ">=2.10.5" }, - { name = "rich", specifier = ">=13.0.0" }, - { name = "structlog", specifier = ">=25.2.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - [[package]] name = "distlib" version = "0.4.0" @@ -1579,7 +1313,7 @@ wheels = [ [[package]] name = "flask" -version = "3.1.2" +version = "3.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "blinker" }, @@ -1589,9 +1323,9 @@ dependencies = [ { name = "markupsafe" }, { name = "werkzeug" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/dc/6d/cfe3c0fcc5e477df242b98bfe186a4c34357b4847e87ecaef04507332dab/flask-3.1.2.tar.gz", hash = "sha256:bf656c15c80190ed628ad08cdfd3aaa35beb087855e2f494910aa3774cc4fd87", size = 720160, upload-time = "2025-08-19T21:03:21.205Z" } +sdist = { url = "https://files.pythonhosted.org/packages/26/00/35d85dcce6c57fdc871f3867d465d780f302a175ea360f62533f12b27e2b/flask-3.1.3.tar.gz", hash = "sha256:0ef0e52b8a9cd932855379197dd8f94047b359ca0a78695144304cb45f87c9eb", size = 759004, upload-time = "2026-02-19T05:00:57.678Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ec/f9/7f9263c5695f4bd0023734af91bedb2ff8209e8de6ead162f35d8dc762fd/flask-3.1.2-py3-none-any.whl", hash = "sha256:ca1d8112ec8a6158cc29ea4858963350011b5c846a414cdb7a954aa9e967d03c", size = 103308, upload-time = "2025-08-19T21:03:19.499Z" }, + { url = "https://files.pythonhosted.org/packages/7f/9c/34f6962f9b9e9c71f6e5ed806e0d0ff03c9d1b0b2340088a0cf4bce09b18/flask-3.1.3-py3-none-any.whl", hash = "sha256:f4bcbefc124291925f1a26446da31a5178f9483862233b23c0c96a20701f670c", size = 103424, upload-time = "2026-02-19T05:00:56.027Z" }, ] [[package]] @@ -1819,8 +1553,6 @@ version = "0.2.0" source = { editable = "samples/framework-restaurant-demo" } dependencies = [ { name = "genkit" }, - { name = "genkit-plugin-dev-local-vectorstore" }, - { name = "genkit-plugin-firebase" }, { name = "genkit-plugin-google-cloud" }, { name = "genkit-plugin-google-genai" }, { name = "genkit-plugin-ollama" }, @@ -1838,8 +1570,6 @@ dev = [ [package.metadata] requires-dist = [ { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-dev-local-vectorstore", editable = "plugins/dev-local-vectorstore" }, - { name = "genkit-plugin-firebase", editable = "plugins/firebase" }, { name = "genkit-plugin-google-cloud", editable = "plugins/google-cloud" }, { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, { name = "genkit-plugin-ollama", editable = "plugins/ollama" }, @@ -1879,127 +1609,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[[package]] -name = "frozenlist" -version = "1.8.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/2d/f5/c831fac6cc817d26fd54c7eaccd04ef7e0288806943f7cc5bbf69f3ac1f0/frozenlist-1.8.0.tar.gz", hash = "sha256:3ede829ed8d842f6cd48fc7081d7a41001a56f1f38603f9d49bf3020d59a31ad", size = 45875, upload-time = "2025-10-06T05:38:17.865Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/83/4a/557715d5047da48d54e659203b9335be7bfaafda2c3f627b7c47e0b3aaf3/frozenlist-1.8.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:b37f6d31b3dcea7deb5e9696e529a6aa4a898adc33db82da12e4c60a7c4d2011", size = 86230, upload-time = "2025-10-06T05:35:23.699Z" }, - { url = "https://files.pythonhosted.org/packages/a2/fb/c85f9fed3ea8fe8740e5b46a59cc141c23b842eca617da8876cfce5f760e/frozenlist-1.8.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ef2b7b394f208233e471abc541cc6991f907ffd47dc72584acee3147899d6565", size = 49621, upload-time = "2025-10-06T05:35:25.341Z" }, - { url = "https://files.pythonhosted.org/packages/63/70/26ca3f06aace16f2352796b08704338d74b6d1a24ca38f2771afbb7ed915/frozenlist-1.8.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a88f062f072d1589b7b46e951698950e7da00442fc1cacbe17e19e025dc327ad", size = 49889, upload-time = "2025-10-06T05:35:26.797Z" }, - { url = "https://files.pythonhosted.org/packages/5d/ed/c7895fd2fde7f3ee70d248175f9b6cdf792fb741ab92dc59cd9ef3bd241b/frozenlist-1.8.0-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f57fb59d9f385710aa7060e89410aeb5058b99e62f4d16b08b91986b9a2140c2", size = 219464, upload-time = "2025-10-06T05:35:28.254Z" }, - { url = "https://files.pythonhosted.org/packages/6b/83/4d587dccbfca74cb8b810472392ad62bfa100bf8108c7223eb4c4fa2f7b3/frozenlist-1.8.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:799345ab092bee59f01a915620b5d014698547afd011e691a208637312db9186", size = 221649, upload-time = "2025-10-06T05:35:29.454Z" }, - { url = "https://files.pythonhosted.org/packages/6a/c6/fd3b9cd046ec5fff9dab66831083bc2077006a874a2d3d9247dea93ddf7e/frozenlist-1.8.0-cp310-cp310-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:c23c3ff005322a6e16f71bf8692fcf4d5a304aaafe1e262c98c6d4adc7be863e", size = 219188, upload-time = "2025-10-06T05:35:30.951Z" }, - { url = "https://files.pythonhosted.org/packages/ce/80/6693f55eb2e085fc8afb28cf611448fb5b90e98e068fa1d1b8d8e66e5c7d/frozenlist-1.8.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8a76ea0f0b9dfa06f254ee06053d93a600865b3274358ca48a352ce4f0798450", size = 231748, upload-time = "2025-10-06T05:35:32.101Z" }, - { url = "https://files.pythonhosted.org/packages/97/d6/e9459f7c5183854abd989ba384fe0cc1a0fb795a83c033f0571ec5933ca4/frozenlist-1.8.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c7366fe1418a6133d5aa824ee53d406550110984de7637d65a178010f759c6ef", size = 236351, upload-time = "2025-10-06T05:35:33.834Z" }, - { url = "https://files.pythonhosted.org/packages/97/92/24e97474b65c0262e9ecd076e826bfd1d3074adcc165a256e42e7b8a7249/frozenlist-1.8.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:13d23a45c4cebade99340c4165bd90eeb4a56c6d8a9d8aa49568cac19a6d0dc4", size = 218767, upload-time = "2025-10-06T05:35:35.205Z" }, - { url = "https://files.pythonhosted.org/packages/ee/bf/dc394a097508f15abff383c5108cb8ad880d1f64a725ed3b90d5c2fbf0bb/frozenlist-1.8.0-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:e4a3408834f65da56c83528fb52ce7911484f0d1eaf7b761fc66001db1646eff", size = 235887, upload-time = "2025-10-06T05:35:36.354Z" }, - { url = "https://files.pythonhosted.org/packages/40/90/25b201b9c015dbc999a5baf475a257010471a1fa8c200c843fd4abbee725/frozenlist-1.8.0-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:42145cd2748ca39f32801dad54aeea10039da6f86e303659db90db1c4b614c8c", size = 228785, upload-time = "2025-10-06T05:35:37.949Z" }, - { url = "https://files.pythonhosted.org/packages/84/f4/b5bc148df03082f05d2dd30c089e269acdbe251ac9a9cf4e727b2dbb8a3d/frozenlist-1.8.0-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:e2de870d16a7a53901e41b64ffdf26f2fbb8917b3e6ebf398098d72c5b20bd7f", size = 230312, upload-time = "2025-10-06T05:35:39.178Z" }, - { url = "https://files.pythonhosted.org/packages/db/4b/87e95b5d15097c302430e647136b7d7ab2398a702390cf4c8601975709e7/frozenlist-1.8.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:20e63c9493d33ee48536600d1a5c95eefc870cd71e7ab037763d1fbb89cc51e7", size = 217650, upload-time = "2025-10-06T05:35:40.377Z" }, - { url = "https://files.pythonhosted.org/packages/e5/70/78a0315d1fea97120591a83e0acd644da638c872f142fd72a6cebee825f3/frozenlist-1.8.0-cp310-cp310-win32.whl", hash = "sha256:adbeebaebae3526afc3c96fad434367cafbfd1b25d72369a9e5858453b1bb71a", size = 39659, upload-time = "2025-10-06T05:35:41.863Z" }, - { url = "https://files.pythonhosted.org/packages/66/aa/3f04523fb189a00e147e60c5b2205126118f216b0aa908035c45336e27e4/frozenlist-1.8.0-cp310-cp310-win_amd64.whl", hash = "sha256:667c3777ca571e5dbeb76f331562ff98b957431df140b54c85fd4d52eea8d8f6", size = 43837, upload-time = "2025-10-06T05:35:43.205Z" }, - { url = "https://files.pythonhosted.org/packages/39/75/1135feecdd7c336938bd55b4dc3b0dfc46d85b9be12ef2628574b28de776/frozenlist-1.8.0-cp310-cp310-win_arm64.whl", hash = "sha256:80f85f0a7cc86e7a54c46d99c9e1318ff01f4687c172ede30fd52d19d1da1c8e", size = 39989, upload-time = "2025-10-06T05:35:44.596Z" }, - { url = "https://files.pythonhosted.org/packages/bc/03/077f869d540370db12165c0aa51640a873fb661d8b315d1d4d67b284d7ac/frozenlist-1.8.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:09474e9831bc2b2199fad6da3c14c7b0fbdd377cce9d3d77131be28906cb7d84", size = 86912, upload-time = "2025-10-06T05:35:45.98Z" }, - { url = "https://files.pythonhosted.org/packages/df/b5/7610b6bd13e4ae77b96ba85abea1c8cb249683217ef09ac9e0ae93f25a91/frozenlist-1.8.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:17c883ab0ab67200b5f964d2b9ed6b00971917d5d8a92df149dc2c9779208ee9", size = 50046, upload-time = "2025-10-06T05:35:47.009Z" }, - { url = "https://files.pythonhosted.org/packages/6e/ef/0e8f1fe32f8a53dd26bdd1f9347efe0778b0fddf62789ea683f4cc7d787d/frozenlist-1.8.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:fa47e444b8ba08fffd1c18e8cdb9a75db1b6a27f17507522834ad13ed5922b93", size = 50119, upload-time = "2025-10-06T05:35:48.38Z" }, - { url = "https://files.pythonhosted.org/packages/11/b1/71a477adc7c36e5fb628245dfbdea2166feae310757dea848d02bd0689fd/frozenlist-1.8.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2552f44204b744fba866e573be4c1f9048d6a324dfe14475103fd51613eb1d1f", size = 231067, upload-time = "2025-10-06T05:35:49.97Z" }, - { url = "https://files.pythonhosted.org/packages/45/7e/afe40eca3a2dc19b9904c0f5d7edfe82b5304cb831391edec0ac04af94c2/frozenlist-1.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:957e7c38f250991e48a9a73e6423db1bb9dd14e722a10f6b8bb8e16a0f55f695", size = 233160, upload-time = "2025-10-06T05:35:51.729Z" }, - { url = "https://files.pythonhosted.org/packages/a6/aa/7416eac95603ce428679d273255ffc7c998d4132cfae200103f164b108aa/frozenlist-1.8.0-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:8585e3bb2cdea02fc88ffa245069c36555557ad3609e83be0ec71f54fd4abb52", size = 228544, upload-time = "2025-10-06T05:35:53.246Z" }, - { url = "https://files.pythonhosted.org/packages/8b/3d/2a2d1f683d55ac7e3875e4263d28410063e738384d3adc294f5ff3d7105e/frozenlist-1.8.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:edee74874ce20a373d62dc28b0b18b93f645633c2943fd90ee9d898550770581", size = 243797, upload-time = "2025-10-06T05:35:54.497Z" }, - { url = "https://files.pythonhosted.org/packages/78/1e/2d5565b589e580c296d3bb54da08d206e797d941a83a6fdea42af23be79c/frozenlist-1.8.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c9a63152fe95756b85f31186bddf42e4c02c6321207fd6601a1c89ebac4fe567", size = 247923, upload-time = "2025-10-06T05:35:55.861Z" }, - { url = "https://files.pythonhosted.org/packages/aa/c3/65872fcf1d326a7f101ad4d86285c403c87be7d832b7470b77f6d2ed5ddc/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:b6db2185db9be0a04fecf2f241c70b63b1a242e2805be291855078f2b404dd6b", size = 230886, upload-time = "2025-10-06T05:35:57.399Z" }, - { url = "https://files.pythonhosted.org/packages/a0/76/ac9ced601d62f6956f03cc794f9e04c81719509f85255abf96e2510f4265/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:f4be2e3d8bc8aabd566f8d5b8ba7ecc09249d74ba3c9ed52e54dc23a293f0b92", size = 245731, upload-time = "2025-10-06T05:35:58.563Z" }, - { url = "https://files.pythonhosted.org/packages/b9/49/ecccb5f2598daf0b4a1415497eba4c33c1e8ce07495eb07d2860c731b8d5/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:c8d1634419f39ea6f5c427ea2f90ca85126b54b50837f31497f3bf38266e853d", size = 241544, upload-time = "2025-10-06T05:35:59.719Z" }, - { url = "https://files.pythonhosted.org/packages/53/4b/ddf24113323c0bbcc54cb38c8b8916f1da7165e07b8e24a717b4a12cbf10/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:1a7fa382a4a223773ed64242dbe1c9c326ec09457e6b8428efb4118c685c3dfd", size = 241806, upload-time = "2025-10-06T05:36:00.959Z" }, - { url = "https://files.pythonhosted.org/packages/a7/fb/9b9a084d73c67175484ba2789a59f8eebebd0827d186a8102005ce41e1ba/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:11847b53d722050808926e785df837353bd4d75f1d494377e59b23594d834967", size = 229382, upload-time = "2025-10-06T05:36:02.22Z" }, - { url = "https://files.pythonhosted.org/packages/95/a3/c8fb25aac55bf5e12dae5c5aa6a98f85d436c1dc658f21c3ac73f9fa95e5/frozenlist-1.8.0-cp311-cp311-win32.whl", hash = "sha256:27c6e8077956cf73eadd514be8fb04d77fc946a7fe9f7fe167648b0b9085cc25", size = 39647, upload-time = "2025-10-06T05:36:03.409Z" }, - { url = "https://files.pythonhosted.org/packages/0a/f5/603d0d6a02cfd4c8f2a095a54672b3cf967ad688a60fb9faf04fc4887f65/frozenlist-1.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:ac913f8403b36a2c8610bbfd25b8013488533e71e62b4b4adce9c86c8cea905b", size = 44064, upload-time = "2025-10-06T05:36:04.368Z" }, - { url = "https://files.pythonhosted.org/packages/5d/16/c2c9ab44e181f043a86f9a8f84d5124b62dbcb3a02c0977ec72b9ac1d3e0/frozenlist-1.8.0-cp311-cp311-win_arm64.whl", hash = "sha256:d4d3214a0f8394edfa3e303136d0575eece0745ff2b47bd2cb2e66dd92d4351a", size = 39937, upload-time = "2025-10-06T05:36:05.669Z" }, - { url = "https://files.pythonhosted.org/packages/69/29/948b9aa87e75820a38650af445d2ef2b6b8a6fab1a23b6bb9e4ef0be2d59/frozenlist-1.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:78f7b9e5d6f2fdb88cdde9440dc147259b62b9d3b019924def9f6478be254ac1", size = 87782, upload-time = "2025-10-06T05:36:06.649Z" }, - { url = "https://files.pythonhosted.org/packages/64/80/4f6e318ee2a7c0750ed724fa33a4bdf1eacdc5a39a7a24e818a773cd91af/frozenlist-1.8.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:229bf37d2e4acdaf808fd3f06e854a4a7a3661e871b10dc1f8f1896a3b05f18b", size = 50594, upload-time = "2025-10-06T05:36:07.69Z" }, - { url = "https://files.pythonhosted.org/packages/2b/94/5c8a2b50a496b11dd519f4a24cb5496cf125681dd99e94c604ccdea9419a/frozenlist-1.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f833670942247a14eafbb675458b4e61c82e002a148f49e68257b79296e865c4", size = 50448, upload-time = "2025-10-06T05:36:08.78Z" }, - { url = "https://files.pythonhosted.org/packages/6a/bd/d91c5e39f490a49df14320f4e8c80161cfcce09f1e2cde1edd16a551abb3/frozenlist-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:494a5952b1c597ba44e0e78113a7266e656b9794eec897b19ead706bd7074383", size = 242411, upload-time = "2025-10-06T05:36:09.801Z" }, - { url = "https://files.pythonhosted.org/packages/8f/83/f61505a05109ef3293dfb1ff594d13d64a2324ac3482be2cedc2be818256/frozenlist-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96f423a119f4777a4a056b66ce11527366a8bb92f54e541ade21f2374433f6d4", size = 243014, upload-time = "2025-10-06T05:36:11.394Z" }, - { url = "https://files.pythonhosted.org/packages/d8/cb/cb6c7b0f7d4023ddda30cf56b8b17494eb3a79e3fda666bf735f63118b35/frozenlist-1.8.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3462dd9475af2025c31cc61be6652dfa25cbfb56cbbf52f4ccfe029f38decaf8", size = 234909, upload-time = "2025-10-06T05:36:12.598Z" }, - { url = "https://files.pythonhosted.org/packages/31/c5/cd7a1f3b8b34af009fb17d4123c5a778b44ae2804e3ad6b86204255f9ec5/frozenlist-1.8.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c4c800524c9cd9bac5166cd6f55285957fcfc907db323e193f2afcd4d9abd69b", size = 250049, upload-time = "2025-10-06T05:36:14.065Z" }, - { url = "https://files.pythonhosted.org/packages/c0/01/2f95d3b416c584a1e7f0e1d6d31998c4a795f7544069ee2e0962a4b60740/frozenlist-1.8.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d6a5df73acd3399d893dafc71663ad22534b5aa4f94e8a2fabfe856c3c1b6a52", size = 256485, upload-time = "2025-10-06T05:36:15.39Z" }, - { url = "https://files.pythonhosted.org/packages/ce/03/024bf7720b3abaebcff6d0793d73c154237b85bdf67b7ed55e5e9596dc9a/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:405e8fe955c2280ce66428b3ca55e12b3c4e9c336fb2103a4937e891c69a4a29", size = 237619, upload-time = "2025-10-06T05:36:16.558Z" }, - { url = "https://files.pythonhosted.org/packages/69/fa/f8abdfe7d76b731f5d8bd217827cf6764d4f1d9763407e42717b4bed50a0/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:908bd3f6439f2fef9e85031b59fd4f1297af54415fb60e4254a95f75b3cab3f3", size = 250320, upload-time = "2025-10-06T05:36:17.821Z" }, - { url = "https://files.pythonhosted.org/packages/f5/3c/b051329f718b463b22613e269ad72138cc256c540f78a6de89452803a47d/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:294e487f9ec720bd8ffcebc99d575f7eff3568a08a253d1ee1a0378754b74143", size = 246820, upload-time = "2025-10-06T05:36:19.046Z" }, - { url = "https://files.pythonhosted.org/packages/0f/ae/58282e8f98e444b3f4dd42448ff36fa38bef29e40d40f330b22e7108f565/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:74c51543498289c0c43656701be6b077f4b265868fa7f8a8859c197006efb608", size = 250518, upload-time = "2025-10-06T05:36:20.763Z" }, - { url = "https://files.pythonhosted.org/packages/8f/96/007e5944694d66123183845a106547a15944fbbb7154788cbf7272789536/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:776f352e8329135506a1d6bf16ac3f87bc25b28e765949282dcc627af36123aa", size = 239096, upload-time = "2025-10-06T05:36:22.129Z" }, - { url = "https://files.pythonhosted.org/packages/66/bb/852b9d6db2fa40be96f29c0d1205c306288f0684df8fd26ca1951d461a56/frozenlist-1.8.0-cp312-cp312-win32.whl", hash = "sha256:433403ae80709741ce34038da08511d4a77062aa924baf411ef73d1146e74faf", size = 39985, upload-time = "2025-10-06T05:36:23.661Z" }, - { url = "https://files.pythonhosted.org/packages/b8/af/38e51a553dd66eb064cdf193841f16f077585d4d28394c2fa6235cb41765/frozenlist-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:34187385b08f866104f0c0617404c8eb08165ab1272e884abc89c112e9c00746", size = 44591, upload-time = "2025-10-06T05:36:24.958Z" }, - { url = "https://files.pythonhosted.org/packages/a7/06/1dc65480ab147339fecc70797e9c2f69d9cea9cf38934ce08df070fdb9cb/frozenlist-1.8.0-cp312-cp312-win_arm64.whl", hash = "sha256:fe3c58d2f5db5fbd18c2987cba06d51b0529f52bc3a6cdc33d3f4eab725104bd", size = 40102, upload-time = "2025-10-06T05:36:26.333Z" }, - { url = "https://files.pythonhosted.org/packages/2d/40/0832c31a37d60f60ed79e9dfb5a92e1e2af4f40a16a29abcc7992af9edff/frozenlist-1.8.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8d92f1a84bb12d9e56f818b3a746f3efba93c1b63c8387a73dde655e1e42282a", size = 85717, upload-time = "2025-10-06T05:36:27.341Z" }, - { url = "https://files.pythonhosted.org/packages/30/ba/b0b3de23f40bc55a7057bd38434e25c34fa48e17f20ee273bbde5e0650f3/frozenlist-1.8.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:96153e77a591c8adc2ee805756c61f59fef4cf4073a9275ee86fe8cba41241f7", size = 49651, upload-time = "2025-10-06T05:36:28.855Z" }, - { url = "https://files.pythonhosted.org/packages/0c/ab/6e5080ee374f875296c4243c381bbdef97a9ac39c6e3ce1d5f7d42cb78d6/frozenlist-1.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f21f00a91358803399890ab167098c131ec2ddd5f8f5fd5fe9c9f2c6fcd91e40", size = 49417, upload-time = "2025-10-06T05:36:29.877Z" }, - { url = "https://files.pythonhosted.org/packages/d5/4e/e4691508f9477ce67da2015d8c00acd751e6287739123113a9fca6f1604e/frozenlist-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:fb30f9626572a76dfe4293c7194a09fb1fe93ba94c7d4f720dfae3b646b45027", size = 234391, upload-time = "2025-10-06T05:36:31.301Z" }, - { url = "https://files.pythonhosted.org/packages/40/76/c202df58e3acdf12969a7895fd6f3bc016c642e6726aa63bd3025e0fc71c/frozenlist-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eaa352d7047a31d87dafcacbabe89df0aa506abb5b1b85a2fb91bc3faa02d822", size = 233048, upload-time = "2025-10-06T05:36:32.531Z" }, - { url = "https://files.pythonhosted.org/packages/f9/c0/8746afb90f17b73ca5979c7a3958116e105ff796e718575175319b5bb4ce/frozenlist-1.8.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:03ae967b4e297f58f8c774c7eabcce57fe3c2434817d4385c50661845a058121", size = 226549, upload-time = "2025-10-06T05:36:33.706Z" }, - { url = "https://files.pythonhosted.org/packages/7e/eb/4c7eefc718ff72f9b6c4893291abaae5fbc0c82226a32dcd8ef4f7a5dbef/frozenlist-1.8.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f6292f1de555ffcc675941d65fffffb0a5bcd992905015f85d0592201793e0e5", size = 239833, upload-time = "2025-10-06T05:36:34.947Z" }, - { url = "https://files.pythonhosted.org/packages/c2/4e/e5c02187cf704224f8b21bee886f3d713ca379535f16893233b9d672ea71/frozenlist-1.8.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:29548f9b5b5e3460ce7378144c3010363d8035cea44bc0bf02d57f5a685e084e", size = 245363, upload-time = "2025-10-06T05:36:36.534Z" }, - { url = "https://files.pythonhosted.org/packages/1f/96/cb85ec608464472e82ad37a17f844889c36100eed57bea094518bf270692/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ec3cc8c5d4084591b4237c0a272cc4f50a5b03396a47d9caaf76f5d7b38a4f11", size = 229314, upload-time = "2025-10-06T05:36:38.582Z" }, - { url = "https://files.pythonhosted.org/packages/5d/6f/4ae69c550e4cee66b57887daeebe006fe985917c01d0fff9caab9883f6d0/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:517279f58009d0b1f2e7c1b130b377a349405da3f7621ed6bfae50b10adf20c1", size = 243365, upload-time = "2025-10-06T05:36:40.152Z" }, - { url = "https://files.pythonhosted.org/packages/7a/58/afd56de246cf11780a40a2c28dc7cbabbf06337cc8ddb1c780a2d97e88d8/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:db1e72ede2d0d7ccb213f218df6a078a9c09a7de257c2fe8fcef16d5925230b1", size = 237763, upload-time = "2025-10-06T05:36:41.355Z" }, - { url = "https://files.pythonhosted.org/packages/cb/36/cdfaf6ed42e2644740d4a10452d8e97fa1c062e2a8006e4b09f1b5fd7d63/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:b4dec9482a65c54a5044486847b8a66bf10c9cb4926d42927ec4e8fd5db7fed8", size = 240110, upload-time = "2025-10-06T05:36:42.716Z" }, - { url = "https://files.pythonhosted.org/packages/03/a8/9ea226fbefad669f11b52e864c55f0bd57d3c8d7eb07e9f2e9a0b39502e1/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:21900c48ae04d13d416f0e1e0c4d81f7931f73a9dfa0b7a8746fb2fe7dd970ed", size = 233717, upload-time = "2025-10-06T05:36:44.251Z" }, - { url = "https://files.pythonhosted.org/packages/1e/0b/1b5531611e83ba7d13ccc9988967ea1b51186af64c42b7a7af465dcc9568/frozenlist-1.8.0-cp313-cp313-win32.whl", hash = "sha256:8b7b94a067d1c504ee0b16def57ad5738701e4ba10cec90529f13fa03c833496", size = 39628, upload-time = "2025-10-06T05:36:45.423Z" }, - { url = "https://files.pythonhosted.org/packages/d8/cf/174c91dbc9cc49bc7b7aab74d8b734e974d1faa8f191c74af9b7e80848e6/frozenlist-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:878be833caa6a3821caf85eb39c5ba92d28e85df26d57afb06b35b2efd937231", size = 43882, upload-time = "2025-10-06T05:36:46.796Z" }, - { url = "https://files.pythonhosted.org/packages/c1/17/502cd212cbfa96eb1388614fe39a3fc9ab87dbbe042b66f97acb57474834/frozenlist-1.8.0-cp313-cp313-win_arm64.whl", hash = "sha256:44389d135b3ff43ba8cc89ff7f51f5a0bb6b63d829c8300f79a2fe4fe61bcc62", size = 39676, upload-time = "2025-10-06T05:36:47.8Z" }, - { url = "https://files.pythonhosted.org/packages/d2/5c/3bbfaa920dfab09e76946a5d2833a7cbdf7b9b4a91c714666ac4855b88b4/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e25ac20a2ef37e91c1b39938b591457666a0fa835c7783c3a8f33ea42870db94", size = 89235, upload-time = "2025-10-06T05:36:48.78Z" }, - { url = "https://files.pythonhosted.org/packages/d2/d6/f03961ef72166cec1687e84e8925838442b615bd0b8854b54923ce5b7b8a/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:07cdca25a91a4386d2e76ad992916a85038a9b97561bf7a3fd12d5d9ce31870c", size = 50742, upload-time = "2025-10-06T05:36:49.837Z" }, - { url = "https://files.pythonhosted.org/packages/1e/bb/a6d12b7ba4c3337667d0e421f7181c82dda448ce4e7ad7ecd249a16fa806/frozenlist-1.8.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4e0c11f2cc6717e0a741f84a527c52616140741cd812a50422f83dc31749fb52", size = 51725, upload-time = "2025-10-06T05:36:50.851Z" }, - { url = "https://files.pythonhosted.org/packages/bc/71/d1fed0ffe2c2ccd70b43714c6cab0f4188f09f8a67a7914a6b46ee30f274/frozenlist-1.8.0-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b3210649ee28062ea6099cfda39e147fa1bc039583c8ee4481cb7811e2448c51", size = 284533, upload-time = "2025-10-06T05:36:51.898Z" }, - { url = "https://files.pythonhosted.org/packages/c9/1f/fb1685a7b009d89f9bf78a42d94461bc06581f6e718c39344754a5d9bada/frozenlist-1.8.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:581ef5194c48035a7de2aefc72ac6539823bb71508189e5de01d60c9dcd5fa65", size = 292506, upload-time = "2025-10-06T05:36:53.101Z" }, - { url = "https://files.pythonhosted.org/packages/e6/3b/b991fe1612703f7e0d05c0cf734c1b77aaf7c7d321df4572e8d36e7048c8/frozenlist-1.8.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3ef2d026f16a2b1866e1d86fc4e1291e1ed8a387b2c333809419a2f8b3a77b82", size = 274161, upload-time = "2025-10-06T05:36:54.309Z" }, - { url = "https://files.pythonhosted.org/packages/ca/ec/c5c618767bcdf66e88945ec0157d7f6c4a1322f1473392319b7a2501ded7/frozenlist-1.8.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5500ef82073f599ac84d888e3a8c1f77ac831183244bfd7f11eaa0289fb30714", size = 294676, upload-time = "2025-10-06T05:36:55.566Z" }, - { url = "https://files.pythonhosted.org/packages/7c/ce/3934758637d8f8a88d11f0585d6495ef54b2044ed6ec84492a91fa3b27aa/frozenlist-1.8.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:50066c3997d0091c411a66e710f4e11752251e6d2d73d70d8d5d4c76442a199d", size = 300638, upload-time = "2025-10-06T05:36:56.758Z" }, - { url = "https://files.pythonhosted.org/packages/fc/4f/a7e4d0d467298f42de4b41cbc7ddaf19d3cfeabaf9ff97c20c6c7ee409f9/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:5c1c8e78426e59b3f8005e9b19f6ff46e5845895adbde20ece9218319eca6506", size = 283067, upload-time = "2025-10-06T05:36:57.965Z" }, - { url = "https://files.pythonhosted.org/packages/dc/48/c7b163063d55a83772b268e6d1affb960771b0e203b632cfe09522d67ea5/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:eefdba20de0d938cec6a89bd4d70f346a03108a19b9df4248d3cf0d88f1b0f51", size = 292101, upload-time = "2025-10-06T05:36:59.237Z" }, - { url = "https://files.pythonhosted.org/packages/9f/d0/2366d3c4ecdc2fd391e0afa6e11500bfba0ea772764d631bbf82f0136c9d/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:cf253e0e1c3ceb4aaff6df637ce033ff6535fb8c70a764a8f46aafd3d6ab798e", size = 289901, upload-time = "2025-10-06T05:37:00.811Z" }, - { url = "https://files.pythonhosted.org/packages/b8/94/daff920e82c1b70e3618a2ac39fbc01ae3e2ff6124e80739ce5d71c9b920/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:032efa2674356903cd0261c4317a561a6850f3ac864a63fc1583147fb05a79b0", size = 289395, upload-time = "2025-10-06T05:37:02.115Z" }, - { url = "https://files.pythonhosted.org/packages/e3/20/bba307ab4235a09fdcd3cc5508dbabd17c4634a1af4b96e0f69bfe551ebd/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6da155091429aeba16851ecb10a9104a108bcd32f6c1642867eadaee401c1c41", size = 283659, upload-time = "2025-10-06T05:37:03.711Z" }, - { url = "https://files.pythonhosted.org/packages/fd/00/04ca1c3a7a124b6de4f8a9a17cc2fcad138b4608e7a3fc5877804b8715d7/frozenlist-1.8.0-cp313-cp313t-win32.whl", hash = "sha256:0f96534f8bfebc1a394209427d0f8a63d343c9779cda6fc25e8e121b5fd8555b", size = 43492, upload-time = "2025-10-06T05:37:04.915Z" }, - { url = "https://files.pythonhosted.org/packages/59/5e/c69f733a86a94ab10f68e496dc6b7e8bc078ebb415281d5698313e3af3a1/frozenlist-1.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5d63a068f978fc69421fb0e6eb91a9603187527c86b7cd3f534a5b77a592b888", size = 48034, upload-time = "2025-10-06T05:37:06.343Z" }, - { url = "https://files.pythonhosted.org/packages/16/6c/be9d79775d8abe79b05fa6d23da99ad6e7763a1d080fbae7290b286093fd/frozenlist-1.8.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bf0a7e10b077bf5fb9380ad3ae8ce20ef919a6ad93b4552896419ac7e1d8e042", size = 41749, upload-time = "2025-10-06T05:37:07.431Z" }, - { url = "https://files.pythonhosted.org/packages/f1/c8/85da824b7e7b9b6e7f7705b2ecaf9591ba6f79c1177f324c2735e41d36a2/frozenlist-1.8.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:cee686f1f4cadeb2136007ddedd0aaf928ab95216e7691c63e50a8ec066336d0", size = 86127, upload-time = "2025-10-06T05:37:08.438Z" }, - { url = "https://files.pythonhosted.org/packages/8e/e8/a1185e236ec66c20afd72399522f142c3724c785789255202d27ae992818/frozenlist-1.8.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:119fb2a1bd47307e899c2fac7f28e85b9a543864df47aa7ec9d3c1b4545f096f", size = 49698, upload-time = "2025-10-06T05:37:09.48Z" }, - { url = "https://files.pythonhosted.org/packages/a1/93/72b1736d68f03fda5fdf0f2180fb6caaae3894f1b854d006ac61ecc727ee/frozenlist-1.8.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4970ece02dbc8c3a92fcc5228e36a3e933a01a999f7094ff7c23fbd2beeaa67c", size = 49749, upload-time = "2025-10-06T05:37:10.569Z" }, - { url = "https://files.pythonhosted.org/packages/a7/b2/fabede9fafd976b991e9f1b9c8c873ed86f202889b864756f240ce6dd855/frozenlist-1.8.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:cba69cb73723c3f329622e34bdbf5ce1f80c21c290ff04256cff1cd3c2036ed2", size = 231298, upload-time = "2025-10-06T05:37:11.993Z" }, - { url = "https://files.pythonhosted.org/packages/3a/3b/d9b1e0b0eed36e70477ffb8360c49c85c8ca8ef9700a4e6711f39a6e8b45/frozenlist-1.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:778a11b15673f6f1df23d9586f83c4846c471a8af693a22e066508b77d201ec8", size = 232015, upload-time = "2025-10-06T05:37:13.194Z" }, - { url = "https://files.pythonhosted.org/packages/dc/94/be719d2766c1138148564a3960fc2c06eb688da592bdc25adcf856101be7/frozenlist-1.8.0-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0325024fe97f94c41c08872db482cf8ac4800d80e79222c6b0b7b162d5b13686", size = 225038, upload-time = "2025-10-06T05:37:14.577Z" }, - { url = "https://files.pythonhosted.org/packages/e4/09/6712b6c5465f083f52f50cf74167b92d4ea2f50e46a9eea0523d658454ae/frozenlist-1.8.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:97260ff46b207a82a7567b581ab4190bd4dfa09f4db8a8b49d1a958f6aa4940e", size = 240130, upload-time = "2025-10-06T05:37:15.781Z" }, - { url = "https://files.pythonhosted.org/packages/f8/d4/cd065cdcf21550b54f3ce6a22e143ac9e4836ca42a0de1022da8498eac89/frozenlist-1.8.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:54b2077180eb7f83dd52c40b2750d0a9f175e06a42e3213ce047219de902717a", size = 242845, upload-time = "2025-10-06T05:37:17.037Z" }, - { url = "https://files.pythonhosted.org/packages/62/c3/f57a5c8c70cd1ead3d5d5f776f89d33110b1addae0ab010ad774d9a44fb9/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2f05983daecab868a31e1da44462873306d3cbfd76d1f0b5b69c473d21dbb128", size = 229131, upload-time = "2025-10-06T05:37:18.221Z" }, - { url = "https://files.pythonhosted.org/packages/6c/52/232476fe9cb64f0742f3fde2b7d26c1dac18b6d62071c74d4ded55e0ef94/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:33f48f51a446114bc5d251fb2954ab0164d5be02ad3382abcbfe07e2531d650f", size = 240542, upload-time = "2025-10-06T05:37:19.771Z" }, - { url = "https://files.pythonhosted.org/packages/5f/85/07bf3f5d0fb5414aee5f47d33c6f5c77bfe49aac680bfece33d4fdf6a246/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:154e55ec0655291b5dd1b8731c637ecdb50975a2ae70c606d100750a540082f7", size = 237308, upload-time = "2025-10-06T05:37:20.969Z" }, - { url = "https://files.pythonhosted.org/packages/11/99/ae3a33d5befd41ac0ca2cc7fd3aa707c9c324de2e89db0e0f45db9a64c26/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:4314debad13beb564b708b4a496020e5306c7333fa9a3ab90374169a20ffab30", size = 238210, upload-time = "2025-10-06T05:37:22.252Z" }, - { url = "https://files.pythonhosted.org/packages/b2/60/b1d2da22f4970e7a155f0adde9b1435712ece01b3cd45ba63702aea33938/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:073f8bf8becba60aa931eb3bc420b217bb7d5b8f4750e6f8b3be7f3da85d38b7", size = 231972, upload-time = "2025-10-06T05:37:23.5Z" }, - { url = "https://files.pythonhosted.org/packages/3f/ab/945b2f32de889993b9c9133216c068b7fcf257d8595a0ac420ac8677cab0/frozenlist-1.8.0-cp314-cp314-win32.whl", hash = "sha256:bac9c42ba2ac65ddc115d930c78d24ab8d4f465fd3fc473cdedfccadb9429806", size = 40536, upload-time = "2025-10-06T05:37:25.581Z" }, - { url = "https://files.pythonhosted.org/packages/59/ad/9caa9b9c836d9ad6f067157a531ac48b7d36499f5036d4141ce78c230b1b/frozenlist-1.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:3e0761f4d1a44f1d1a47996511752cf3dcec5bbdd9cc2b4fe595caf97754b7a0", size = 44330, upload-time = "2025-10-06T05:37:26.928Z" }, - { url = "https://files.pythonhosted.org/packages/82/13/e6950121764f2676f43534c555249f57030150260aee9dcf7d64efda11dd/frozenlist-1.8.0-cp314-cp314-win_arm64.whl", hash = "sha256:d1eaff1d00c7751b7c6662e9c5ba6eb2c17a2306ba5e2a37f24ddf3cc953402b", size = 40627, upload-time = "2025-10-06T05:37:28.075Z" }, - { url = "https://files.pythonhosted.org/packages/c0/c7/43200656ecc4e02d3f8bc248df68256cd9572b3f0017f0a0c4e93440ae23/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:d3bb933317c52d7ea5004a1c442eef86f426886fba134ef8cf4226ea6ee1821d", size = 89238, upload-time = "2025-10-06T05:37:29.373Z" }, - { url = "https://files.pythonhosted.org/packages/d1/29/55c5f0689b9c0fb765055629f472c0de484dcaf0acee2f7707266ae3583c/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:8009897cdef112072f93a0efdce29cd819e717fd2f649ee3016efd3cd885a7ed", size = 50738, upload-time = "2025-10-06T05:37:30.792Z" }, - { url = "https://files.pythonhosted.org/packages/ba/7d/b7282a445956506fa11da8c2db7d276adcbf2b17d8bb8407a47685263f90/frozenlist-1.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2c5dcbbc55383e5883246d11fd179782a9d07a986c40f49abe89ddf865913930", size = 51739, upload-time = "2025-10-06T05:37:32.127Z" }, - { url = "https://files.pythonhosted.org/packages/62/1c/3d8622e60d0b767a5510d1d3cf21065b9db874696a51ea6d7a43180a259c/frozenlist-1.8.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:39ecbc32f1390387d2aa4f5a995e465e9e2f79ba3adcac92d68e3e0afae6657c", size = 284186, upload-time = "2025-10-06T05:37:33.21Z" }, - { url = "https://files.pythonhosted.org/packages/2d/14/aa36d5f85a89679a85a1d44cd7a6657e0b1c75f61e7cad987b203d2daca8/frozenlist-1.8.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92db2bf818d5cc8d9c1f1fc56b897662e24ea5adb36ad1f1d82875bd64e03c24", size = 292196, upload-time = "2025-10-06T05:37:36.107Z" }, - { url = "https://files.pythonhosted.org/packages/05/23/6bde59eb55abd407d34f77d39a5126fb7b4f109a3f611d3929f14b700c66/frozenlist-1.8.0-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2dc43a022e555de94c3b68a4ef0b11c4f747d12c024a520c7101709a2144fb37", size = 273830, upload-time = "2025-10-06T05:37:37.663Z" }, - { url = "https://files.pythonhosted.org/packages/d2/3f/22cff331bfad7a8afa616289000ba793347fcd7bc275f3b28ecea2a27909/frozenlist-1.8.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cb89a7f2de3602cfed448095bab3f178399646ab7c61454315089787df07733a", size = 294289, upload-time = "2025-10-06T05:37:39.261Z" }, - { url = "https://files.pythonhosted.org/packages/a4/89/5b057c799de4838b6c69aa82b79705f2027615e01be996d2486a69ca99c4/frozenlist-1.8.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:33139dc858c580ea50e7e60a1b0ea003efa1fd42e6ec7fdbad78fff65fad2fd2", size = 300318, upload-time = "2025-10-06T05:37:43.213Z" }, - { url = "https://files.pythonhosted.org/packages/30/de/2c22ab3eb2a8af6d69dc799e48455813bab3690c760de58e1bf43b36da3e/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:168c0969a329b416119507ba30b9ea13688fafffac1b7822802537569a1cb0ef", size = 282814, upload-time = "2025-10-06T05:37:45.337Z" }, - { url = "https://files.pythonhosted.org/packages/59/f7/970141a6a8dbd7f556d94977858cfb36fa9b66e0892c6dd780d2219d8cd8/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:28bd570e8e189d7f7b001966435f9dac6718324b5be2990ac496cf1ea9ddb7fe", size = 291762, upload-time = "2025-10-06T05:37:46.657Z" }, - { url = "https://files.pythonhosted.org/packages/c1/15/ca1adae83a719f82df9116d66f5bb28bb95557b3951903d39135620ef157/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:b2a095d45c5d46e5e79ba1e5b9cb787f541a8dee0433836cea4b96a2c439dcd8", size = 289470, upload-time = "2025-10-06T05:37:47.946Z" }, - { url = "https://files.pythonhosted.org/packages/ac/83/dca6dc53bf657d371fbc88ddeb21b79891e747189c5de990b9dfff2ccba1/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:eab8145831a0d56ec9c4139b6c3e594c7a83c2c8be25d5bcf2d86136a532287a", size = 289042, upload-time = "2025-10-06T05:37:49.499Z" }, - { url = "https://files.pythonhosted.org/packages/96/52/abddd34ca99be142f354398700536c5bd315880ed0a213812bc491cff5e4/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:974b28cf63cc99dfb2188d8d222bc6843656188164848c4f679e63dae4b0708e", size = 283148, upload-time = "2025-10-06T05:37:50.745Z" }, - { url = "https://files.pythonhosted.org/packages/af/d3/76bd4ed4317e7119c2b7f57c3f6934aba26d277acc6309f873341640e21f/frozenlist-1.8.0-cp314-cp314t-win32.whl", hash = "sha256:342c97bf697ac5480c0a7ec73cd700ecfa5a8a40ac923bd035484616efecc2df", size = 44676, upload-time = "2025-10-06T05:37:52.222Z" }, - { url = "https://files.pythonhosted.org/packages/89/76/c615883b7b521ead2944bb3480398cbb07e12b7b4e4d073d3752eb721558/frozenlist-1.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:06be8f67f39c8b1dc671f5d83aaefd3358ae5cdcf8314552c57e7ed3e6475bdd", size = 49451, upload-time = "2025-10-06T05:37:53.425Z" }, - { url = "https://files.pythonhosted.org/packages/e0/a3/5982da14e113d07b325230f95060e2169f5311b1017ea8af2a29b374c289/frozenlist-1.8.0-cp314-cp314t-win_arm64.whl", hash = "sha256:102e6314ca4da683dca92e3b1355490fed5f313b768500084fbe6371fddfdb79", size = 42507, upload-time = "2025-10-06T05:37:54.513Z" }, - { url = "https://files.pythonhosted.org/packages/9a/9a/e35b4a917281c0b8419d4207f4334c8e8c5dbf4f3f5f9ada73958d937dcc/frozenlist-1.8.0-py3-none-any.whl", hash = "sha256:0c18a16eab41e82c295618a77502e17b195883241c563b00f0aa5106fc4eaa0d", size = 13409, upload-time = "2025-10-06T05:38:16.721Z" }, -] - [[package]] name = "genkit" version = "0.5.1" @@ -2029,9 +1638,6 @@ dependencies = [ ] [package.optional-dependencies] -dev-local-vectorstore = [ - { name = "genkit-plugin-dev-local-vectorstore" }, -] flask = [ { name = "genkit-plugin-flask" }, ] @@ -2057,7 +1663,6 @@ requires-dist = [ { name = "asgiref", specifier = ">=3.8.1" }, { name = "dotpromptz", specifier = ">=0.1.5" }, { name = "genkit-plugin-compat-oai", marker = "extra == 'openai'", editable = "plugins/compat-oai" }, - { name = "genkit-plugin-dev-local-vectorstore", marker = "extra == 'dev-local-vectorstore'", editable = "plugins/dev-local-vectorstore" }, { name = "genkit-plugin-flask", marker = "extra == 'flask'", editable = "plugins/flask" }, { name = "genkit-plugin-google-cloud", marker = "extra == 'google-cloud'", editable = "plugins/google-cloud" }, { name = "genkit-plugin-google-genai", marker = "extra == 'google-genai'", editable = "plugins/google-genai" }, @@ -2082,34 +1687,7 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.34.0" }, { name = "uvloop", marker = "sys_platform != 'win32'", specifier = ">=0.21.0" }, ] -provides-extras = ["dev-local-vectorstore", "flask", "google-cloud", "google-genai", "ollama", "openai", "vertex-ai"] - -[[package]] -name = "genkit-plugin-amazon-bedrock" -version = "0.5.1" -source = { editable = "plugins/amazon-bedrock" } -dependencies = [ - { name = "aioboto3" }, - { name = "boto3" }, - { name = "botocore" }, - { name = "genkit" }, - { name = "opentelemetry-exporter-otlp-proto-http" }, - { name = "opentelemetry-propagator-aws-xray" }, - { name = "opentelemetry-sdk-extension-aws" }, - { name = "strenum", marker = "python_full_version < '3.11'" }, -] - -[package.metadata] -requires-dist = [ - { name = "aioboto3", specifier = ">=13.0.0" }, - { name = "boto3", specifier = ">=1.35.0" }, - { name = "botocore", specifier = ">=1.35.0" }, - { name = "genkit", editable = "packages/genkit" }, - { name = "opentelemetry-exporter-otlp-proto-http", specifier = ">=1.20.0" }, - { name = "opentelemetry-propagator-aws-xray", specifier = ">=1.0.0" }, - { name = "opentelemetry-sdk-extension-aws", specifier = ">=2.1.0" }, - { name = "strenum", marker = "python_full_version < '3.11'", specifier = ">=0.4.15" }, -] +provides-extras = ["flask", "google-cloud", "google-genai", "ollama", "openai", "vertex-ai"] [[package]] name = "genkit-plugin-anthropic" @@ -2143,27 +1721,6 @@ requires-dist = [ { name = "strenum", marker = "python_full_version < '3.11'", specifier = ">=0.4.15" }, ] -[[package]] -name = "genkit-plugin-dev-local-vectorstore" -version = "0.5.1" -source = { editable = "plugins/dev-local-vectorstore" } -dependencies = [ - { name = "aiofiles" }, - { name = "genkit" }, - { name = "pytest-asyncio" }, - { name = "pytest-mock" }, - { name = "strenum", marker = "python_full_version < '3.11'" }, -] - -[package.metadata] -requires-dist = [ - { name = "aiofiles", specifier = ">=24.1.0" }, - { name = "genkit", editable = "packages/genkit" }, - { name = "pytest-asyncio" }, - { name = "pytest-mock" }, - { name = "strenum", marker = "python_full_version < '3.11'", specifier = ">=0.4.15" }, -] - [[package]] name = "genkit-plugin-fastapi" version = "0.5.1" @@ -2181,32 +1738,6 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.10.5" }, ] -[[package]] -name = "genkit-plugin-firebase" -version = "0.5.1" -source = { editable = "plugins/firebase" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-google-cloud" }, - { name = "google-cloud-firestore" }, - { name = "strenum", marker = "python_full_version < '3.11'" }, -] - -[package.optional-dependencies] -telemetry = [ - { name = "genkit-plugin-google-cloud" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-google-cloud", editable = "plugins/google-cloud" }, - { name = "genkit-plugin-google-cloud", marker = "extra == 'telemetry'", editable = "plugins/google-cloud" }, - { name = "google-cloud-firestore" }, - { name = "strenum", marker = "python_full_version < '3.11'", specifier = ">=0.4.15" }, -] -provides-extras = ["telemetry"] - [[package]] name = "genkit-plugin-flask" version = "0.5.1" @@ -2220,7 +1751,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "flask" }, + { name = "flask", specifier = ">=3.1.3" }, { name = "genkit", editable = "packages/genkit" }, { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, { name = "pydantic", specifier = ">=2.10.5" }, @@ -2322,7 +1853,6 @@ dependencies = [ { name = "datamodel-code-generator" }, { name = "fastapi" }, { name = "genkit" }, - { name = "genkit-plugin-amazon-bedrock" }, { name = "genkit-plugin-anthropic" }, { name = "genkit-plugin-google-cloud" }, { name = "genkit-plugin-google-genai" }, @@ -2353,7 +1883,6 @@ requires-dist = [ { name = "datamodel-code-generator" }, { name = "fastapi" }, { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-amazon-bedrock", editable = "plugins/amazon-bedrock" }, { name = "genkit-plugin-anthropic", editable = "plugins/anthropic" }, { name = "genkit-plugin-google-cloud", editable = "plugins/google-cloud" }, { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, @@ -2379,23 +1908,6 @@ requires-dist = [ { name = "uvloop" }, ] -[[package]] -name = "genkit-tools-sample-flows" -version = "0.1.1" -source = { editable = "tools/sample-flows" } -dependencies = [ - { name = "genkit" }, - { name = "httpx" }, - { name = "structlog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "httpx" }, - { name = "structlog" }, -] - [[package]] name = "genkit-workspace" version = "0.1.0" @@ -2404,12 +1916,9 @@ dependencies = [ { name = "conform" }, { name = "dotpromptz" }, { name = "genkit" }, - { name = "genkit-plugin-amazon-bedrock" }, { name = "genkit-plugin-anthropic" }, { name = "genkit-plugin-compat-oai" }, - { name = "genkit-plugin-dev-local-vectorstore" }, { name = "genkit-plugin-fastapi" }, - { name = "genkit-plugin-firebase" }, { name = "genkit-plugin-flask" }, { name = "genkit-plugin-google-cloud" }, { name = "genkit-plugin-google-genai" }, @@ -2477,12 +1986,9 @@ requires-dist = [ { name = "conform", editable = "tools/conform" }, { name = "dotpromptz", specifier = "==0.1.5" }, { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-amazon-bedrock", editable = "plugins/amazon-bedrock" }, { name = "genkit-plugin-anthropic", editable = "plugins/anthropic" }, { name = "genkit-plugin-compat-oai", editable = "plugins/compat-oai" }, - { name = "genkit-plugin-dev-local-vectorstore", editable = "plugins/dev-local-vectorstore" }, { name = "genkit-plugin-fastapi", editable = "plugins/fastapi" }, - { name = "genkit-plugin-firebase", editable = "plugins/firebase" }, { name = "genkit-plugin-flask", editable = "plugins/flask" }, { name = "genkit-plugin-google-cloud", editable = "plugins/google-cloud" }, { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, @@ -2530,7 +2036,7 @@ lint = [ { name = "opentelemetry-instrumentation-fastapi", specifier = ">=0.41b0" }, { name = "opentelemetry-instrumentation-grpc", specifier = ">=0.41b0" }, { name = "pip-audit", specifier = ">=2.7.0" }, - { name = "pypdf", specifier = ">=6.6.2" }, + { name = "pypdf", specifier = ">=6.7.5" }, { name = "pyrefly", specifier = ">=0.15.0" }, { name = "pyright", specifier = ">=1.1.392" }, { name = "pysentry-rs", specifier = ">=0.3.14" }, @@ -3581,15 +3087,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/67/8a/a342b2f0251f3dac4ca17618265d93bf244a2a4d089126e81e4c1056ac50/jiter-0.13.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7bb00b6d26db67a05fe3e12c76edc75f32077fb51deed13822dc648fa373bc19", size = 343768, upload-time = "2026-02-02T12:37:55.055Z" }, ] -[[package]] -name = "jmespath" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, -] - [[package]] name = "json5" version = "0.13.0" @@ -5026,18 +4523,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/f9/8a4ce3901bc52277794e4b18c4ac43dc5929806eff01d22812364132f45f/opentelemetry_instrumentation_logging-0.60b1-py3-none-any.whl", hash = "sha256:f2e18cbc7e1dd3628c80e30d243897fdc93c5b7e0c8ae60abd2b9b6a99f82343", size = 12577, upload-time = "2025-12-11T13:36:08.123Z" }, ] -[[package]] -name = "opentelemetry-propagator-aws-xray" -version = "1.0.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-api" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f4/31/40004e9e55b1e5694ef3a7526f0b7637df44196fc68a8b7d248a3684680f/opentelemetry_propagator_aws_xray-1.0.2.tar.gz", hash = "sha256:6b2cee5479d2ef0172307b66ed2ed151f598a0fd29b3c01133ac87ca06326260", size = 10994, upload-time = "2024-08-05T17:45:57.601Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ea/89/849a0847871fd9745315896ad9e23d6479db84d90b8b36c4c26dc46e92b8/opentelemetry_propagator_aws_xray-1.0.2-py3-none-any.whl", hash = "sha256:1c99181ee228e99bddb638a0c911a297fa21f1c3a0af951f841e79919b5f1934", size = 10856, upload-time = "2024-08-05T17:45:56.492Z" }, -] - [[package]] name = "opentelemetry-proto" version = "1.39.1" @@ -5079,18 +4564,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7c/98/e91cf858f203d86f4eccdf763dcf01cf03f1dae80c3750f7e635bfa206b6/opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c", size = 132565, upload-time = "2025-12-11T13:32:35.069Z" }, ] -[[package]] -name = "opentelemetry-sdk-extension-aws" -version = "2.1.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-sdk" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f5/b3/825c93fe4c238845f1356297abea33d03b2adaafb5ae98fc257b394de124/opentelemetry_sdk_extension_aws-2.1.0.tar.gz", hash = "sha256:ff68ddecc1910f62c019d22ec0f7461713ead7f662d6a2304d4089c1a0b20416", size = 16334, upload-time = "2024-12-24T15:01:57.387Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/02/61/47a6a43b7935d54b5734fbf3fb0357dd5a7d0dfaa9677b7318518fe8d507/opentelemetry_sdk_extension_aws-2.1.0-py3-none-any.whl", hash = "sha256:c7cf6efc275d2c24108a468d954287ce5aab9733bac816a080cfb3117374e63a", size = 18776, upload-time = "2024-12-24T15:01:56.053Z" }, -] - [[package]] name = "opentelemetry-semantic-conventions" version = "0.60b1" @@ -5487,120 +4960,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, ] -[[package]] -name = "propcache" -version = "0.4.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/9e/da/e9fc233cf63743258bff22b3dfa7ea5baef7b5bc324af47a0ad89b8ffc6f/propcache-0.4.1.tar.gz", hash = "sha256:f48107a8c637e80362555f37ecf49abe20370e557cc4ab374f04ec4423c97c3d", size = 46442, upload-time = "2025-10-08T19:49:02.291Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/0e/934b541323035566a9af292dba85a195f7b78179114f2c6ebb24551118a9/propcache-0.4.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:7c2d1fa3201efaf55d730400d945b5b3ab6e672e100ba0f9a409d950ab25d7db", size = 79534, upload-time = "2025-10-08T19:46:02.083Z" }, - { url = "https://files.pythonhosted.org/packages/a1/6b/db0d03d96726d995dc7171286c6ba9d8d14251f37433890f88368951a44e/propcache-0.4.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:1eb2994229cc8ce7fe9b3db88f5465f5fd8651672840b2e426b88cdb1a30aac8", size = 45526, upload-time = "2025-10-08T19:46:03.884Z" }, - { url = "https://files.pythonhosted.org/packages/e4/c3/82728404aea669e1600f304f2609cde9e665c18df5a11cdd57ed73c1dceb/propcache-0.4.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:66c1f011f45a3b33d7bcb22daed4b29c0c9e2224758b6be00686731e1b46f925", size = 47263, upload-time = "2025-10-08T19:46:05.405Z" }, - { url = "https://files.pythonhosted.org/packages/df/1b/39313ddad2bf9187a1432654c38249bab4562ef535ef07f5eb6eb04d0b1b/propcache-0.4.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9a52009f2adffe195d0b605c25ec929d26b36ef986ba85244891dee3b294df21", size = 201012, upload-time = "2025-10-08T19:46:07.165Z" }, - { url = "https://files.pythonhosted.org/packages/5b/01/f1d0b57d136f294a142acf97f4ed58c8e5b974c21e543000968357115011/propcache-0.4.1-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5d4e2366a9c7b837555cf02fb9be2e3167d333aff716332ef1b7c3a142ec40c5", size = 209491, upload-time = "2025-10-08T19:46:08.909Z" }, - { url = "https://files.pythonhosted.org/packages/a1/c8/038d909c61c5bb039070b3fb02ad5cccdb1dde0d714792e251cdb17c9c05/propcache-0.4.1-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9d2b6caef873b4f09e26ea7e33d65f42b944837563a47a94719cc3544319a0db", size = 215319, upload-time = "2025-10-08T19:46:10.7Z" }, - { url = "https://files.pythonhosted.org/packages/08/57/8c87e93142b2c1fa2408e45695205a7ba05fb5db458c0bf5c06ba0e09ea6/propcache-0.4.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2b16ec437a8c8a965ecf95739448dd938b5c7f56e67ea009f4300d8df05f32b7", size = 196856, upload-time = "2025-10-08T19:46:12.003Z" }, - { url = "https://files.pythonhosted.org/packages/42/df/5615fec76aa561987a534759b3686008a288e73107faa49a8ae5795a9f7a/propcache-0.4.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:296f4c8ed03ca7476813fe666c9ea97869a8d7aec972618671b33a38a5182ef4", size = 193241, upload-time = "2025-10-08T19:46:13.495Z" }, - { url = "https://files.pythonhosted.org/packages/d5/21/62949eb3a7a54afe8327011c90aca7e03547787a88fb8bd9726806482fea/propcache-0.4.1-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:1f0978529a418ebd1f49dad413a2b68af33f85d5c5ca5c6ca2a3bed375a7ac60", size = 190552, upload-time = "2025-10-08T19:46:14.938Z" }, - { url = "https://files.pythonhosted.org/packages/30/ee/ab4d727dd70806e5b4de96a798ae7ac6e4d42516f030ee60522474b6b332/propcache-0.4.1-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:fd138803047fb4c062b1c1dd95462f5209456bfab55c734458f15d11da288f8f", size = 200113, upload-time = "2025-10-08T19:46:16.695Z" }, - { url = "https://files.pythonhosted.org/packages/8a/0b/38b46208e6711b016aa8966a3ac793eee0d05c7159d8342aa27fc0bc365e/propcache-0.4.1-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:8c9b3cbe4584636d72ff556d9036e0c9317fa27b3ac1f0f558e7e84d1c9c5900", size = 200778, upload-time = "2025-10-08T19:46:18.023Z" }, - { url = "https://files.pythonhosted.org/packages/cf/81/5abec54355ed344476bee711e9f04815d4b00a311ab0535599204eecc257/propcache-0.4.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f93243fdc5657247533273ac4f86ae106cc6445a0efacb9a1bfe982fcfefd90c", size = 193047, upload-time = "2025-10-08T19:46:19.449Z" }, - { url = "https://files.pythonhosted.org/packages/ec/b6/1f237c04e32063cb034acd5f6ef34ef3a394f75502e72703545631ab1ef6/propcache-0.4.1-cp310-cp310-win32.whl", hash = "sha256:a0ee98db9c5f80785b266eb805016e36058ac72c51a064040f2bc43b61101cdb", size = 38093, upload-time = "2025-10-08T19:46:20.643Z" }, - { url = "https://files.pythonhosted.org/packages/a6/67/354aac4e0603a15f76439caf0427781bcd6797f370377f75a642133bc954/propcache-0.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:1cdb7988c4e5ac7f6d175a28a9aa0c94cb6f2ebe52756a3c0cda98d2809a9e37", size = 41638, upload-time = "2025-10-08T19:46:21.935Z" }, - { url = "https://files.pythonhosted.org/packages/e0/e1/74e55b9fd1a4c209ff1a9a824bf6c8b3d1fc5a1ac3eabe23462637466785/propcache-0.4.1-cp310-cp310-win_arm64.whl", hash = "sha256:d82ad62b19645419fe79dd63b3f9253e15b30e955c0170e5cebc350c1844e581", size = 38229, upload-time = "2025-10-08T19:46:23.368Z" }, - { url = "https://files.pythonhosted.org/packages/8c/d4/4e2c9aaf7ac2242b9358f98dccd8f90f2605402f5afeff6c578682c2c491/propcache-0.4.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:60a8fda9644b7dfd5dece8c61d8a85e271cb958075bfc4e01083c148b61a7caf", size = 80208, upload-time = "2025-10-08T19:46:24.597Z" }, - { url = "https://files.pythonhosted.org/packages/c2/21/d7b68e911f9c8e18e4ae43bdbc1e1e9bbd971f8866eb81608947b6f585ff/propcache-0.4.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:c30b53e7e6bda1d547cabb47c825f3843a0a1a42b0496087bb58d8fedf9f41b5", size = 45777, upload-time = "2025-10-08T19:46:25.733Z" }, - { url = "https://files.pythonhosted.org/packages/d3/1d/11605e99ac8ea9435651ee71ab4cb4bf03f0949586246476a25aadfec54a/propcache-0.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6918ecbd897443087a3b7cd978d56546a812517dcaaca51b49526720571fa93e", size = 47647, upload-time = "2025-10-08T19:46:27.304Z" }, - { url = "https://files.pythonhosted.org/packages/58/1a/3c62c127a8466c9c843bccb503d40a273e5cc69838805f322e2826509e0d/propcache-0.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3d902a36df4e5989763425a8ab9e98cd8ad5c52c823b34ee7ef307fd50582566", size = 214929, upload-time = "2025-10-08T19:46:28.62Z" }, - { url = "https://files.pythonhosted.org/packages/56/b9/8fa98f850960b367c4b8fe0592e7fc341daa7a9462e925228f10a60cf74f/propcache-0.4.1-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a9695397f85973bb40427dedddf70d8dc4a44b22f1650dd4af9eedf443d45165", size = 221778, upload-time = "2025-10-08T19:46:30.358Z" }, - { url = "https://files.pythonhosted.org/packages/46/a6/0ab4f660eb59649d14b3d3d65c439421cf2f87fe5dd68591cbe3c1e78a89/propcache-0.4.1-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2bb07ffd7eaad486576430c89f9b215f9e4be68c4866a96e97db9e97fead85dc", size = 228144, upload-time = "2025-10-08T19:46:32.607Z" }, - { url = "https://files.pythonhosted.org/packages/52/6a/57f43e054fb3d3a56ac9fc532bc684fc6169a26c75c353e65425b3e56eef/propcache-0.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fd6f30fdcf9ae2a70abd34da54f18da086160e4d7d9251f81f3da0ff84fc5a48", size = 210030, upload-time = "2025-10-08T19:46:33.969Z" }, - { url = "https://files.pythonhosted.org/packages/40/e2/27e6feebb5f6b8408fa29f5efbb765cd54c153ac77314d27e457a3e993b7/propcache-0.4.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:fc38cba02d1acba4e2869eef1a57a43dfbd3d49a59bf90dda7444ec2be6a5570", size = 208252, upload-time = "2025-10-08T19:46:35.309Z" }, - { url = "https://files.pythonhosted.org/packages/9e/f8/91c27b22ccda1dbc7967f921c42825564fa5336a01ecd72eb78a9f4f53c2/propcache-0.4.1-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:67fad6162281e80e882fb3ec355398cf72864a54069d060321f6cd0ade95fe85", size = 202064, upload-time = "2025-10-08T19:46:36.993Z" }, - { url = "https://files.pythonhosted.org/packages/f2/26/7f00bd6bd1adba5aafe5f4a66390f243acab58eab24ff1a08bebb2ef9d40/propcache-0.4.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:f10207adf04d08bec185bae14d9606a1444715bc99180f9331c9c02093e1959e", size = 212429, upload-time = "2025-10-08T19:46:38.398Z" }, - { url = "https://files.pythonhosted.org/packages/84/89/fd108ba7815c1117ddca79c228f3f8a15fc82a73bca8b142eb5de13b2785/propcache-0.4.1-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:e9b0d8d0845bbc4cfcdcbcdbf5086886bc8157aa963c31c777ceff7846c77757", size = 216727, upload-time = "2025-10-08T19:46:39.732Z" }, - { url = "https://files.pythonhosted.org/packages/79/37/3ec3f7e3173e73f1d600495d8b545b53802cbf35506e5732dd8578db3724/propcache-0.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:981333cb2f4c1896a12f4ab92a9cc8f09ea664e9b7dbdc4eff74627af3a11c0f", size = 205097, upload-time = "2025-10-08T19:46:41.025Z" }, - { url = "https://files.pythonhosted.org/packages/61/b0/b2631c19793f869d35f47d5a3a56fb19e9160d3c119f15ac7344fc3ccae7/propcache-0.4.1-cp311-cp311-win32.whl", hash = "sha256:f1d2f90aeec838a52f1c1a32fe9a619fefd5e411721a9117fbf82aea638fe8a1", size = 38084, upload-time = "2025-10-08T19:46:42.693Z" }, - { url = "https://files.pythonhosted.org/packages/f4/78/6cce448e2098e9f3bfc91bb877f06aa24b6ccace872e39c53b2f707c4648/propcache-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:364426a62660f3f699949ac8c621aad6977be7126c5807ce48c0aeb8e7333ea6", size = 41637, upload-time = "2025-10-08T19:46:43.778Z" }, - { url = "https://files.pythonhosted.org/packages/9c/e9/754f180cccd7f51a39913782c74717c581b9cc8177ad0e949f4d51812383/propcache-0.4.1-cp311-cp311-win_arm64.whl", hash = "sha256:e53f3a38d3510c11953f3e6a33f205c6d1b001129f972805ca9b42fc308bc239", size = 38064, upload-time = "2025-10-08T19:46:44.872Z" }, - { url = "https://files.pythonhosted.org/packages/a2/0f/f17b1b2b221d5ca28b4b876e8bb046ac40466513960646bda8e1853cdfa2/propcache-0.4.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e153e9cd40cc8945138822807139367f256f89c6810c2634a4f6902b52d3b4e2", size = 80061, upload-time = "2025-10-08T19:46:46.075Z" }, - { url = "https://files.pythonhosted.org/packages/76/47/8ccf75935f51448ba9a16a71b783eb7ef6b9ee60f5d14c7f8a8a79fbeed7/propcache-0.4.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:cd547953428f7abb73c5ad82cbb32109566204260d98e41e5dfdc682eb7f8403", size = 46037, upload-time = "2025-10-08T19:46:47.23Z" }, - { url = "https://files.pythonhosted.org/packages/0a/b6/5c9a0e42df4d00bfb4a3cbbe5cf9f54260300c88a0e9af1f47ca5ce17ac0/propcache-0.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f048da1b4f243fc44f205dfd320933a951b8d89e0afd4c7cacc762a8b9165207", size = 47324, upload-time = "2025-10-08T19:46:48.384Z" }, - { url = "https://files.pythonhosted.org/packages/9e/d3/6c7ee328b39a81ee877c962469f1e795f9db87f925251efeb0545e0020d0/propcache-0.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ec17c65562a827bba85e3872ead335f95405ea1674860d96483a02f5c698fa72", size = 225505, upload-time = "2025-10-08T19:46:50.055Z" }, - { url = "https://files.pythonhosted.org/packages/01/5d/1c53f4563490b1d06a684742cc6076ef944bc6457df6051b7d1a877c057b/propcache-0.4.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:405aac25c6394ef275dee4c709be43745d36674b223ba4eb7144bf4d691b7367", size = 230242, upload-time = "2025-10-08T19:46:51.815Z" }, - { url = "https://files.pythonhosted.org/packages/20/e1/ce4620633b0e2422207c3cb774a0ee61cac13abc6217763a7b9e2e3f4a12/propcache-0.4.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0013cb6f8dde4b2a2f66903b8ba740bdfe378c943c4377a200551ceb27f379e4", size = 238474, upload-time = "2025-10-08T19:46:53.208Z" }, - { url = "https://files.pythonhosted.org/packages/46/4b/3aae6835b8e5f44ea6a68348ad90f78134047b503765087be2f9912140ea/propcache-0.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:15932ab57837c3368b024473a525e25d316d8353016e7cc0e5ba9eb343fbb1cf", size = 221575, upload-time = "2025-10-08T19:46:54.511Z" }, - { url = "https://files.pythonhosted.org/packages/6e/a5/8a5e8678bcc9d3a1a15b9a29165640d64762d424a16af543f00629c87338/propcache-0.4.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:031dce78b9dc099f4c29785d9cf5577a3faf9ebf74ecbd3c856a7b92768c3df3", size = 216736, upload-time = "2025-10-08T19:46:56.212Z" }, - { url = "https://files.pythonhosted.org/packages/f1/63/b7b215eddeac83ca1c6b934f89d09a625aa9ee4ba158338854c87210cc36/propcache-0.4.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:ab08df6c9a035bee56e31af99be621526bd237bea9f32def431c656b29e41778", size = 213019, upload-time = "2025-10-08T19:46:57.595Z" }, - { url = "https://files.pythonhosted.org/packages/57/74/f580099a58c8af587cac7ba19ee7cb418506342fbbe2d4a4401661cca886/propcache-0.4.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:4d7af63f9f93fe593afbf104c21b3b15868efb2c21d07d8732c0c4287e66b6a6", size = 220376, upload-time = "2025-10-08T19:46:59.067Z" }, - { url = "https://files.pythonhosted.org/packages/c4/ee/542f1313aff7eaf19c2bb758c5d0560d2683dac001a1c96d0774af799843/propcache-0.4.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:cfc27c945f422e8b5071b6e93169679e4eb5bf73bbcbf1ba3ae3a83d2f78ebd9", size = 226988, upload-time = "2025-10-08T19:47:00.544Z" }, - { url = "https://files.pythonhosted.org/packages/8f/18/9c6b015dd9c6930f6ce2229e1f02fb35298b847f2087ea2b436a5bfa7287/propcache-0.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:35c3277624a080cc6ec6f847cbbbb5b49affa3598c4535a0a4682a697aaa5c75", size = 215615, upload-time = "2025-10-08T19:47:01.968Z" }, - { url = "https://files.pythonhosted.org/packages/80/9e/e7b85720b98c45a45e1fca6a177024934dc9bc5f4d5dd04207f216fc33ed/propcache-0.4.1-cp312-cp312-win32.whl", hash = "sha256:671538c2262dadb5ba6395e26c1731e1d52534bfe9ae56d0b5573ce539266aa8", size = 38066, upload-time = "2025-10-08T19:47:03.503Z" }, - { url = "https://files.pythonhosted.org/packages/54/09/d19cff2a5aaac632ec8fc03737b223597b1e347416934c1b3a7df079784c/propcache-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:cb2d222e72399fcf5890d1d5cc1060857b9b236adff2792ff48ca2dfd46c81db", size = 41655, upload-time = "2025-10-08T19:47:04.973Z" }, - { url = "https://files.pythonhosted.org/packages/68/ab/6b5c191bb5de08036a8c697b265d4ca76148efb10fa162f14af14fb5f076/propcache-0.4.1-cp312-cp312-win_arm64.whl", hash = "sha256:204483131fb222bdaaeeea9f9e6c6ed0cac32731f75dfc1d4a567fc1926477c1", size = 37789, upload-time = "2025-10-08T19:47:06.077Z" }, - { url = "https://files.pythonhosted.org/packages/bf/df/6d9c1b6ac12b003837dde8a10231a7344512186e87b36e855bef32241942/propcache-0.4.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:43eedf29202c08550aac1d14e0ee619b0430aaef78f85864c1a892294fbc28cf", size = 77750, upload-time = "2025-10-08T19:47:07.648Z" }, - { url = "https://files.pythonhosted.org/packages/8b/e8/677a0025e8a2acf07d3418a2e7ba529c9c33caf09d3c1f25513023c1db56/propcache-0.4.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d62cdfcfd89ccb8de04e0eda998535c406bf5e060ffd56be6c586cbcc05b3311", size = 44780, upload-time = "2025-10-08T19:47:08.851Z" }, - { url = "https://files.pythonhosted.org/packages/89/a4/92380f7ca60f99ebae761936bc48a72a639e8a47b29050615eef757cb2a7/propcache-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:cae65ad55793da34db5f54e4029b89d3b9b9490d8abe1b4c7ab5d4b8ec7ebf74", size = 46308, upload-time = "2025-10-08T19:47:09.982Z" }, - { url = "https://files.pythonhosted.org/packages/2d/48/c5ac64dee5262044348d1d78a5f85dd1a57464a60d30daee946699963eb3/propcache-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:333ddb9031d2704a301ee3e506dc46b1fe5f294ec198ed6435ad5b6a085facfe", size = 208182, upload-time = "2025-10-08T19:47:11.319Z" }, - { url = "https://files.pythonhosted.org/packages/c6/0c/cd762dd011a9287389a6a3eb43aa30207bde253610cca06824aeabfe9653/propcache-0.4.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:fd0858c20f078a32cf55f7e81473d96dcf3b93fd2ccdb3d40fdf54b8573df3af", size = 211215, upload-time = "2025-10-08T19:47:13.146Z" }, - { url = "https://files.pythonhosted.org/packages/30/3e/49861e90233ba36890ae0ca4c660e95df565b2cd15d4a68556ab5865974e/propcache-0.4.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:678ae89ebc632c5c204c794f8dab2837c5f159aeb59e6ed0539500400577298c", size = 218112, upload-time = "2025-10-08T19:47:14.913Z" }, - { url = "https://files.pythonhosted.org/packages/f1/8b/544bc867e24e1bd48f3118cecd3b05c694e160a168478fa28770f22fd094/propcache-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d472aeb4fbf9865e0c6d622d7f4d54a4e101a89715d8904282bb5f9a2f476c3f", size = 204442, upload-time = "2025-10-08T19:47:16.277Z" }, - { url = "https://files.pythonhosted.org/packages/50/a6/4282772fd016a76d3e5c0df58380a5ea64900afd836cec2c2f662d1b9bb3/propcache-0.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4d3df5fa7e36b3225954fba85589da77a0fe6a53e3976de39caf04a0db4c36f1", size = 199398, upload-time = "2025-10-08T19:47:17.962Z" }, - { url = "https://files.pythonhosted.org/packages/3e/ec/d8a7cd406ee1ddb705db2139f8a10a8a427100347bd698e7014351c7af09/propcache-0.4.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:ee17f18d2498f2673e432faaa71698032b0127ebf23ae5974eeaf806c279df24", size = 196920, upload-time = "2025-10-08T19:47:19.355Z" }, - { url = "https://files.pythonhosted.org/packages/f6/6c/f38ab64af3764f431e359f8baf9e0a21013e24329e8b85d2da32e8ed07ca/propcache-0.4.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:580e97762b950f993ae618e167e7be9256b8353c2dcd8b99ec100eb50f5286aa", size = 203748, upload-time = "2025-10-08T19:47:21.338Z" }, - { url = "https://files.pythonhosted.org/packages/d6/e3/fa846bd70f6534d647886621388f0a265254d30e3ce47e5c8e6e27dbf153/propcache-0.4.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:501d20b891688eb8e7aa903021f0b72d5a55db40ffaab27edefd1027caaafa61", size = 205877, upload-time = "2025-10-08T19:47:23.059Z" }, - { url = "https://files.pythonhosted.org/packages/e2/39/8163fc6f3133fea7b5f2827e8eba2029a0277ab2c5beee6c1db7b10fc23d/propcache-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9a0bd56e5b100aef69bd8562b74b46254e7c8812918d3baa700c8a8009b0af66", size = 199437, upload-time = "2025-10-08T19:47:24.445Z" }, - { url = "https://files.pythonhosted.org/packages/93/89/caa9089970ca49c7c01662bd0eeedfe85494e863e8043565aeb6472ce8fe/propcache-0.4.1-cp313-cp313-win32.whl", hash = "sha256:bcc9aaa5d80322bc2fb24bb7accb4a30f81e90ab8d6ba187aec0744bc302ad81", size = 37586, upload-time = "2025-10-08T19:47:25.736Z" }, - { url = "https://files.pythonhosted.org/packages/f5/ab/f76ec3c3627c883215b5c8080debb4394ef5a7a29be811f786415fc1e6fd/propcache-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:381914df18634f5494334d201e98245c0596067504b9372d8cf93f4bb23e025e", size = 40790, upload-time = "2025-10-08T19:47:26.847Z" }, - { url = "https://files.pythonhosted.org/packages/59/1b/e71ae98235f8e2ba5004d8cb19765a74877abf189bc53fc0c80d799e56c3/propcache-0.4.1-cp313-cp313-win_arm64.whl", hash = "sha256:8873eb4460fd55333ea49b7d189749ecf6e55bf85080f11b1c4530ed3034cba1", size = 37158, upload-time = "2025-10-08T19:47:27.961Z" }, - { url = "https://files.pythonhosted.org/packages/83/ce/a31bbdfc24ee0dcbba458c8175ed26089cf109a55bbe7b7640ed2470cfe9/propcache-0.4.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:92d1935ee1f8d7442da9c0c4fa7ac20d07e94064184811b685f5c4fada64553b", size = 81451, upload-time = "2025-10-08T19:47:29.445Z" }, - { url = "https://files.pythonhosted.org/packages/25/9c/442a45a470a68456e710d96cacd3573ef26a1d0a60067e6a7d5e655621ed/propcache-0.4.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:473c61b39e1460d386479b9b2f337da492042447c9b685f28be4f74d3529e566", size = 46374, upload-time = "2025-10-08T19:47:30.579Z" }, - { url = "https://files.pythonhosted.org/packages/f4/bf/b1d5e21dbc3b2e889ea4327044fb16312a736d97640fb8b6aa3f9c7b3b65/propcache-0.4.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:c0ef0aaafc66fbd87842a3fe3902fd889825646bc21149eafe47be6072725835", size = 48396, upload-time = "2025-10-08T19:47:31.79Z" }, - { url = "https://files.pythonhosted.org/packages/f4/04/5b4c54a103d480e978d3c8a76073502b18db0c4bc17ab91b3cb5092ad949/propcache-0.4.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f95393b4d66bfae908c3ca8d169d5f79cd65636ae15b5e7a4f6e67af675adb0e", size = 275950, upload-time = "2025-10-08T19:47:33.481Z" }, - { url = "https://files.pythonhosted.org/packages/b4/c1/86f846827fb969c4b78b0af79bba1d1ea2156492e1b83dea8b8a6ae27395/propcache-0.4.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c07fda85708bc48578467e85099645167a955ba093be0a2dcba962195676e859", size = 273856, upload-time = "2025-10-08T19:47:34.906Z" }, - { url = "https://files.pythonhosted.org/packages/36/1d/fc272a63c8d3bbad6878c336c7a7dea15e8f2d23a544bda43205dfa83ada/propcache-0.4.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:af223b406d6d000830c6f65f1e6431783fc3f713ba3e6cc8c024d5ee96170a4b", size = 280420, upload-time = "2025-10-08T19:47:36.338Z" }, - { url = "https://files.pythonhosted.org/packages/07/0c/01f2219d39f7e53d52e5173bcb09c976609ba30209912a0680adfb8c593a/propcache-0.4.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a78372c932c90ee474559c5ddfffd718238e8673c340dc21fe45c5b8b54559a0", size = 263254, upload-time = "2025-10-08T19:47:37.692Z" }, - { url = "https://files.pythonhosted.org/packages/2d/18/cd28081658ce597898f0c4d174d4d0f3c5b6d4dc27ffafeef835c95eb359/propcache-0.4.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:564d9f0d4d9509e1a870c920a89b2fec951b44bf5ba7d537a9e7c1ccec2c18af", size = 261205, upload-time = "2025-10-08T19:47:39.659Z" }, - { url = "https://files.pythonhosted.org/packages/7a/71/1f9e22eb8b8316701c2a19fa1f388c8a3185082607da8e406a803c9b954e/propcache-0.4.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:17612831fda0138059cc5546f4d12a2aacfb9e47068c06af35c400ba58ba7393", size = 247873, upload-time = "2025-10-08T19:47:41.084Z" }, - { url = "https://files.pythonhosted.org/packages/4a/65/3d4b61f36af2b4eddba9def857959f1016a51066b4f1ce348e0cf7881f58/propcache-0.4.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:41a89040cb10bd345b3c1a873b2bf36413d48da1def52f268a055f7398514874", size = 262739, upload-time = "2025-10-08T19:47:42.51Z" }, - { url = "https://files.pythonhosted.org/packages/2a/42/26746ab087faa77c1c68079b228810436ccd9a5ce9ac85e2b7307195fd06/propcache-0.4.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:e35b88984e7fa64aacecea39236cee32dd9bd8c55f57ba8a75cf2399553f9bd7", size = 263514, upload-time = "2025-10-08T19:47:43.927Z" }, - { url = "https://files.pythonhosted.org/packages/94/13/630690fe201f5502d2403dd3cfd451ed8858fe3c738ee88d095ad2ff407b/propcache-0.4.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6f8b465489f927b0df505cbe26ffbeed4d6d8a2bbc61ce90eb074ff129ef0ab1", size = 257781, upload-time = "2025-10-08T19:47:45.448Z" }, - { url = "https://files.pythonhosted.org/packages/92/f7/1d4ec5841505f423469efbfc381d64b7b467438cd5a4bbcbb063f3b73d27/propcache-0.4.1-cp313-cp313t-win32.whl", hash = "sha256:2ad890caa1d928c7c2965b48f3a3815c853180831d0e5503d35cf00c472f4717", size = 41396, upload-time = "2025-10-08T19:47:47.202Z" }, - { url = "https://files.pythonhosted.org/packages/48/f0/615c30622316496d2cbbc29f5985f7777d3ada70f23370608c1d3e081c1f/propcache-0.4.1-cp313-cp313t-win_amd64.whl", hash = "sha256:f7ee0e597f495cf415bcbd3da3caa3bd7e816b74d0d52b8145954c5e6fd3ff37", size = 44897, upload-time = "2025-10-08T19:47:48.336Z" }, - { url = "https://files.pythonhosted.org/packages/fd/ca/6002e46eccbe0e33dcd4069ef32f7f1c9e243736e07adca37ae8c4830ec3/propcache-0.4.1-cp313-cp313t-win_arm64.whl", hash = "sha256:929d7cbe1f01bb7baffb33dc14eb5691c95831450a26354cd210a8155170c93a", size = 39789, upload-time = "2025-10-08T19:47:49.876Z" }, - { url = "https://files.pythonhosted.org/packages/8e/5c/bca52d654a896f831b8256683457ceddd490ec18d9ec50e97dfd8fc726a8/propcache-0.4.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3f7124c9d820ba5548d431afb4632301acf965db49e666aa21c305cbe8c6de12", size = 78152, upload-time = "2025-10-08T19:47:51.051Z" }, - { url = "https://files.pythonhosted.org/packages/65/9b/03b04e7d82a5f54fb16113d839f5ea1ede58a61e90edf515f6577c66fa8f/propcache-0.4.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c0d4b719b7da33599dfe3b22d3db1ef789210a0597bc650b7cee9c77c2be8c5c", size = 44869, upload-time = "2025-10-08T19:47:52.594Z" }, - { url = "https://files.pythonhosted.org/packages/b2/fa/89a8ef0468d5833a23fff277b143d0573897cf75bd56670a6d28126c7d68/propcache-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:9f302f4783709a78240ebc311b793f123328716a60911d667e0c036bc5dcbded", size = 46596, upload-time = "2025-10-08T19:47:54.073Z" }, - { url = "https://files.pythonhosted.org/packages/86/bd/47816020d337f4a746edc42fe8d53669965138f39ee117414c7d7a340cfe/propcache-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c80ee5802e3fb9ea37938e7eecc307fb984837091d5fd262bb37238b1ae97641", size = 206981, upload-time = "2025-10-08T19:47:55.715Z" }, - { url = "https://files.pythonhosted.org/packages/df/f6/c5fa1357cc9748510ee55f37173eb31bfde6d94e98ccd9e6f033f2fc06e1/propcache-0.4.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ed5a841e8bb29a55fb8159ed526b26adc5bdd7e8bd7bf793ce647cb08656cdf4", size = 211490, upload-time = "2025-10-08T19:47:57.499Z" }, - { url = "https://files.pythonhosted.org/packages/80/1e/e5889652a7c4a3846683401a48f0f2e5083ce0ec1a8a5221d8058fbd1adf/propcache-0.4.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:55c72fd6ea2da4c318e74ffdf93c4fe4e926051133657459131a95c846d16d44", size = 215371, upload-time = "2025-10-08T19:47:59.317Z" }, - { url = "https://files.pythonhosted.org/packages/b2/f2/889ad4b2408f72fe1a4f6a19491177b30ea7bf1a0fd5f17050ca08cfc882/propcache-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8326e144341460402713f91df60ade3c999d601e7eb5ff8f6f7862d54de0610d", size = 201424, upload-time = "2025-10-08T19:48:00.67Z" }, - { url = "https://files.pythonhosted.org/packages/27/73/033d63069b57b0812c8bd19f311faebeceb6ba31b8f32b73432d12a0b826/propcache-0.4.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:060b16ae65bc098da7f6d25bf359f1f31f688384858204fe5d652979e0015e5b", size = 197566, upload-time = "2025-10-08T19:48:02.604Z" }, - { url = "https://files.pythonhosted.org/packages/dc/89/ce24f3dc182630b4e07aa6d15f0ff4b14ed4b9955fae95a0b54c58d66c05/propcache-0.4.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:89eb3fa9524f7bec9de6e83cf3faed9d79bffa560672c118a96a171a6f55831e", size = 193130, upload-time = "2025-10-08T19:48:04.499Z" }, - { url = "https://files.pythonhosted.org/packages/a9/24/ef0d5fd1a811fb5c609278d0209c9f10c35f20581fcc16f818da959fc5b4/propcache-0.4.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:dee69d7015dc235f526fe80a9c90d65eb0039103fe565776250881731f06349f", size = 202625, upload-time = "2025-10-08T19:48:06.213Z" }, - { url = "https://files.pythonhosted.org/packages/f5/02/98ec20ff5546f68d673df2f7a69e8c0d076b5abd05ca882dc7ee3a83653d/propcache-0.4.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:5558992a00dfd54ccbc64a32726a3357ec93825a418a401f5cc67df0ac5d9e49", size = 204209, upload-time = "2025-10-08T19:48:08.432Z" }, - { url = "https://files.pythonhosted.org/packages/a0/87/492694f76759b15f0467a2a93ab68d32859672b646aa8a04ce4864e7932d/propcache-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c9b822a577f560fbd9554812526831712c1436d2c046cedee4c3796d3543b144", size = 197797, upload-time = "2025-10-08T19:48:09.968Z" }, - { url = "https://files.pythonhosted.org/packages/ee/36/66367de3575db1d2d3f3d177432bd14ee577a39d3f5d1b3d5df8afe3b6e2/propcache-0.4.1-cp314-cp314-win32.whl", hash = "sha256:ab4c29b49d560fe48b696cdcb127dd36e0bc2472548f3bf56cc5cb3da2b2984f", size = 38140, upload-time = "2025-10-08T19:48:11.232Z" }, - { url = "https://files.pythonhosted.org/packages/0c/2a/a758b47de253636e1b8aef181c0b4f4f204bf0dd964914fb2af90a95b49b/propcache-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:5a103c3eb905fcea0ab98be99c3a9a5ab2de60228aa5aceedc614c0281cf6153", size = 41257, upload-time = "2025-10-08T19:48:12.707Z" }, - { url = "https://files.pythonhosted.org/packages/34/5e/63bd5896c3fec12edcbd6f12508d4890d23c265df28c74b175e1ef9f4f3b/propcache-0.4.1-cp314-cp314-win_arm64.whl", hash = "sha256:74c1fb26515153e482e00177a1ad654721bf9207da8a494a0c05e797ad27b992", size = 38097, upload-time = "2025-10-08T19:48:13.923Z" }, - { url = "https://files.pythonhosted.org/packages/99/85/9ff785d787ccf9bbb3f3106f79884a130951436f58392000231b4c737c80/propcache-0.4.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:824e908bce90fb2743bd6b59db36eb4f45cd350a39637c9f73b1c1ea66f5b75f", size = 81455, upload-time = "2025-10-08T19:48:15.16Z" }, - { url = "https://files.pythonhosted.org/packages/90/85/2431c10c8e7ddb1445c1f7c4b54d886e8ad20e3c6307e7218f05922cad67/propcache-0.4.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:c2b5e7db5328427c57c8e8831abda175421b709672f6cfc3d630c3b7e2146393", size = 46372, upload-time = "2025-10-08T19:48:16.424Z" }, - { url = "https://files.pythonhosted.org/packages/01/20/b0972d902472da9bcb683fa595099911f4d2e86e5683bcc45de60dd05dc3/propcache-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6f6ff873ed40292cd4969ef5310179afd5db59fdf055897e282485043fc80ad0", size = 48411, upload-time = "2025-10-08T19:48:17.577Z" }, - { url = "https://files.pythonhosted.org/packages/e2/e3/7dc89f4f21e8f99bad3d5ddb3a3389afcf9da4ac69e3deb2dcdc96e74169/propcache-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:49a2dc67c154db2c1463013594c458881a069fcf98940e61a0569016a583020a", size = 275712, upload-time = "2025-10-08T19:48:18.901Z" }, - { url = "https://files.pythonhosted.org/packages/20/67/89800c8352489b21a8047c773067644e3897f02ecbbd610f4d46b7f08612/propcache-0.4.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:005f08e6a0529984491e37d8dbc3dd86f84bd78a8ceb5fa9a021f4c48d4984be", size = 273557, upload-time = "2025-10-08T19:48:20.762Z" }, - { url = "https://files.pythonhosted.org/packages/e2/a1/b52b055c766a54ce6d9c16d9aca0cad8059acd9637cdf8aa0222f4a026ef/propcache-0.4.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5c3310452e0d31390da9035c348633b43d7e7feb2e37be252be6da45abd1abcc", size = 280015, upload-time = "2025-10-08T19:48:22.592Z" }, - { url = "https://files.pythonhosted.org/packages/48/c8/33cee30bd890672c63743049f3c9e4be087e6780906bfc3ec58528be59c1/propcache-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c3c70630930447f9ef1caac7728c8ad1c56bc5015338b20fed0d08ea2480b3a", size = 262880, upload-time = "2025-10-08T19:48:23.947Z" }, - { url = "https://files.pythonhosted.org/packages/0c/b1/8f08a143b204b418285c88b83d00edbd61afbc2c6415ffafc8905da7038b/propcache-0.4.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8e57061305815dfc910a3634dcf584f08168a8836e6999983569f51a8544cd89", size = 260938, upload-time = "2025-10-08T19:48:25.656Z" }, - { url = "https://files.pythonhosted.org/packages/cf/12/96e4664c82ca2f31e1c8dff86afb867348979eb78d3cb8546a680287a1e9/propcache-0.4.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:521a463429ef54143092c11a77e04056dd00636f72e8c45b70aaa3140d639726", size = 247641, upload-time = "2025-10-08T19:48:27.207Z" }, - { url = "https://files.pythonhosted.org/packages/18/ed/e7a9cfca28133386ba52278136d42209d3125db08d0a6395f0cba0c0285c/propcache-0.4.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:120c964da3fdc75e3731aa392527136d4ad35868cc556fd09bb6d09172d9a367", size = 262510, upload-time = "2025-10-08T19:48:28.65Z" }, - { url = "https://files.pythonhosted.org/packages/f5/76/16d8bf65e8845dd62b4e2b57444ab81f07f40caa5652b8969b87ddcf2ef6/propcache-0.4.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:d8f353eb14ee3441ee844ade4277d560cdd68288838673273b978e3d6d2c8f36", size = 263161, upload-time = "2025-10-08T19:48:30.133Z" }, - { url = "https://files.pythonhosted.org/packages/e7/70/c99e9edb5d91d5ad8a49fa3c1e8285ba64f1476782fed10ab251ff413ba1/propcache-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ab2943be7c652f09638800905ee1bab2c544e537edb57d527997a24c13dc1455", size = 257393, upload-time = "2025-10-08T19:48:31.567Z" }, - { url = "https://files.pythonhosted.org/packages/08/02/87b25304249a35c0915d236575bc3574a323f60b47939a2262b77632a3ee/propcache-0.4.1-cp314-cp314t-win32.whl", hash = "sha256:05674a162469f31358c30bcaa8883cb7829fa3110bf9c0991fe27d7896c42d85", size = 42546, upload-time = "2025-10-08T19:48:32.872Z" }, - { url = "https://files.pythonhosted.org/packages/cb/ef/3c6ecf8b317aa982f309835e8f96987466123c6e596646d4e6a1dfcd080f/propcache-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:990f6b3e2a27d683cb7602ed6c86f15ee6b43b1194736f9baaeb93d0016633b1", size = 46259, upload-time = "2025-10-08T19:48:34.226Z" }, - { url = "https://files.pythonhosted.org/packages/c4/2d/346e946d4951f37eca1e4f55be0f0174c52cd70720f84029b02f296f4a38/propcache-0.4.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ecef2343af4cc68e05131e45024ba34f6095821988a9d0a02aa7c73fcc448aa9", size = 40428, upload-time = "2025-10-08T19:48:35.441Z" }, - { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, -] - [[package]] name = "proto-plus" version = "1.27.1" @@ -5628,36 +4987,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/57/bf/2086963c69bdac3d7cff1cc7ff79b8ce5ea0bec6797a017e1be338a46248/protobuf-6.33.5-py3-none-any.whl", hash = "sha256:69915a973dd0f60f31a08b8318b73eab2bd6a392c79184b3612226b0a3f8ec02", size = 170687, upload-time = "2026-01-29T21:51:32.557Z" }, ] -[[package]] -name = "provider-amazon-bedrock-hello" -version = "0.2.0" -source = { editable = "samples/provider-amazon-bedrock-hello" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-amazon-bedrock" }, - { name = "pydantic" }, - { name = "rich" }, - { name = "structlog" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-amazon-bedrock", editable = "plugins/amazon-bedrock" }, - { name = "pydantic", specifier = ">=2.0.0" }, - { name = "rich", specifier = ">=13.0.0" }, - { name = "structlog", specifier = ">=24.0.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - [[package]] name = "provider-anthropic-hello" version = "0.2.0" @@ -5720,36 +5049,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[[package]] -name = "provider-firestore-retriever" -version = "0.2.0" -source = { editable = "samples/provider-firestore-retriever" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-firebase" }, - { name = "genkit-plugin-google-genai" }, - { name = "google-cloud-firestore" }, - { name = "rich" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-firebase", editable = "plugins/firebase" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "google-cloud-firestore" }, - { name = "rich", specifier = ">=13.0.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - [[package]] name = "provider-google-genai-code-execution" version = "0.2.0" @@ -5988,112 +5287,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[[package]] -name = "provider-vertex-ai-rerank-eval" -version = "0.2.0" -source = { editable = "samples/provider-vertex-ai-rerank-eval" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-google-genai" }, - { name = "pydantic" }, - { name = "rich" }, - { name = "structlog" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "pydantic", specifier = ">=2.10.5" }, - { name = "rich", specifier = ">=13.0.0" }, - { name = "structlog", specifier = ">=25.2.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - -[[package]] -name = "provider-vertex-ai-vector-search-bigquery" -version = "0.2.0" -source = { editable = "samples/provider-vertex-ai-vector-search-bigquery" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-google-genai" }, - { name = "genkit-plugin-vertex-ai" }, - { name = "google-cloud-aiplatform" }, - { name = "google-cloud-bigquery" }, - { name = "pydantic" }, - { name = "rich" }, - { name = "strenum", marker = "python_full_version < '3.11'" }, - { name = "structlog" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "genkit-plugin-vertex-ai", editable = "plugins/vertex-ai" }, - { name = "google-cloud-aiplatform" }, - { name = "google-cloud-bigquery" }, - { name = "pydantic", specifier = ">=2.10.5" }, - { name = "rich", specifier = ">=13.0.0" }, - { name = "strenum", marker = "python_full_version < '3.11'", specifier = ">=0.4.15" }, - { name = "structlog", specifier = ">=25.2.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - -[[package]] -name = "provider-vertex-ai-vector-search-firestore" -version = "0.2.0" -source = { editable = "samples/provider-vertex-ai-vector-search-firestore" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-google-genai" }, - { name = "genkit-plugin-vertex-ai" }, - { name = "google-cloud-aiplatform" }, - { name = "google-cloud-firestore" }, - { name = "pydantic" }, - { name = "rich" }, - { name = "strenum", marker = "python_full_version < '3.11'" }, - { name = "structlog" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "genkit-plugin-vertex-ai", editable = "plugins/vertex-ai" }, - { name = "google-cloud-aiplatform" }, - { name = "google-cloud-firestore" }, - { name = "pydantic", specifier = ">=2.10.5" }, - { name = "rich", specifier = ">=13.0.0" }, - { name = "strenum", marker = "python_full_version < '3.11'", specifier = ">=0.4.15" }, - { name = "structlog", specifier = ">=25.2.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - [[package]] name = "psutil" version = "7.2.2" @@ -6439,14 +5632,14 @@ wheels = [ [[package]] name = "pypdf" -version = "6.7.0" +version = "6.8.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/10/45/8340de1c752bfda2da912ea0fa8c9a432f7de3f6315e82f1c0847811dff6/pypdf-6.7.0.tar.gz", hash = "sha256:eb95e244d9f434e6cfd157272283339ef586e593be64ee699c620f756d5c3f7e", size = 5299947, upload-time = "2026-02-08T14:47:11.897Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b4/a3/e705b0805212b663a4c27b861c8a603dba0f8b4bb281f96f8e746576a50d/pypdf-6.8.0.tar.gz", hash = "sha256:cb7eaeaa4133ce76f762184069a854e03f4d9a08568f0e0623f7ea810407833b", size = 5307831, upload-time = "2026-03-09T13:37:40.591Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ed/f1/c92e75a0eb18bb10845e792054ded113010de958b6d4998e201c029417bb/pypdf-6.7.0-py3-none-any.whl", hash = "sha256:62e85036d50839cbdf45b8067c2c1a1b925517514d7cba4cbe8755a6c2829bc9", size = 330557, upload-time = "2026-02-08T14:47:10.111Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ec/4ccf3bb86b1afe5d7176e1c8abcdbf22b53dd682ec2eda50e1caadcf6846/pypdf-6.8.0-py3-none-any.whl", hash = "sha256:2a025080a8dd73f48123c89c57174a5ff3806c71763ee4e49572dc90454943c7", size = 332177, upload-time = "2026-03-09T13:37:38.774Z" }, ] [[package]] @@ -7219,18 +6412,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/07/5bda6a85b220c64c65686bc85bd0bbb23b29c62b3a9f9433fa55f17cda93/ruff-0.15.1-py3-none-win_arm64.whl", hash = "sha256:5ff7d5f0f88567850f45081fac8f4ec212be8d0b963e385c3f7d0d2eb4899416", size = 10874604, upload-time = "2026-02-12T23:09:05.515Z" }, ] -[[package]] -name = "s3transfer" -version = "0.14.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "botocore" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/62/74/8d69dcb7a9efe8baa2046891735e5dfe433ad558ae23d9e3c14c633d1d58/s3transfer-0.14.0.tar.gz", hash = "sha256:eff12264e7c8b4985074ccce27a3b38a485bb7f7422cc8046fee9be4983e4125", size = 151547, upload-time = "2025-09-09T19:23:31.089Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/48/f0/ae7ca09223a81a1d890b2557186ea015f6e0502e9b8cb8e1813f1d8cfa4e/s3transfer-0.14.0-py3-none-any.whl", hash = "sha256:ea3b790c7077558ed1f02a3072fb3cb992bbbd253392f4b6e9e8976941c7d456", size = 85712, upload-time = "2025-09-09T19:23:30.041Z" }, -] - [[package]] name = "secretstorage" version = "3.5.0" @@ -8054,25 +7235,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[[package]] -name = "web-fastapi-minimal-devui" -version = "0.1.0" -source = { editable = "samples/web-fastapi-minimal-devui" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-compat-oai" }, - { name = "genkit-plugin-fastapi" }, - { name = "uvicorn", extra = ["standard"] }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-compat-oai", editable = "plugins/compat-oai" }, - { name = "genkit-plugin-fastapi", editable = "plugins/fastapi" }, - { name = "uvicorn", extras = ["standard"], specifier = ">=0.34.0" }, -] - [[package]] name = "web-flask-hello" version = "0.2.0" @@ -8191,14 +7353,14 @@ wheels = [ [[package]] name = "werkzeug" -version = "3.1.5" +version = "3.1.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markupsafe" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5a/70/1469ef1d3542ae7c2c7b72bd5e3a4e6ee69d7978fa8a3af05a38eca5becf/werkzeug-3.1.5.tar.gz", hash = "sha256:6a548b0e88955dd07ccb25539d7d0cc97417ee9e179677d22c7041c8f078ce67", size = 864754, upload-time = "2026-01-08T17:49:23.247Z" } +sdist = { url = "https://files.pythonhosted.org/packages/61/f1/ee81806690a87dab5f5653c1f146c92bc066d7f4cebc603ef88eb9e13957/werkzeug-3.1.6.tar.gz", hash = "sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25", size = 864736, upload-time = "2026-02-19T15:17:18.884Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ad/e4/8d97cca767bcc1be76d16fb76951608305561c6e056811587f36cb1316a8/werkzeug-3.1.5-py3-none-any.whl", hash = "sha256:5111e36e91086ece91f93268bb39b4a35c1e6f1feac762c9c822ded0a4e322dc", size = 225025, upload-time = "2026-01-08T17:49:21.859Z" }, + { url = "https://files.pythonhosted.org/packages/4d/ec/d58832f89ede95652fd01f4f24236af7d32b70cab2196dfcc2d2fd13c5c2/werkzeug-3.1.6-py3-none-any.whl", hash = "sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131", size = 225166, upload-time = "2026-02-19T15:17:17.475Z" }, ] [[package]] @@ -8291,132 +7453,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" }, ] -[[package]] -name = "yarl" -version = "1.22.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "idna" }, - { name = "multidict" }, - { name = "propcache" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/57/63/0c6ebca57330cd313f6102b16dd57ffaf3ec4c83403dcb45dbd15c6f3ea1/yarl-1.22.0.tar.gz", hash = "sha256:bebf8557577d4401ba8bd9ff33906f1376c877aa78d1fe216ad01b4d6745af71", size = 187169, upload-time = "2025-10-06T14:12:55.963Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/43/a2204825342f37c337f5edb6637040fa14e365b2fcc2346960201d457579/yarl-1.22.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:c7bd6683587567e5a49ee6e336e0612bec8329be1b7d4c8af5687dcdeb67ee1e", size = 140517, upload-time = "2025-10-06T14:08:42.494Z" }, - { url = "https://files.pythonhosted.org/packages/44/6f/674f3e6f02266428c56f704cd2501c22f78e8b2eeb23f153117cc86fb28a/yarl-1.22.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5cdac20da754f3a723cceea5b3448e1a2074866406adeb4ef35b469d089adb8f", size = 93495, upload-time = "2025-10-06T14:08:46.2Z" }, - { url = "https://files.pythonhosted.org/packages/b8/12/5b274d8a0f30c07b91b2f02cba69152600b47830fcfb465c108880fcee9c/yarl-1.22.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:07a524d84df0c10f41e3ee918846e1974aba4ec017f990dc735aad487a0bdfdf", size = 94400, upload-time = "2025-10-06T14:08:47.855Z" }, - { url = "https://files.pythonhosted.org/packages/e2/7f/df1b6949b1fa1aa9ff6de6e2631876ad4b73c4437822026e85d8acb56bb1/yarl-1.22.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e1b329cb8146d7b736677a2440e422eadd775d1806a81db2d4cded80a48efc1a", size = 347545, upload-time = "2025-10-06T14:08:49.683Z" }, - { url = "https://files.pythonhosted.org/packages/84/09/f92ed93bd6cd77872ab6c3462df45ca45cd058d8f1d0c9b4f54c1704429f/yarl-1.22.0-cp310-cp310-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:75976c6945d85dbb9ee6308cd7ff7b1fb9409380c82d6119bd778d8fcfe2931c", size = 319598, upload-time = "2025-10-06T14:08:51.215Z" }, - { url = "https://files.pythonhosted.org/packages/c3/97/ac3f3feae7d522cf7ccec3d340bb0b2b61c56cb9767923df62a135092c6b/yarl-1.22.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:80ddf7a5f8c86cb3eb4bc9028b07bbbf1f08a96c5c0bc1244be5e8fefcb94147", size = 363893, upload-time = "2025-10-06T14:08:53.144Z" }, - { url = "https://files.pythonhosted.org/packages/06/49/f3219097403b9c84a4d079b1d7bda62dd9b86d0d6e4428c02d46ab2c77fc/yarl-1.22.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d332fc2e3c94dad927f2112395772a4e4fedbcf8f80efc21ed7cdfae4d574fdb", size = 371240, upload-time = "2025-10-06T14:08:55.036Z" }, - { url = "https://files.pythonhosted.org/packages/35/9f/06b765d45c0e44e8ecf0fe15c9eacbbde342bb5b7561c46944f107bfb6c3/yarl-1.22.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cf71bf877efeac18b38d3930594c0948c82b64547c1cf420ba48722fe5509f6", size = 346965, upload-time = "2025-10-06T14:08:56.722Z" }, - { url = "https://files.pythonhosted.org/packages/c5/69/599e7cea8d0fcb1694323b0db0dda317fa3162f7b90166faddecf532166f/yarl-1.22.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:663e1cadaddae26be034a6ab6072449a8426ddb03d500f43daf952b74553bba0", size = 342026, upload-time = "2025-10-06T14:08:58.563Z" }, - { url = "https://files.pythonhosted.org/packages/95/6f/9dfd12c8bc90fea9eab39832ee32ea48f8e53d1256252a77b710c065c89f/yarl-1.22.0-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:6dcbb0829c671f305be48a7227918cfcd11276c2d637a8033a99a02b67bf9eda", size = 335637, upload-time = "2025-10-06T14:09:00.506Z" }, - { url = "https://files.pythonhosted.org/packages/57/2e/34c5b4eb9b07e16e873db5b182c71e5f06f9b5af388cdaa97736d79dd9a6/yarl-1.22.0-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:f0d97c18dfd9a9af4490631905a3f131a8e4c9e80a39353919e2cfed8f00aedc", size = 359082, upload-time = "2025-10-06T14:09:01.936Z" }, - { url = "https://files.pythonhosted.org/packages/31/71/fa7e10fb772d273aa1f096ecb8ab8594117822f683bab7d2c5a89914c92a/yarl-1.22.0-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:437840083abe022c978470b942ff832c3940b2ad3734d424b7eaffcd07f76737", size = 357811, upload-time = "2025-10-06T14:09:03.445Z" }, - { url = "https://files.pythonhosted.org/packages/26/da/11374c04e8e1184a6a03cf9c8f5688d3e5cec83ed6f31ad3481b3207f709/yarl-1.22.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:a899cbd98dce6f5d8de1aad31cb712ec0a530abc0a86bd6edaa47c1090138467", size = 351223, upload-time = "2025-10-06T14:09:05.401Z" }, - { url = "https://files.pythonhosted.org/packages/82/8f/e2d01f161b0c034a30410e375e191a5d27608c1f8693bab1a08b089ca096/yarl-1.22.0-cp310-cp310-win32.whl", hash = "sha256:595697f68bd1f0c1c159fcb97b661fc9c3f5db46498043555d04805430e79bea", size = 82118, upload-time = "2025-10-06T14:09:11.148Z" }, - { url = "https://files.pythonhosted.org/packages/62/46/94c76196642dbeae634c7a61ba3da88cd77bed875bf6e4a8bed037505aa6/yarl-1.22.0-cp310-cp310-win_amd64.whl", hash = "sha256:cb95a9b1adaa48e41815a55ae740cfda005758104049a640a398120bf02515ca", size = 86852, upload-time = "2025-10-06T14:09:12.958Z" }, - { url = "https://files.pythonhosted.org/packages/af/af/7df4f179d3b1a6dcb9a4bd2ffbc67642746fcafdb62580e66876ce83fff4/yarl-1.22.0-cp310-cp310-win_arm64.whl", hash = "sha256:b85b982afde6df99ecc996990d4ad7ccbdbb70e2a4ba4de0aecde5922ba98a0b", size = 82012, upload-time = "2025-10-06T14:09:14.664Z" }, - { url = "https://files.pythonhosted.org/packages/4d/27/5ab13fc84c76a0250afd3d26d5936349a35be56ce5785447d6c423b26d92/yarl-1.22.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:1ab72135b1f2db3fed3997d7e7dc1b80573c67138023852b6efb336a5eae6511", size = 141607, upload-time = "2025-10-06T14:09:16.298Z" }, - { url = "https://files.pythonhosted.org/packages/6a/a1/d065d51d02dc02ce81501d476b9ed2229d9a990818332242a882d5d60340/yarl-1.22.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:669930400e375570189492dc8d8341301578e8493aec04aebc20d4717f899dd6", size = 94027, upload-time = "2025-10-06T14:09:17.786Z" }, - { url = "https://files.pythonhosted.org/packages/c1/da/8da9f6a53f67b5106ffe902c6fa0164e10398d4e150d85838b82f424072a/yarl-1.22.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:792a2af6d58177ef7c19cbf0097aba92ca1b9cb3ffdd9c7470e156c8f9b5e028", size = 94963, upload-time = "2025-10-06T14:09:19.662Z" }, - { url = "https://files.pythonhosted.org/packages/68/fe/2c1f674960c376e29cb0bec1249b117d11738db92a6ccc4a530b972648db/yarl-1.22.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ea66b1c11c9150f1372f69afb6b8116f2dd7286f38e14ea71a44eee9ec51b9d", size = 368406, upload-time = "2025-10-06T14:09:21.402Z" }, - { url = "https://files.pythonhosted.org/packages/95/26/812a540e1c3c6418fec60e9bbd38e871eaba9545e94fa5eff8f4a8e28e1e/yarl-1.22.0-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3e2daa88dc91870215961e96a039ec73e4937da13cf77ce17f9cad0c18df3503", size = 336581, upload-time = "2025-10-06T14:09:22.98Z" }, - { url = "https://files.pythonhosted.org/packages/0b/f5/5777b19e26fdf98563985e481f8be3d8a39f8734147a6ebf459d0dab5a6b/yarl-1.22.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ba440ae430c00eee41509353628600212112cd5018d5def7e9b05ea7ac34eb65", size = 388924, upload-time = "2025-10-06T14:09:24.655Z" }, - { url = "https://files.pythonhosted.org/packages/86/08/24bd2477bd59c0bbd994fe1d93b126e0472e4e3df5a96a277b0a55309e89/yarl-1.22.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e6438cc8f23a9c1478633d216b16104a586b9761db62bfacb6425bac0a36679e", size = 392890, upload-time = "2025-10-06T14:09:26.617Z" }, - { url = "https://files.pythonhosted.org/packages/46/00/71b90ed48e895667ecfb1eaab27c1523ee2fa217433ed77a73b13205ca4b/yarl-1.22.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c52a6e78aef5cf47a98ef8e934755abf53953379b7d53e68b15ff4420e6683d", size = 365819, upload-time = "2025-10-06T14:09:28.544Z" }, - { url = "https://files.pythonhosted.org/packages/30/2d/f715501cae832651d3282387c6a9236cd26bd00d0ff1e404b3dc52447884/yarl-1.22.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:3b06bcadaac49c70f4c88af4ffcfbe3dc155aab3163e75777818092478bcbbe7", size = 363601, upload-time = "2025-10-06T14:09:30.568Z" }, - { url = "https://files.pythonhosted.org/packages/f8/f9/a678c992d78e394e7126ee0b0e4e71bd2775e4334d00a9278c06a6cce96a/yarl-1.22.0-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:6944b2dc72c4d7f7052683487e3677456050ff77fcf5e6204e98caf785ad1967", size = 358072, upload-time = "2025-10-06T14:09:32.528Z" }, - { url = "https://files.pythonhosted.org/packages/2c/d1/b49454411a60edb6fefdcad4f8e6dbba7d8019e3a508a1c5836cba6d0781/yarl-1.22.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:d5372ca1df0f91a86b047d1277c2aaf1edb32d78bbcefffc81b40ffd18f027ed", size = 385311, upload-time = "2025-10-06T14:09:34.634Z" }, - { url = "https://files.pythonhosted.org/packages/87/e5/40d7a94debb8448c7771a916d1861d6609dddf7958dc381117e7ba36d9e8/yarl-1.22.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:51af598701f5299012b8416486b40fceef8c26fc87dc6d7d1f6fc30609ea0aa6", size = 381094, upload-time = "2025-10-06T14:09:36.268Z" }, - { url = "https://files.pythonhosted.org/packages/35/d8/611cc282502381ad855448643e1ad0538957fc82ae83dfe7762c14069e14/yarl-1.22.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b266bd01fedeffeeac01a79ae181719ff848a5a13ce10075adbefc8f1daee70e", size = 370944, upload-time = "2025-10-06T14:09:37.872Z" }, - { url = "https://files.pythonhosted.org/packages/2d/df/fadd00fb1c90e1a5a8bd731fa3d3de2e165e5a3666a095b04e31b04d9cb6/yarl-1.22.0-cp311-cp311-win32.whl", hash = "sha256:a9b1ba5610a4e20f655258d5a1fdc7ebe3d837bb0e45b581398b99eb98b1f5ca", size = 81804, upload-time = "2025-10-06T14:09:39.359Z" }, - { url = "https://files.pythonhosted.org/packages/b5/f7/149bb6f45f267cb5c074ac40c01c6b3ea6d8a620d34b337f6321928a1b4d/yarl-1.22.0-cp311-cp311-win_amd64.whl", hash = "sha256:078278b9b0b11568937d9509b589ee83ef98ed6d561dfe2020e24a9fd08eaa2b", size = 86858, upload-time = "2025-10-06T14:09:41.068Z" }, - { url = "https://files.pythonhosted.org/packages/2b/13/88b78b93ad3f2f0b78e13bfaaa24d11cbc746e93fe76d8c06bf139615646/yarl-1.22.0-cp311-cp311-win_arm64.whl", hash = "sha256:b6a6f620cfe13ccec221fa312139135166e47ae169f8253f72a0abc0dae94376", size = 81637, upload-time = "2025-10-06T14:09:42.712Z" }, - { url = "https://files.pythonhosted.org/packages/75/ff/46736024fee3429b80a165a732e38e5d5a238721e634ab41b040d49f8738/yarl-1.22.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e340382d1afa5d32b892b3ff062436d592ec3d692aeea3bef3a5cfe11bbf8c6f", size = 142000, upload-time = "2025-10-06T14:09:44.631Z" }, - { url = "https://files.pythonhosted.org/packages/5a/9a/b312ed670df903145598914770eb12de1bac44599549b3360acc96878df8/yarl-1.22.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f1e09112a2c31ffe8d80be1b0988fa6a18c5d5cad92a9ffbb1c04c91bfe52ad2", size = 94338, upload-time = "2025-10-06T14:09:46.372Z" }, - { url = "https://files.pythonhosted.org/packages/ba/f5/0601483296f09c3c65e303d60c070a5c19fcdbc72daa061e96170785bc7d/yarl-1.22.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:939fe60db294c786f6b7c2d2e121576628468f65453d86b0fe36cb52f987bd74", size = 94909, upload-time = "2025-10-06T14:09:48.648Z" }, - { url = "https://files.pythonhosted.org/packages/60/41/9a1fe0b73dbcefce72e46cf149b0e0a67612d60bfc90fb59c2b2efdfbd86/yarl-1.22.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e1651bf8e0398574646744c1885a41198eba53dc8a9312b954073f845c90a8df", size = 372940, upload-time = "2025-10-06T14:09:50.089Z" }, - { url = "https://files.pythonhosted.org/packages/17/7a/795cb6dfee561961c30b800f0ed616b923a2ec6258b5def2a00bf8231334/yarl-1.22.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:b8a0588521a26bf92a57a1705b77b8b59044cdceccac7151bd8d229e66b8dedb", size = 345825, upload-time = "2025-10-06T14:09:52.142Z" }, - { url = "https://files.pythonhosted.org/packages/d7/93/a58f4d596d2be2ae7bab1a5846c4d270b894958845753b2c606d666744d3/yarl-1.22.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:42188e6a615c1a75bcaa6e150c3fe8f3e8680471a6b10150c5f7e83f47cc34d2", size = 386705, upload-time = "2025-10-06T14:09:54.128Z" }, - { url = "https://files.pythonhosted.org/packages/61/92/682279d0e099d0e14d7fd2e176bd04f48de1484f56546a3e1313cd6c8e7c/yarl-1.22.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f6d2cb59377d99718913ad9a151030d6f83ef420a2b8f521d94609ecc106ee82", size = 396518, upload-time = "2025-10-06T14:09:55.762Z" }, - { url = "https://files.pythonhosted.org/packages/db/0f/0d52c98b8a885aeda831224b78f3be7ec2e1aa4a62091f9f9188c3c65b56/yarl-1.22.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:50678a3b71c751d58d7908edc96d332af328839eea883bb554a43f539101277a", size = 377267, upload-time = "2025-10-06T14:09:57.958Z" }, - { url = "https://files.pythonhosted.org/packages/22/42/d2685e35908cbeaa6532c1fc73e89e7f2efb5d8a7df3959ea8e37177c5a3/yarl-1.22.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1e8fbaa7cec507aa24ea27a01456e8dd4b6fab829059b69844bd348f2d467124", size = 365797, upload-time = "2025-10-06T14:09:59.527Z" }, - { url = "https://files.pythonhosted.org/packages/a2/83/cf8c7bcc6355631762f7d8bdab920ad09b82efa6b722999dfb05afa6cfac/yarl-1.22.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:433885ab5431bc3d3d4f2f9bd15bfa1614c522b0f1405d62c4f926ccd69d04fa", size = 365535, upload-time = "2025-10-06T14:10:01.139Z" }, - { url = "https://files.pythonhosted.org/packages/25/e1/5302ff9b28f0c59cac913b91fe3f16c59a033887e57ce9ca5d41a3a94737/yarl-1.22.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:b790b39c7e9a4192dc2e201a282109ed2985a1ddbd5ac08dc56d0e121400a8f7", size = 382324, upload-time = "2025-10-06T14:10:02.756Z" }, - { url = "https://files.pythonhosted.org/packages/bf/cd/4617eb60f032f19ae3a688dc990d8f0d89ee0ea378b61cac81ede3e52fae/yarl-1.22.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:31f0b53913220599446872d757257be5898019c85e7971599065bc55065dc99d", size = 383803, upload-time = "2025-10-06T14:10:04.552Z" }, - { url = "https://files.pythonhosted.org/packages/59/65/afc6e62bb506a319ea67b694551dab4a7e6fb7bf604e9bd9f3e11d575fec/yarl-1.22.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a49370e8f711daec68d09b821a34e1167792ee2d24d405cbc2387be4f158b520", size = 374220, upload-time = "2025-10-06T14:10:06.489Z" }, - { url = "https://files.pythonhosted.org/packages/e7/3d/68bf18d50dc674b942daec86a9ba922d3113d8399b0e52b9897530442da2/yarl-1.22.0-cp312-cp312-win32.whl", hash = "sha256:70dfd4f241c04bd9239d53b17f11e6ab672b9f1420364af63e8531198e3f5fe8", size = 81589, upload-time = "2025-10-06T14:10:09.254Z" }, - { url = "https://files.pythonhosted.org/packages/c8/9a/6ad1a9b37c2f72874f93e691b2e7ecb6137fb2b899983125db4204e47575/yarl-1.22.0-cp312-cp312-win_amd64.whl", hash = "sha256:8884d8b332a5e9b88e23f60bb166890009429391864c685e17bd73a9eda9105c", size = 87213, upload-time = "2025-10-06T14:10:11.369Z" }, - { url = "https://files.pythonhosted.org/packages/44/c5/c21b562d1680a77634d748e30c653c3ca918beb35555cff24986fff54598/yarl-1.22.0-cp312-cp312-win_arm64.whl", hash = "sha256:ea70f61a47f3cc93bdf8b2f368ed359ef02a01ca6393916bc8ff877427181e74", size = 81330, upload-time = "2025-10-06T14:10:13.112Z" }, - { url = "https://files.pythonhosted.org/packages/ea/f3/d67de7260456ee105dc1d162d43a019ecad6b91e2f51809d6cddaa56690e/yarl-1.22.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8dee9c25c74997f6a750cd317b8ca63545169c098faee42c84aa5e506c819b53", size = 139980, upload-time = "2025-10-06T14:10:14.601Z" }, - { url = "https://files.pythonhosted.org/packages/01/88/04d98af0b47e0ef42597b9b28863b9060bb515524da0a65d5f4db160b2d5/yarl-1.22.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:01e73b85a5434f89fc4fe27dcda2aff08ddf35e4d47bbbea3bdcd25321af538a", size = 93424, upload-time = "2025-10-06T14:10:16.115Z" }, - { url = "https://files.pythonhosted.org/packages/18/91/3274b215fd8442a03975ce6bee5fe6aa57a8326b29b9d3d56234a1dca244/yarl-1.22.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:22965c2af250d20c873cdbee8ff958fb809940aeb2e74ba5f20aaf6b7ac8c70c", size = 93821, upload-time = "2025-10-06T14:10:17.993Z" }, - { url = "https://files.pythonhosted.org/packages/61/3a/caf4e25036db0f2da4ca22a353dfeb3c9d3c95d2761ebe9b14df8fc16eb0/yarl-1.22.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b4f15793aa49793ec8d1c708ab7f9eded1aa72edc5174cae703651555ed1b601", size = 373243, upload-time = "2025-10-06T14:10:19.44Z" }, - { url = "https://files.pythonhosted.org/packages/6e/9e/51a77ac7516e8e7803b06e01f74e78649c24ee1021eca3d6a739cb6ea49c/yarl-1.22.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5542339dcf2747135c5c85f68680353d5cb9ffd741c0f2e8d832d054d41f35a", size = 342361, upload-time = "2025-10-06T14:10:21.124Z" }, - { url = "https://files.pythonhosted.org/packages/d4/f8/33b92454789dde8407f156c00303e9a891f1f51a0330b0fad7c909f87692/yarl-1.22.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5c401e05ad47a75869c3ab3e35137f8468b846770587e70d71e11de797d113df", size = 387036, upload-time = "2025-10-06T14:10:22.902Z" }, - { url = "https://files.pythonhosted.org/packages/d9/9a/c5db84ea024f76838220280f732970aa4ee154015d7f5c1bfb60a267af6f/yarl-1.22.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:243dda95d901c733f5b59214d28b0120893d91777cb8aa043e6ef059d3cddfe2", size = 397671, upload-time = "2025-10-06T14:10:24.523Z" }, - { url = "https://files.pythonhosted.org/packages/11/c9/cd8538dc2e7727095e0c1d867bad1e40c98f37763e6d995c1939f5fdc7b1/yarl-1.22.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bec03d0d388060058f5d291a813f21c011041938a441c593374da6077fe21b1b", size = 377059, upload-time = "2025-10-06T14:10:26.406Z" }, - { url = "https://files.pythonhosted.org/packages/a1/b9/ab437b261702ced75122ed78a876a6dec0a1b0f5e17a4ac7a9a2482d8abe/yarl-1.22.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b0748275abb8c1e1e09301ee3cf90c8a99678a4e92e4373705f2a2570d581273", size = 365356, upload-time = "2025-10-06T14:10:28.461Z" }, - { url = "https://files.pythonhosted.org/packages/b2/9d/8e1ae6d1d008a9567877b08f0ce4077a29974c04c062dabdb923ed98e6fe/yarl-1.22.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:47fdb18187e2a4e18fda2c25c05d8251a9e4a521edaed757fef033e7d8498d9a", size = 361331, upload-time = "2025-10-06T14:10:30.541Z" }, - { url = "https://files.pythonhosted.org/packages/ca/5a/09b7be3905962f145b73beb468cdd53db8aa171cf18c80400a54c5b82846/yarl-1.22.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:c7044802eec4524fde550afc28edda0dd5784c4c45f0be151a2d3ba017daca7d", size = 382590, upload-time = "2025-10-06T14:10:33.352Z" }, - { url = "https://files.pythonhosted.org/packages/aa/7f/59ec509abf90eda5048b0bc3e2d7b5099dffdb3e6b127019895ab9d5ef44/yarl-1.22.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:139718f35149ff544caba20fce6e8a2f71f1e39b92c700d8438a0b1d2a631a02", size = 385316, upload-time = "2025-10-06T14:10:35.034Z" }, - { url = "https://files.pythonhosted.org/packages/e5/84/891158426bc8036bfdfd862fabd0e0fa25df4176ec793e447f4b85cf1be4/yarl-1.22.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e1b51bebd221006d3d2f95fbe124b22b247136647ae5dcc8c7acafba66e5ee67", size = 374431, upload-time = "2025-10-06T14:10:37.76Z" }, - { url = "https://files.pythonhosted.org/packages/bb/49/03da1580665baa8bef5e8ed34c6df2c2aca0a2f28bf397ed238cc1bbc6f2/yarl-1.22.0-cp313-cp313-win32.whl", hash = "sha256:d3e32536234a95f513bd374e93d717cf6b2231a791758de6c509e3653f234c95", size = 81555, upload-time = "2025-10-06T14:10:39.649Z" }, - { url = "https://files.pythonhosted.org/packages/9a/ee/450914ae11b419eadd067c6183ae08381cfdfcb9798b90b2b713bbebddda/yarl-1.22.0-cp313-cp313-win_amd64.whl", hash = "sha256:47743b82b76d89a1d20b83e60d5c20314cbd5ba2befc9cda8f28300c4a08ed4d", size = 86965, upload-time = "2025-10-06T14:10:41.313Z" }, - { url = "https://files.pythonhosted.org/packages/98/4d/264a01eae03b6cf629ad69bae94e3b0e5344741e929073678e84bf7a3e3b/yarl-1.22.0-cp313-cp313-win_arm64.whl", hash = "sha256:5d0fcda9608875f7d052eff120c7a5da474a6796fe4d83e152e0e4d42f6d1a9b", size = 81205, upload-time = "2025-10-06T14:10:43.167Z" }, - { url = "https://files.pythonhosted.org/packages/88/fc/6908f062a2f77b5f9f6d69cecb1747260831ff206adcbc5b510aff88df91/yarl-1.22.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:719ae08b6972befcba4310e49edb1161a88cdd331e3a694b84466bd938a6ab10", size = 146209, upload-time = "2025-10-06T14:10:44.643Z" }, - { url = "https://files.pythonhosted.org/packages/65/47/76594ae8eab26210b4867be6f49129861ad33da1f1ebdf7051e98492bf62/yarl-1.22.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:47d8a5c446df1c4db9d21b49619ffdba90e77c89ec6e283f453856c74b50b9e3", size = 95966, upload-time = "2025-10-06T14:10:46.554Z" }, - { url = "https://files.pythonhosted.org/packages/ab/ce/05e9828a49271ba6b5b038b15b3934e996980dd78abdfeb52a04cfb9467e/yarl-1.22.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:cfebc0ac8333520d2d0423cbbe43ae43c8838862ddb898f5ca68565e395516e9", size = 97312, upload-time = "2025-10-06T14:10:48.007Z" }, - { url = "https://files.pythonhosted.org/packages/d1/c5/7dffad5e4f2265b29c9d7ec869c369e4223166e4f9206fc2243ee9eea727/yarl-1.22.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4398557cbf484207df000309235979c79c4356518fd5c99158c7d38203c4da4f", size = 361967, upload-time = "2025-10-06T14:10:49.997Z" }, - { url = "https://files.pythonhosted.org/packages/50/b2/375b933c93a54bff7fc041e1a6ad2c0f6f733ffb0c6e642ce56ee3b39970/yarl-1.22.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2ca6fd72a8cd803be290d42f2dec5cdcd5299eeb93c2d929bf060ad9efaf5de0", size = 323949, upload-time = "2025-10-06T14:10:52.004Z" }, - { url = "https://files.pythonhosted.org/packages/66/50/bfc2a29a1d78644c5a7220ce2f304f38248dc94124a326794e677634b6cf/yarl-1.22.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ca1f59c4e1ab6e72f0a23c13fca5430f889634166be85dbf1013683e49e3278e", size = 361818, upload-time = "2025-10-06T14:10:54.078Z" }, - { url = "https://files.pythonhosted.org/packages/46/96/f3941a46af7d5d0f0498f86d71275696800ddcdd20426298e572b19b91ff/yarl-1.22.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6c5010a52015e7c70f86eb967db0f37f3c8bd503a695a49f8d45700144667708", size = 372626, upload-time = "2025-10-06T14:10:55.767Z" }, - { url = "https://files.pythonhosted.org/packages/c1/42/8b27c83bb875cd89448e42cd627e0fb971fa1675c9ec546393d18826cb50/yarl-1.22.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d7672ecf7557476642c88497c2f8d8542f8e36596e928e9bcba0e42e1e7d71f", size = 341129, upload-time = "2025-10-06T14:10:57.985Z" }, - { url = "https://files.pythonhosted.org/packages/49/36/99ca3122201b382a3cf7cc937b95235b0ac944f7e9f2d5331d50821ed352/yarl-1.22.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:3b7c88eeef021579d600e50363e0b6ee4f7f6f728cd3486b9d0f3ee7b946398d", size = 346776, upload-time = "2025-10-06T14:10:59.633Z" }, - { url = "https://files.pythonhosted.org/packages/85/b4/47328bf996acd01a4c16ef9dcd2f59c969f495073616586f78cd5f2efb99/yarl-1.22.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:f4afb5c34f2c6fecdcc182dfcfc6af6cccf1aa923eed4d6a12e9d96904e1a0d8", size = 334879, upload-time = "2025-10-06T14:11:01.454Z" }, - { url = "https://files.pythonhosted.org/packages/c2/ad/b77d7b3f14a4283bffb8e92c6026496f6de49751c2f97d4352242bba3990/yarl-1.22.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:59c189e3e99a59cf8d83cbb31d4db02d66cda5a1a4374e8a012b51255341abf5", size = 350996, upload-time = "2025-10-06T14:11:03.452Z" }, - { url = "https://files.pythonhosted.org/packages/81/c8/06e1d69295792ba54d556f06686cbd6a7ce39c22307100e3fb4a2c0b0a1d/yarl-1.22.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:5a3bf7f62a289fa90f1990422dc8dff5a458469ea71d1624585ec3a4c8d6960f", size = 356047, upload-time = "2025-10-06T14:11:05.115Z" }, - { url = "https://files.pythonhosted.org/packages/4b/b8/4c0e9e9f597074b208d18cef227d83aac36184bfbc6eab204ea55783dbc5/yarl-1.22.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:de6b9a04c606978fdfe72666fa216ffcf2d1a9f6a381058d4378f8d7b1e5de62", size = 342947, upload-time = "2025-10-06T14:11:08.137Z" }, - { url = "https://files.pythonhosted.org/packages/e0/e5/11f140a58bf4c6ad7aca69a892bff0ee638c31bea4206748fc0df4ebcb3a/yarl-1.22.0-cp313-cp313t-win32.whl", hash = "sha256:1834bb90991cc2999f10f97f5f01317f99b143284766d197e43cd5b45eb18d03", size = 86943, upload-time = "2025-10-06T14:11:10.284Z" }, - { url = "https://files.pythonhosted.org/packages/31/74/8b74bae38ed7fe6793d0c15a0c8207bbb819cf287788459e5ed230996cdd/yarl-1.22.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff86011bd159a9d2dfc89c34cfd8aff12875980e3bd6a39ff097887520e60249", size = 93715, upload-time = "2025-10-06T14:11:11.739Z" }, - { url = "https://files.pythonhosted.org/packages/69/66/991858aa4b5892d57aef7ee1ba6b4d01ec3b7eb3060795d34090a3ca3278/yarl-1.22.0-cp313-cp313t-win_arm64.whl", hash = "sha256:7861058d0582b847bc4e3a4a4c46828a410bca738673f35a29ba3ca5db0b473b", size = 83857, upload-time = "2025-10-06T14:11:13.586Z" }, - { url = "https://files.pythonhosted.org/packages/46/b3/e20ef504049f1a1c54a814b4b9bed96d1ac0e0610c3b4da178f87209db05/yarl-1.22.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:34b36c2c57124530884d89d50ed2c1478697ad7473efd59cfd479945c95650e4", size = 140520, upload-time = "2025-10-06T14:11:15.465Z" }, - { url = "https://files.pythonhosted.org/packages/e4/04/3532d990fdbab02e5ede063676b5c4260e7f3abea2151099c2aa745acc4c/yarl-1.22.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:0dd9a702591ca2e543631c2a017e4a547e38a5c0f29eece37d9097e04a7ac683", size = 93504, upload-time = "2025-10-06T14:11:17.106Z" }, - { url = "https://files.pythonhosted.org/packages/11/63/ff458113c5c2dac9a9719ac68ee7c947cb621432bcf28c9972b1c0e83938/yarl-1.22.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:594fcab1032e2d2cc3321bb2e51271e7cd2b516c7d9aee780ece81b07ff8244b", size = 94282, upload-time = "2025-10-06T14:11:19.064Z" }, - { url = "https://files.pythonhosted.org/packages/a7/bc/315a56aca762d44a6aaaf7ad253f04d996cb6b27bad34410f82d76ea8038/yarl-1.22.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f3d7a87a78d46a2e3d5b72587ac14b4c16952dd0887dbb051451eceac774411e", size = 372080, upload-time = "2025-10-06T14:11:20.996Z" }, - { url = "https://files.pythonhosted.org/packages/3f/3f/08e9b826ec2e099ea6e7c69a61272f4f6da62cb5b1b63590bb80ca2e4a40/yarl-1.22.0-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:852863707010316c973162e703bddabec35e8757e67fcb8ad58829de1ebc8590", size = 338696, upload-time = "2025-10-06T14:11:22.847Z" }, - { url = "https://files.pythonhosted.org/packages/e3/9f/90360108e3b32bd76789088e99538febfea24a102380ae73827f62073543/yarl-1.22.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:131a085a53bfe839a477c0845acf21efc77457ba2bcf5899618136d64f3303a2", size = 387121, upload-time = "2025-10-06T14:11:24.889Z" }, - { url = "https://files.pythonhosted.org/packages/98/92/ab8d4657bd5b46a38094cfaea498f18bb70ce6b63508fd7e909bd1f93066/yarl-1.22.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:078a8aefd263f4d4f923a9677b942b445a2be970ca24548a8102689a3a8ab8da", size = 394080, upload-time = "2025-10-06T14:11:27.307Z" }, - { url = "https://files.pythonhosted.org/packages/f5/e7/d8c5a7752fef68205296201f8ec2bf718f5c805a7a7e9880576c67600658/yarl-1.22.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bca03b91c323036913993ff5c738d0842fc9c60c4648e5c8d98331526df89784", size = 372661, upload-time = "2025-10-06T14:11:29.387Z" }, - { url = "https://files.pythonhosted.org/packages/b6/2e/f4d26183c8db0bb82d491b072f3127fb8c381a6206a3a56332714b79b751/yarl-1.22.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:68986a61557d37bb90d3051a45b91fa3d5c516d177dfc6dd6f2f436a07ff2b6b", size = 364645, upload-time = "2025-10-06T14:11:31.423Z" }, - { url = "https://files.pythonhosted.org/packages/80/7c/428e5812e6b87cd00ee8e898328a62c95825bf37c7fa87f0b6bb2ad31304/yarl-1.22.0-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:4792b262d585ff0dff6bcb787f8492e40698443ec982a3568c2096433660c694", size = 355361, upload-time = "2025-10-06T14:11:33.055Z" }, - { url = "https://files.pythonhosted.org/packages/ec/2a/249405fd26776f8b13c067378ef4d7dd49c9098d1b6457cdd152a99e96a9/yarl-1.22.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:ebd4549b108d732dba1d4ace67614b9545b21ece30937a63a65dd34efa19732d", size = 381451, upload-time = "2025-10-06T14:11:35.136Z" }, - { url = "https://files.pythonhosted.org/packages/67/a8/fb6b1adbe98cf1e2dd9fad71003d3a63a1bc22459c6e15f5714eb9323b93/yarl-1.22.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:f87ac53513d22240c7d59203f25cc3beac1e574c6cd681bbfd321987b69f95fd", size = 383814, upload-time = "2025-10-06T14:11:37.094Z" }, - { url = "https://files.pythonhosted.org/packages/d9/f9/3aa2c0e480fb73e872ae2814c43bc1e734740bb0d54e8cb2a95925f98131/yarl-1.22.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:22b029f2881599e2f1b06f8f1db2ee63bd309e2293ba2d566e008ba12778b8da", size = 370799, upload-time = "2025-10-06T14:11:38.83Z" }, - { url = "https://files.pythonhosted.org/packages/50/3c/af9dba3b8b5eeb302f36f16f92791f3ea62e3f47763406abf6d5a4a3333b/yarl-1.22.0-cp314-cp314-win32.whl", hash = "sha256:6a635ea45ba4ea8238463b4f7d0e721bad669f80878b7bfd1f89266e2ae63da2", size = 82990, upload-time = "2025-10-06T14:11:40.624Z" }, - { url = "https://files.pythonhosted.org/packages/ac/30/ac3a0c5bdc1d6efd1b41fa24d4897a4329b3b1e98de9449679dd327af4f0/yarl-1.22.0-cp314-cp314-win_amd64.whl", hash = "sha256:0d6e6885777af0f110b0e5d7e5dda8b704efed3894da26220b7f3d887b839a79", size = 88292, upload-time = "2025-10-06T14:11:42.578Z" }, - { url = "https://files.pythonhosted.org/packages/df/0a/227ab4ff5b998a1b7410abc7b46c9b7a26b0ca9e86c34ba4b8d8bc7c63d5/yarl-1.22.0-cp314-cp314-win_arm64.whl", hash = "sha256:8218f4e98d3c10d683584cb40f0424f4b9fd6e95610232dd75e13743b070ee33", size = 82888, upload-time = "2025-10-06T14:11:44.863Z" }, - { url = "https://files.pythonhosted.org/packages/06/5e/a15eb13db90abd87dfbefb9760c0f3f257ac42a5cac7e75dbc23bed97a9f/yarl-1.22.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:45c2842ff0e0d1b35a6bf1cd6c690939dacb617a70827f715232b2e0494d55d1", size = 146223, upload-time = "2025-10-06T14:11:46.796Z" }, - { url = "https://files.pythonhosted.org/packages/18/82/9665c61910d4d84f41a5bf6837597c89e665fa88aa4941080704645932a9/yarl-1.22.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:d947071e6ebcf2e2bee8fce76e10faca8f7a14808ca36a910263acaacef08eca", size = 95981, upload-time = "2025-10-06T14:11:48.845Z" }, - { url = "https://files.pythonhosted.org/packages/5d/9a/2f65743589809af4d0a6d3aa749343c4b5f4c380cc24a8e94a3c6625a808/yarl-1.22.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:334b8721303e61b00019474cc103bdac3d7b1f65e91f0bfedeec2d56dfe74b53", size = 97303, upload-time = "2025-10-06T14:11:50.897Z" }, - { url = "https://files.pythonhosted.org/packages/b0/ab/5b13d3e157505c43c3b43b5a776cbf7b24a02bc4cccc40314771197e3508/yarl-1.22.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1e7ce67c34138a058fd092f67d07a72b8e31ff0c9236e751957465a24b28910c", size = 361820, upload-time = "2025-10-06T14:11:52.549Z" }, - { url = "https://files.pythonhosted.org/packages/fb/76/242a5ef4677615cf95330cfc1b4610e78184400699bdda0acb897ef5e49a/yarl-1.22.0-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:d77e1b2c6d04711478cb1c4ab90db07f1609ccf06a287d5607fcd90dc9863acf", size = 323203, upload-time = "2025-10-06T14:11:54.225Z" }, - { url = "https://files.pythonhosted.org/packages/8c/96/475509110d3f0153b43d06164cf4195c64d16999e0c7e2d8a099adcd6907/yarl-1.22.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c4647674b6150d2cae088fc07de2738a84b8bcedebef29802cf0b0a82ab6face", size = 363173, upload-time = "2025-10-06T14:11:56.069Z" }, - { url = "https://files.pythonhosted.org/packages/c9/66/59db471aecfbd559a1fd48aedd954435558cd98c7d0da8b03cc6c140a32c/yarl-1.22.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:efb07073be061c8f79d03d04139a80ba33cbd390ca8f0297aae9cce6411e4c6b", size = 373562, upload-time = "2025-10-06T14:11:58.783Z" }, - { url = "https://files.pythonhosted.org/packages/03/1f/c5d94abc91557384719da10ff166b916107c1b45e4d0423a88457071dd88/yarl-1.22.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e51ac5435758ba97ad69617e13233da53908beccc6cfcd6c34bbed8dcbede486", size = 339828, upload-time = "2025-10-06T14:12:00.686Z" }, - { url = "https://files.pythonhosted.org/packages/5f/97/aa6a143d3afba17b6465733681c70cf175af89f76ec8d9286e08437a7454/yarl-1.22.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:33e32a0dd0c8205efa8e83d04fc9f19313772b78522d1bdc7d9aed706bfd6138", size = 347551, upload-time = "2025-10-06T14:12:02.628Z" }, - { url = "https://files.pythonhosted.org/packages/43/3c/45a2b6d80195959239a7b2a8810506d4eea5487dce61c2a3393e7fc3c52e/yarl-1.22.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:bf4a21e58b9cde0e401e683ebd00f6ed30a06d14e93f7c8fd059f8b6e8f87b6a", size = 334512, upload-time = "2025-10-06T14:12:04.871Z" }, - { url = "https://files.pythonhosted.org/packages/86/a0/c2ab48d74599c7c84cb104ebd799c5813de252bea0f360ffc29d270c2caa/yarl-1.22.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:e4b582bab49ac33c8deb97e058cd67c2c50dac0dd134874106d9c774fd272529", size = 352400, upload-time = "2025-10-06T14:12:06.624Z" }, - { url = "https://files.pythonhosted.org/packages/32/75/f8919b2eafc929567d3d8411f72bdb1a2109c01caaab4ebfa5f8ffadc15b/yarl-1.22.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:0b5bcc1a9c4839e7e30b7b30dd47fe5e7e44fb7054ec29b5bb8d526aa1041093", size = 357140, upload-time = "2025-10-06T14:12:08.362Z" }, - { url = "https://files.pythonhosted.org/packages/cf/72/6a85bba382f22cf78add705d8c3731748397d986e197e53ecc7835e76de7/yarl-1.22.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:c0232bce2170103ec23c454e54a57008a9a72b5d1c3105dc2496750da8cfa47c", size = 341473, upload-time = "2025-10-06T14:12:10.994Z" }, - { url = "https://files.pythonhosted.org/packages/35/18/55e6011f7c044dc80b98893060773cefcfdbf60dfefb8cb2f58b9bacbd83/yarl-1.22.0-cp314-cp314t-win32.whl", hash = "sha256:8009b3173bcd637be650922ac455946197d858b3630b6d8787aa9e5c4564533e", size = 89056, upload-time = "2025-10-06T14:12:13.317Z" }, - { url = "https://files.pythonhosted.org/packages/f9/86/0f0dccb6e59a9e7f122c5afd43568b1d31b8ab7dda5f1b01fb5c7025c9a9/yarl-1.22.0-cp314-cp314t-win_amd64.whl", hash = "sha256:9fb17ea16e972c63d25d4a97f016d235c78dd2344820eb35bc034bc32012ee27", size = 96292, upload-time = "2025-10-06T14:12:15.398Z" }, - { url = "https://files.pythonhosted.org/packages/48/b7/503c98092fb3b344a179579f55814b613c1fbb1c23b3ec14a7b008a66a6e/yarl-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:9f6d73c1436b934e3f01df1e1b21ff765cd1d28c77dfb9ace207f746d4610ee1", size = 85171, upload-time = "2025-10-06T14:12:16.935Z" }, - { url = "https://files.pythonhosted.org/packages/73/ae/b48f95715333080afb75a4504487cbe142cae1268afc482d06692d605ae6/yarl-1.22.0-py3-none-any.whl", hash = "sha256:1380560bdba02b6b6c90de54133c81c9f2a453dee9912fe58c1dcced1edb7cff", size = 46814, upload-time = "2025-10-06T14:12:53.872Z" }, -] - [[package]] name = "zipp" version = "3.23.0"