chore(py): Update generated schema typing files after local changes

Author: MengqinShen-GL

py/packages/genkit/src/genkit/ai/_aio.py

@@ -338,24 +338,23 @@ async def retrieve(
 
     async def embed(
         self,
-        embedder: str | EmbedderRef | None = None, # Allow EmbedderRef
+        embedder: str | EmbedderRef | None = None,
         documents: list[Document] | None = None,
         options: dict[str, Any] | None = None,
     ) -> EmbedResponse:
-
         embedder_name: str
         embedder_config: dict[str, Any] = {}
 
         if isinstance(embedder, EmbedderRef):
             embedder_name = embedder.name
             embedder_config = embedder.config or {}
             if embedder.version:
-                embedder_config['version'] = embedder.version # Handle version from ref
+                embedder_config['version'] = embedder.version  # Handle version from ref
         elif isinstance(embedder, str):
             embedder_name = embedder
         else:
             # Handle case where embedder is None
-            raise ValueError("Embedder must be specified as a string name or an EmbedderRef.")
+            raise ValueError('Embedder must be specified as a string name or an EmbedderRef.')
 
         # Merge options passed to embed() with config from EmbedderRef
         final_options = {**(embedder_config or {}), **(options or {})}

py/packages/genkit/src/genkit/ai/_registry.py

@@ -478,12 +478,12 @@ def define_embedder(
             if options.dimensions:
                 embedder_metadata['embedder']['dimensions'] = options.dimensions
             if options.supports:
-                embedder_metadata['embedder']['supports'] = options.supports.model_dump(exclude_none=True,
-                                                                                        by_alias=True)
+                embedder_metadata['embedder']['supports'] = options.supports.model_dump(
+                    exclude_none=True, by_alias=True
+                )
             if options.config_schema:
                 embedder_metadata['embedder']['customOptions'] = to_json_schema(options.config_schema)
 
-        # Fallback to default if 'embedder' key doesn't exist
         if 'embedder' not in embedder_metadata:
             embedder_metadata['embedder'] = {}
 

py/packages/genkit/src/genkit/blocks/embedding.py

@@ -22,33 +22,14 @@
 from genkit.ai import ActionKind
 from genkit.core.action import ActionMetadata
 from genkit.core.schema import to_json_schema
-from genkit.core.typing import EmbedRequest, EmbedResponse, EmbedderOptions, EmbedderRef, EmbedderSupports,EmbedderFn
+from genkit.core.typing import EmbedRequest, EmbedResponse, EmbedderOptions, EmbedderRef, EmbedderSupports, EmbedderFn
 from pydantic import BaseModel
-# type EmbedderFn = Callable[[EmbedRequest], EmbedResponse]
-# EmbedderFn = Callable[[EmbedRequest], EmbedResponse]
 
 
-# def embedder_action_metadata(
-#     name: str,
-#     info: dict[str, Any] | None = None,
-#     config_schema: Any | None = None,
-# ) -> ActionMetadata:
-#     """Generates an ActionMetadata for embedders."""
-#     info = info if info is not None else {}
-#     return ActionMetadata(
-#         kind=ActionKind.EMBEDDER,
-#         name=name,
-#         input_json_schema=to_json_schema(EmbedRequest),
-#         output_json_schema=to_json_schema(EmbedResponse),
-#         metadata={'embedder': {**info, 'customOptions': to_json_schema(config_schema) if config_schema else None}},
-#     )
-
 def embedder_action_metadata(
     name: str,
-    #info: dict[str, Any] | None = None,
     options: EmbedderOptions | None = None,
 ) -> ActionMetadata:
-
     options = options if options is not None else EmbedderOptions()
     embedder_metadata_dict = {'embedder': {}}
 
@@ -60,18 +41,19 @@ def embedder_action_metadata(
     if options.supports:
         embedder_metadata_dict['embedder']['supports'] = options.supports.model_dump(exclude_none=True, by_alias=True)
 
-    embedder_metadata_dict['embedder']['customOptions'] = to_json_schema(options.config_schema) if options.config_schema else None
+    embedder_metadata_dict['embedder']['customOptions'] = (
+        to_json_schema(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,
-
     )
 
-# New helper function
+
 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)

py/packages/genkit/src/genkit/core/typing.py

@@ -33,7 +33,7 @@
     from enum import StrEnum  # noqa
 
 
-from typing import Any, Awaitable, Callable, Protocol, TypeVar
+from typing import Any
 
 from pydantic import BaseModel, ConfigDict, Field, RootModel
 
@@ -1013,47 +1013,3 @@ class ModelResponse(BaseModel):
     raw: Raw | None = None
     request: Request | None = None
     operation: Operation | None = None
-
-T = TypeVar('T')
-P = TypeVar('P')
-
-class EmbedderFn(Protocol[T, P]):
-    """Protocol for embedder function."""
-
-    def __call__(
-        self,
-        request: EmbedRequest,
-        options: dict[str, Any] | None = None,
-    ) -> Awaitable[EmbedResponse]:
-        """Embeds a list of documents.
-
-        Args:
-            request: The embed request.
-            options: Configuration options for the embedder.
-
-        Returns:
-            The embed response.
-        """
-        ...
-
-class EmbedderRef(BaseModel):
-
-    model_config = ConfigDict(extra='forbid', populate_by_name=True)
-    name: str
-    config: dict[str, Any] | None = None
-    version: str | None = None
-
-class EmbedderSupports(BaseModel):
-
-    model_config = ConfigDict(extra='forbid', populate_by_name=True)
-    input: list[str] | None = None
-    multilingual: bool | None = None
-
-class EmbedderOptions(BaseModel):
-
-    model_config = ConfigDict(extra='forbid', populate_by_name=True)
-    dimensions: int | None = None
-    label: str | None = None
-    supports: EmbedderSupports | None = None
-    config_schema: dict[str, Any] | None = Field(None, alias='configSchema')
-

py/packages/genkit/tests/genkit/blocks/embedding_test.py

@@ -18,7 +18,15 @@
 
 from genkit.blocks.embedding import embedder_action_metadata, create_embedder_ref
 from genkit.core.action import ActionMetadata, Action
-from genkit.core.typing import EmbedderOptions, EmbedderSupports, EmbedderRef, EmbedRequest, EmbedResponse, Embedding, Part
+from genkit.core.typing import (
+    EmbedderOptions,
+    EmbedderSupports,
+    EmbedderRef,
+    EmbedRequest,
+    EmbedResponse,
+    Embedding,
+    Part,
+)
 from genkit.blocks.document import Document
 from genkit.core.schema import to_json_schema
 from genkit.core.action.types import ActionResponse
@@ -28,19 +36,6 @@
 from pydantic import BaseModel
 
 
-# def test_embedder_action_metadata():
-#     """Test for embedder_action_metadata."""
-#     action_metadata = embedder_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 == {'embedder': {'customOptions': None, 'label': 'test_label'}}
-
 def test_embedder_action_metadata():
     """Test for embedder_action_metadata with basic options."""
     options = EmbedderOptions(label='Test Embedder', dimensions=128)
@@ -60,8 +55,10 @@ def test_embedder_action_metadata():
         }
     }
 
+
 def test_embedder_action_metadata_with_supports_and_config_schema():
     """Test for embedder_action_metadata with supports and config_schema."""
+
     class CustomConfig(BaseModel):
         param1: str
         param2: int
@@ -70,7 +67,7 @@ class CustomConfig(BaseModel):
         label='Advanced Embedder',
         dimensions=256,
         supports=EmbedderSupports(input=['text', 'image'], multilingual=True),
-        config_schema=to_json_schema(CustomConfig) # Pass the JSON schema directly
+        config_schema=to_json_schema(CustomConfig),  # Pass the JSON schema directly
     )
     action_metadata = embedder_action_metadata(
         name='advanced_model',
@@ -94,19 +91,22 @@ class CustomConfig(BaseModel):
         'required': ['param1', 'param2'],
     }
 
