Add DisplayName and Labels support to Sessions #4165
Conversation
Summary of ChangesHello @ashubham, 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 significantly improves session management by introducing 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
|
adk-bot
added
the
services
![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 adds support for display_name and labels to sessions, which is a valuable feature for identifying and filtering them. The changes are consistently applied across all session service implementations (in-memory, database, sqlite, and vertex_ai). The new fields are also added to the data models and database schemas, with backward compatibility in mind for sqlite_session_service through schema alteration. The addition of unit tests is also great.
My main feedback is regarding performance in the database-backed session services. Both DatabaseSessionService and SqliteSessionService perform label-based filtering in Python after fetching all data from the database. This can be inefficient. I've left comments with suggestions to move this filtering logic to the database query itself.
I've also included a couple of suggestions for the in_memory_session_service to improve code conciseness and reduce duplication.
Overall, this is a solid contribution. Addressing the performance concerns in the database services would make it even better.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for display_name and labels in sessions, which is a valuable feature for organizing and filtering. The changes are comprehensive, consistently applied across all session service implementations (in-memory, database, sqlite, and vertex_ai). The data models, database schemas, and service methods have been updated correctly. Label-based filtering is implemented with backend-specific optimizations, and a non-disruptive schema migration for SQLite is included to maintain backward compatibility. The addition of unit tests for the new functionality is also great. Overall, this is a well-executed feature addition. I have one suggestion to improve the efficiency of the schema check in the SqliteSessionService.
| async def _ensure_new_columns(self, db: aiosqlite.Connection) -> None: | ||
| """Ensures display_name and labels columns exist in the sessions table.""" | ||
| async with db.execute("PRAGMA table_info(sessions)") as cursor: | ||
| columns = [row[1] async for row in cursor] | ||
|
|
||
| if "display_name" not in columns: | ||
| await db.execute("ALTER TABLE sessions ADD COLUMN display_name TEXT") | ||
| if "labels" not in columns: | ||
| await db.execute( | ||
| "ALTER TABLE sessions ADD COLUMN labels TEXT NOT NULL DEFAULT '{}'" | ||
| ) |
There was a problem hiding this comment.
This method is called on every database connection, which is inefficient as it runs a PRAGMA query each time. To optimize this, the check should run only once per service instance.
You can achieve this with an instance flag and a lock for thread safety. First, add self._columns_ensured = False and self._columns_ensured_lock = asyncio.Lock() to the __init__ method. Then, modify this method to use them as shown in the suggestion.
At the call site in _get_db_connection, you would then wrap the call:
# Ensure new columns exist for existing databases
if not self._columns_ensured:
await self._ensure_new_columns(db)| async def _ensure_new_columns(self, db: aiosqlite.Connection) -> None: | |
| """Ensures display_name and labels columns exist in the sessions table.""" | |
| async with db.execute("PRAGMA table_info(sessions)") as cursor: | |
| columns = [row[1] async for row in cursor] | |
| if "display_name" not in columns: | |
| await db.execute("ALTER TABLE sessions ADD COLUMN display_name TEXT") | |
| if "labels" not in columns: | |
| await db.execute( | |
| "ALTER TABLE sessions ADD COLUMN labels TEXT NOT NULL DEFAULT '{}'" | |
| ) | |
| async def _ensure_new_columns(self, db: aiosqlite.Connection) -> None: | |
| """Ensures display_name and labels columns exist in the sessions table.""" | |
| async with self._columns_ensured_lock: | |
| if self._columns_ensured: | |
| return | |
| async with db.execute("PRAGMA table_info(sessions)") as cursor: | |
| columns = [row[1] async for row in cursor] | |
| if "display_name" not in columns: | |
| await db.execute("ALTER TABLE sessions ADD COLUMN display_name TEXT") | |
| if "labels" not in columns: | |
| await db.execute( | |
| "ALTER TABLE sessions ADD COLUMN labels TEXT NOT NULL DEFAULT '{}'" | |
| ) | |
| self._columns_ensured = True |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds support for display_name and labels to sessions, which is a great feature for organizing and filtering. The changes are implemented consistently across different session service backends (InMemory, Database, Sqlite, VertexAi). The addition of tests for the new functionality is also good. I have one main suggestion regarding the SqliteSessionService implementation to improve its robustness.
| if labels_filter: | ||
| for key, value in labels_filter.items(): | ||
| query_parts.append("AND json_extract(labels, ?)=?") | ||
| params.extend([f"$.{key}", value]) |
There was a problem hiding this comment.
The current construction of the JSONPath f"$.{key}" is not safe if the label key contains special characters like . which are interpreted by json_extract. To handle all valid string keys correctly, the key should be properly quoted within the JSONPath expression. For example, a key "a.b" would be interpreted as nested access, not as a literal key. Using json.dumps(key) will correctly quote and escape the key, ensuring it's treated as a single literal key in the JSONPath.
| params.extend([f"$.{key}", value]) | |
| params.extend([f"$.{json.dumps(key)}", value]) |
ryanaiagent
added
the
request clarification
3 participants
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
Problem:
ADK does not support
display_nameandlabelsin theSessionclass. Both of these are required for a title for a session and labeling the sessions for filtering when listed.Solution:
Support
display_nameandlabelsnatively. A further improvement could be to autogenerate titles based onmodelwhen usingLLMAgent.Testing Plan
Ran all the Unit/Integration Tests. Also added basic tests. The changes are done in a backward compatible way by making the the two fields optional. This ensures none of the existing tests needed updating.
Unit Tests:
Please include a summary of passed
pytestresults.Manual End-to-End (E2E) Tests:
display_nameand/orlabelsusing the API.sessionobject available in theContext.Checklist
Additional context
The new fields align with what is present in the
Sessionresource on the VertexAI Sessions service defined here