feat: PluginV2 refactor #4132
feat: PluginV2 refactor #4132
Conversation
Summary of ChangesHello @huangjeff5, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a fundamental refactoring of the plugin system, transitioning to an asynchronous, lazy-loading architecture. This change streamlines how external functionalities are integrated and managed within the framework, improving efficiency and flexibility by resolving actions on demand rather than during initial setup. The update impacts the core registry, reflection mechanisms, and all existing plugins, ensuring a consistent and modern approach to extensibility. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
| self._runtime_dir.mkdir(parents=True, exist_ok=True) | ||
| runtime_file_path = _create_and_write_runtime_file(self._runtime_dir, self.spec) | ||
| _register_atexit_cleanup_handler(runtime_file_path) | ||
| self._runtime_file_path = _create_and_write_runtime_file(self._runtime_dir, self.spec) |
There was a problem hiding this comment.
investigate some potential weirdness in runtime file
seems like first time we start (no runtime file) there is some weird slowness
|
|
||
| from pydantic import BaseModel, ConfigDict, Field | ||
|
|
||
| if sys.version_info < (3, 11): |
There was a problem hiding this comment.
why did we erase this? not sure if needed.
| from strenum import StrEnum | ||
| else: | ||
| from enum import StrEnum | ||
| from enum import StrEnum |
There was a problem hiding this comment.
why did we erase this? not sure if needed.
| try: | ||
| key = f'/{meta.kind.value}/{meta.name}' | ||
| except Exception: | ||
| # Defensive: skip unexpected plugin metadata objects. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request is a major refactoring of the plugin system to a new asynchronous "V2" API. The changes are extensive, touching core registry logic, all existing plugins, and tests. The new plugin interface is cleaner and async-first. Most of the changes correctly adapt the codebase to this new model, including significant improvements to streaming logic in the reflection server and bug fixes in request handling. However, the refactoring appears to be incomplete for the VertexAIVectorSearch plugin, as its new init method is empty, effectively disabling its functionality. This is a critical issue that needs to be addressed.
| async def init(self) -> list: | ||
| """Initialize plugin with the retriever specified. | ||
| Register actions with the registry making them available for use in the Genkit framework. | ||
| Returns: | ||
| Empty list (using lazy loading via resolve). | ||
| """ | ||
| return [] | ||
|
|
There was a problem hiding this comment.
The refactoring of the VertexAIVectorSearch plugin to the new V2 API appears to be incomplete. The previous initialize method, which was responsible for registering the retriever action, has been replaced with an init method that currently returns an empty list. This means the plugin no longer registers its retriever and is effectively non-functional.
To fix this, the init method should create and return the retriever Action. However, this requires access to the Genkit instance (for embedding), which is not currently passed to the plugin's __init__. A similar plugin, FirestoreVectorStore, was updated to accept the ai instance in its constructor. A similar change is likely needed here.
Assuming the __init__ is updated to accept and store an ai instance as self.ai, the init method should be implemented like this:
async def init(self) -> list:
"""Initialize plugin with the retriever specified.
Creates and returns the retriever action.
"""
from genkit.core.action import Action
from genkit.core.action.types import ActionKind
from genkit.plugins.vertex_ai.vector_search.retriever import RetrieverOptionsSchema
from genkit.core.schema import to_json_schema
retriever = self.retriever_cls(
ai=self.ai,
name=self.name,
match_service_client_generator=self._match_service_client_generator,
embedder=self.embedder,
embedder_options=self.embedder_options,
**self.retriever_extra_args,
)
return [
Action(
kind=ActionKind.RETRIEVER,
name=vertexai_name(self.name),
fn=retriever.retrieve,
metadata={
'retriever': {
'customOptions': to_json_schema(RetrieverOptionsSchema),
}
},
)
]
|
|
||
| def send_chunk(chunk): | ||
| chunk_json = dump_json(chunk) | ||
| logger.info(f'📤 Sending chunk: {chunk_json[:200]}...') |
| ) | ||
| self.wfile.write(bytes('\n', encoding)) | ||
| self.wfile.flush() # Flush immediately for streaming | ||
| logger.info('✅ Chunk flushed') |
| def send_chunk(chunk): | ||
| """Callback that puts chunks on the queue.""" | ||
| out = dump_json(chunk) | ||
| chunk_queue.put_nowait(f'{out}\n') |
There was a problem hiding this comment.
not sure waht put_nowait does
| self._entries[action.kind] = {} | ||
| self._entries[action.kind][name] = action | ||
|
|
||
| async def resolve_action(self, kind: ActionKind, name: str) -> Action | None: |
There was a problem hiding this comment.
does this method skip registered actions? seems like it
| True if the query flag is enabled, False otherwise. | ||
| """ | ||
| return query_params.get(flag, ['false'])[0] == 'true' | ||
| return query_params.get(flag, 'false') == 'true' |
There was a problem hiding this comment.
think this is right but double check
| ) | ||
|
|
||
| override_messages = [ | ||
| [ |
There was a problem hiding this comment.
?? is this dead code
|
|
||
| got = registry.list_actions({}, allowed_kind) | ||
| assert got == expected | ||
| # NOTE: This test is disabled because the old synchronous list_actions method |
| resolver_calls.append([kind, name]) | ||
| resolver_calls.append([action_type, name]) | ||
|
|
||
| from genkit.core.action import Action |
There was a problem hiding this comment.
import at top level
| """ | ||
| model = AnthropicModel(model_name=model_name, client=self._anthropic_client) | ||
| model_info = get_model_info(model_name) | ||
| from genkit.core.action import Action |
| Returns: | ||
| List of ActionMetadata for all supported models. | ||
| """ | ||
| from genkit.blocks.model import model_action_metadata |
| The fully qualified OpenAi-Compat action name. | ||
| """ | ||
| return f'openai/{name}' | ||
| return f'openai-compat/{name}' |
There was a problem hiding this comment.
keep ti as openai
| Action object for the model. | ||
| """ | ||
| from genkit.core.action import Action | ||
| from genkit.core.schema import to_json_schema |
| model=clean_name, | ||
| api_key=self.api_key, | ||
| registry=ai, | ||
| registry=None, # We don't need registry anymore |
There was a problem hiding this comment.
see to remove registry parameter from this
| # Create retriever action | ||
| retriever = DevLocalVectorStoreRetriever( | ||
| ai=ai, | ||
| ai=None, |
There was a problem hiding this comment.
why pass ai=None
| embedder_options: dict[str, Any] | None = None, | ||
| distance_measure: DistanceMeasure = DistanceMeasure.COSINE, | ||
| metadata_fields: list[str] | MetadataTransformFn | None = None, | ||
| ai: Genkit | None = None, |
There was a problem hiding this comment.
why ai gets passedi n now?
| location=self.location, | ||
| project_id=self.project_id, | ||
| registry=ai, | ||
| registry=None, |
There was a problem hiding this comment.
why registry is an arg here
| response = await ai.generate( | ||
| prompt=f'Generate and run code for the task: {task}', | ||
| config=GeminiConfigSchema(temperature=1, code_execution=True), | ||
| config={'temperature': 1, 'code_execution': True}, |
There was a problem hiding this comment.
why did we get rid fo this again
1 participant
No description provided.