+
 def test_embedder_action_metadata_no_options():
     """Test embedder_action_metadata when no options are provided."""
     action_metadata = embedder_action_metadata(name='default_model')
     assert isinstance(action_metadata, ActionMetadata)
     assert action_metadata.metadata == {'embedder': {'customOptions': None, 'dimensions': None}}
 
+
 def test_create_embedder_ref_basic():
     """Test basic creation of EmbedderRef."""
     ref = create_embedder_ref('my-embedder')
     assert ref.name == 'my-embedder'
     assert ref.config is None
     assert ref.version is None
 
+
 def test_create_embedder_ref_with_config():
     """Test creation of EmbedderRef with configuration."""
     config = {'temperature': 0.5, 'max_tokens': 100}
@@ -115,13 +115,15 @@ def test_create_embedder_ref_with_config():
     assert ref.config == config
     assert ref.version is None
 
+
 def test_create_embedder_ref_with_version():
     """Test creation of EmbedderRef with a version."""
     ref = create_embedder_ref('versioned-embedder', version='v1.0')
     assert ref.name == 'versioned-embedder'
     assert ref.config is None
     assert ref.version == 'v1.0'
 
+
 def test_create_embedder_ref_with_config_and_version():
     """Test creation of EmbedderRef with both config and version."""
     config = {'task_type': 'retrieval'}
@@ -130,8 +132,10 @@ def test_create_embedder_ref_with_config_and_version():
     assert ref.config == config
     assert ref.version == 'beta'
 
+
 class MockGenkitRegistry:
     """A mock registry to simulate action lookup."""
+
     def __init__(self):
         self.actions = {}
 
@@ -145,7 +149,7 @@ def register_action(self, name, kind, fn, metadata, description):
         async def mock_arun_side_effect(request, *args, **kwargs):
             # Call the actual (fake) embedder function directly
             embed_response = await fn(request)
-            return ActionResponse(response=embed_response, trace_id="mock_trace_id")
+            return ActionResponse(response=embed_response, trace_id='mock_trace_id')
 
         mock_action.arun = AsyncMock(side_effect=mock_arun_side_effect)
         self.actions[(kind, name)] = mock_action
