feat: Implement a vertex based task store for the 1.0 branch#791
feat: Implement a vertex based task store for the 1.0 branch#791gaborfeher wants to merge 2 commits intoa2aproject:1.0-devfrom
Conversation
Summary of ChangesHello, 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 integrates the system with Google Cloud's Vertex AI Agent Engine by introducing a new task storage implementation. It provides the necessary components to save, retrieve, and update tasks directly within Vertex AI, enhancing the platform's capabilities for managing AI-driven workflows. The changes include core logic for data conversion and a robust testing suite to ensure reliability. Highlights
Changelog
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
|
![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 introduces a new TaskStore implementation using Google Cloud Vertex AI, along with the necessary data converters and tests. The implementation is well-structured and includes comprehensive tests for both fake and real backends. My feedback includes suggestions for improving code readability and maintainability, such as refactoring a complex method to use list comprehensions, using f-strings for path construction, and improving test setup practices. I've also noted a few minor points regarding docstrings and exception handling, and added a style guide reference for import placement.
Note: Security Review did not run due to the size of the PR.
| for artifact in previous_task_artifacts.values(): | ||
| if artifact.artifact_id not in task_artifacts: | ||
| if not task_artifact_change.deleted_artifact_ids: | ||
| task_artifact_change.deleted_artifact_ids = [] | ||
| task_artifact_change.deleted_artifact_ids.append( | ||
| artifact.artifact_id | ||
| ) | ||
| for artifact in task_artifacts.values(): | ||
| if artifact.artifact_id not in previous_task_artifacts: | ||
| if not task_artifact_change.added_artifacts: | ||
| task_artifact_change.added_artifacts = [] | ||
| task_artifact_change.added_artifacts.append( | ||
| vertex_task_converter.to_stored_artifact(artifact) | ||
| ) | ||
| elif artifact != previous_task_artifacts[artifact.artifact_id]: | ||
| if not task_artifact_change.updated_artifacts: | ||
| task_artifact_change.updated_artifacts = [] | ||
| task_artifact_change.updated_artifacts.append( | ||
| vertex_task_converter.to_stored_artifact(artifact) | ||
| ) |
There was a problem hiding this comment.
The logic for detecting artifact changes can be made more efficient and readable by using list comprehensions. This avoids repeated checks and list initializations inside loops.
deleted_ids = [
artifact.artifact_id
for artifact in previous_task_artifacts.values()
if artifact.artifact_id not in task_artifacts
]
if deleted_ids:
task_artifact_change.deleted_artifact_ids = deleted_ids
added_artifacts = [
vertex_task_converter.to_stored_artifact(artifact)
for artifact in task_artifacts.values()
if artifact.artifact_id not in previous_task_artifacts
]
if added_artifacts:
task_artifact_change.added_artifacts = added_artifacts
updated_artifacts = [
vertex_task_converter.to_stored_artifact(artifact)
for artifact_id, artifact in task_artifacts.items()
if artifact_id in previous_task_artifacts
and artifact != previous_task_artifacts[artifact_id]
]
if updated_artifacts:
task_artifact_change.updated_artifacts = updated_artifacts| if not events: | ||
| return | ||
| await self._client.aio.agent_engines.a2a_tasks.events.append( | ||
| name=self._agent_engine_resource_id + '/a2aTasks/' + task.id, |
There was a problem hiding this comment.
Using an f-string for constructing the resource name is more readable and idiomatic than string concatenation. This also applies on line 190.
| name=self._agent_engine_resource_id + '/a2aTasks/' + task.id, | |
| name=f'{self._agent_engine_resource_id}/a2aTasks/{task.id}', |
| params: ListTasksRequest, | ||
| context: ServerCallContext | None = None, | ||
| ) -> ListTasksResponse: | ||
| """Retrieves a list of tasks from the store.""" |
There was a problem hiding this comment.
The list method raises NotImplementedError. It would be helpful to update the docstring to explain why, similar to how it's done for the delete method. This clarifies whether the feature is unsupported by the backend or planned for future implementation.
| """Retrieves a list of tasks from the store.""" | |
| """The backend doesn't support listing tasks, so this is not implemented.""" |
| try: | ||
| del task.output.artifacts[:] | ||
| task.output.artifacts.extend(current_artifacts) | ||
| except Exception: |
There was a problem hiding this comment.
Catching a broad Exception can hide unexpected errors during testing. It's better to catch more specific exceptions. If the goal is to handle cases where task.output.artifacts might not be a list that supports item deletion, consider catching AttributeError or TypeError instead.
| 'VERTEX_API_VERSION', | ||
| ] | ||
| ) | ||
| import sys |
There was a problem hiding this comment.
According to PEP 8, imports should be at the top of the file. Please move import sys to the top with the other imports.
References
- PEP 8 recommends that imports are always placed at the top of the file, just after any module comments and docstrings, and before module globals and constants. (link)
| sys.path.append(os.path.dirname(__file__)) | ||
| from fake_vertex_client import FakeVertexClient |
There was a problem hiding this comment.
Modifying sys.path at runtime can make tests less predictable and harder to maintain. Consider making the tests/contrib/tasks directory a Python package by adding an empty __init__.py file. This would allow you to use a relative import like from .fake_vertex_client import FakeVertexClient without manipulating the path.
gaborfeher
force-pushed
the
1.0-dev
branch
3 times, most recently
from
1d6e5c0 to
2f2ffd5
Compare
gaborfeher
force-pushed
the
1.0-dev
branch
13 times, most recently
from
5baa501 to
7c2456b
Compare
1 participant
For #751