@@ -154,62 +158,52 @@ async def mock_arun_side_effect(request, *args, **kwargs):
     def lookup_action(self, kind, name):
         return self.actions.get((kind, name))
 
+
 @pytest.fixture
 def mock_genkit_instance():
     """Fixture for a Genkit instance with a mock registry."""
     registry = MockGenkitRegistry()
-    genkit_instance = Genkit()  # Create Genkit instance without registry argument
-    genkit_instance.registry = registry  # Assign the mock registry to the instance
+    genkit_instance = Genkit()
+    genkit_instance.registry = registry
     return genkit_instance, registry
 
+
 @pytest.mark.asyncio
 async def test_embed_with_embedder_ref(mock_genkit_instance):
     """Test the embed method using EmbedderRef."""
     genkit_instance, registry = mock_genkit_instance
 
-    # Define an embedder in the mock registry
     async def fake_embedder_fn(request: EmbedRequest) -> EmbedResponse:
         return EmbedResponse(embeddings=[Embedding(embedding=[1.0, 2.0, 3.0])])
 
     embedder_options = EmbedderOptions(
         label='Fake Embedder',
         dimensions=3,
         supports=EmbedderSupports(input=['text']),
-        config_schema={'type': 'object', 'properties': {'param': {'type': 'string'}}}
+        config_schema={'type': 'object', 'properties': {'param': {'type': 'string'}}},
     )
     registry.register_action(
         name='my-plugin/my-embedder',
         kind='embedder',
         fn=fake_embedder_fn,
         metadata=embedder_action_metadata('my-plugin/my-embedder', options=embedder_options).metadata,
-        description='A fake embedder for testing'
+        description='A fake embedder for testing',
     )
 
-    # Create an EmbedderRef
-    embedder_ref = create_embedder_ref(
-        'my-plugin/my-embedder',
-        config={'param': 'value'},
-        version='v1'
-    )
+    embedder_ref = create_embedder_ref('my-plugin/my-embedder', config={'param': 'value'}, version='v1')
 
     documents = [Document.from_text('hello world')]
 
-    # Call the embed method
     response = await genkit_instance.embed(
-        embedder=embedder_ref,
-        documents=documents,
-        options={'additional_option': True}
+        embedder=embedder_ref, documents=documents, options={'additional_option': True}
     )
 
-    # Assertions
     assert response.embeddings[0].embedding == [1.0, 2.0, 3.0]
 
-    # Verify that lookup_action was called correctly
     embed_action = registry.lookup_action('embedder', 'my-plugin/my-embedder')
     assert embed_action is not None
     embed_action.arun.assert_called_once()
 
-    # Get the EmbedRequest passed to the mock action
     called_request = embed_action.arun.call_args[0][0]
     assert isinstance(called_request, EmbedRequest)
     assert called_request.input == documents
@@ -231,27 +225,26 @@ async def fake_embedder_fn(request: EmbedRequest) -> EmbedResponse:
         kind='embedder',
         fn=fake_embedder_fn,
         metadata=embedder_action_metadata('another-embedder', options=embedder_options).metadata,
-        description='Another fake embedder'
+        description='Another fake embedder',
     )
 
     documents = [Document.from_text('test text')]
 
     response = await genkit_instance.embed(
-        embedder='another-embedder',
-        documents=documents,
-        options={'custom_setting': 'high'}
+        embedder='another-embedder', documents=documents, options={'custom_setting': 'high'}
     )
 
     assert response.embeddings[0].embedding == [4.0, 5.0, 6.0]
     embed_action = registry.lookup_action('embedder', 'another-embedder')
     called_request = embed_action.arun.call_args[0][0]
     assert called_request.options == {'custom_setting': 'high'}
 
+
 @pytest.mark.asyncio
 async def test_embed_missing_embedder_raises_error(mock_genkit_instance):
     """Test that embedding with a missing embedder raises an error."""
     genkit_instance, _ = mock_genkit_instance
     documents = [Document.from_text('some text')]
 
-    with pytest.raises(ValueError, match="Embedder must be specified as a string name or an EmbedderRef."):
+    with pytest.raises(ValueError, match='Embedder must be specified as a string name or an EmbedderRef.'):
         await genkit_instance.embed(documents=documents)

py/packages/genkit/tests/genkit/blocks/prompt_test.py

@@ -202,7 +202,8 @@ def test_tool(input: ToolInput):
     ),
 ]
 
-@pytest.mark.skip(reason="issues when running on CI")
+
+@pytest.mark.skip(reason='issues when running on CI')
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     'test_case, prompt, input, input_option, context, want_rendered',

py/plugins/compat-oai/src/genkit/plugins/compat_oai/openai_plugin.py

@@ -196,7 +196,7 @@ def list_actions(self) -> list[ActionMetadata]:
                             config_schema=to_json_schema(Embedding),
                             label=f'OpenAI Embedding - {_name}',
                             supports=EmbedderSupports(input=['text']),
-                        )
+                        ),
                     )
                 )
             else: