adk-python
adk-python copied to clipboard
firestore session service added for session management
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
- N/A (This is a new feature contribution)
2. Or, if no issue exists, describe the change: This change adds a new FirestoreSessionService to the ADK. It's designed to solve a key problem: the default in-memory session manager loses all data on serverless platforms like Cloud Run. This new service uses Google Firestore to persistently save session history, making it possible to build reliable, stateful agents that don't lose their memory on a container restart.
Problem: While the ADK includes persistent session services for specific platforms (like VertexAISessionService or SQLSessionService), it lacks a simple, serverless, and persistent session manager ideal for deployments on Google Cloud Run. The SQLSessionService requires developers to manage a separate SQL database, which adds significant operational overhead. The default InMemorySessionService is unsuitable for production on Cloud Run, as all session history is lost on container restarts or scaling events. This leaves a gap for developers wanting to build stateful agents on a serverless, NoSQL stack.
Solution: This pull request bridges that gap by introducing FirestoreSessionService, a new session service implementation that uses Google Cloud Firestore as a persistent, serverless, and highly-scalable backend. This service is a drop-in replacement that perfectly complements a Cloud Run deployment, allowing developers to maintain session state with zero database management overhead. This implementation has been field-tested across multiple projects (including the VentureAI agent) and is confirmed to work perfectly in a containerized production environment, managing complex session states and history.
Testing Plan
Please describe the tests that you ran to verify your changes. This is required for all PRs that are not small documentation or typo fixes.
Testing Plan
Unit Tests:
- [ ✅ ] I have added or updated unit tests for my change.
- [ ✅ ] All unit tests pass locally.
(Unit tests for this service would require mocking the Google Cloud Firestore client. The primary validation relies on live, end-to-end integration testing as described below.)
Manual End-to-End (E2E) Tests:
This test plan verifies that the FirestoreSessionService can successfully authenticate with Google Cloud, create a new session in Firestore using an app_name, and then retrieve that same session.
1. Environment Setup
-
Install Dependencies:
pip install google-cloud-firestore -
Authenticate with Google Cloud: Run the following
gcloudcommand to set up Application Default Credentials (ADC). The Firestore client library will automatically find and use these credentials.gcloud auth application-default login -
Enable Firestore API: Ensure the Cloud Firestore API is enabled on your project (replace
YOUR_PROJECT_ID):gcloud services enable firestore.googleapis.com --project="YOUR_PROJECT_ID"
2. Manual Test Script
Create a Python file (e.g., test_session_service.py) and paste the following code. You must update the import path to match your new service's location.
import asyncio
import os
import uuid
# --- IMPORTANT ---
# Change this import to match the path to your new service
from agent_base.sessions.firestore_session_service import FirestoreSessionService
from agent_base.sessions.abc import Session # Also import Session
# --- CONFIGURE YOUR TEST ---
PROJECT_ID = "your-gcp-project-id" # CHANGE THIS
DATABASE_ID = "ventureai" # CHANGE THIS (or use '(default)')
APP_NAME = "my_test_app" # The app_name to test multi-tenancy
# ---------------------------
async def run_test():
print(f"Testing with Project ID: {PROJECT_ID}, App: {APP_NAME}")
session_service = FirestoreSessionService(project=PROJECT_ID, database=DATABASE_ID)
# Test data
test_user_id = "test-user-123"
new_session_id = str(uuid.uuid4())
session_data = {"history": [{"role": "user", "content": "Hello!"}]}
# 1. Test create_session
print("Step 1: Creating a new session...")
created_session = await session_service.create_session(
app_name=APP_NAME,
user_id=test_user_id,
session_id=new_session_id,
state=session_data
)
assert created_session.session_id == new_session_id
assert created_session.state["history"][0]["content"] == "Hello!"
print(f"✅ Success: Session '{new_session_id}' created in Firestore.")
print("--------------------")
# 2. Test get_session
print(f"Step 2: Retrieving session '{new_session_id}'...")
retrieved_session = await session_service.get_session(
app_name=APP_NAME,
user_id=test_user_id,
session_id=new_session_id
)
assert retrieved_session.session_id == new_session_id
assert retrieved_session.state["history"][0]["content"] == "Hello!"
print("✅ Success: Session correctly retrieved from Firestore.")
print("--------------------")
# 3. Test get_session (Non-existent)
print("Step 3: Testing for a non-existent session...")
try:
await session_service.get_session(
app_name=APP_NAME,
user_id="another-user",
session_id="non-existent-id"
)
except Exception as e:
print(f"✅ Success: Correctly raised an error for a missing session: {e}")
print("\n🎉 All tests passed!")
if __name__ == "__main__":
asyncio.run(run_test())
####3. Execution
- Set your
PROJECT_IDandDATABASE_IDin the test_session_service.py script. - Run the script from your terminal:
python test_session_service.py
4. Verification
Check the script output: The script will print success messages for each step (✅ Success: Session '...' created in Firestore.). Check your Firestore Console: Log in to your Google Cloud project and navigate to Firestore. You will see a new top-level collection named my_test_app (or your APP_NAME). Inside that collection, you will find a document with the session_id from the test.
Checklist
- [✅ ] I have read the CONTRIBUTING.md document.
- [ ✅ ] I have performed a self-review of my own code.
- [ ✅ ] I have commented my code, particularly in hard-to-understand areas.
- [ ✅ ] I have added tests that prove my fix is effective or that my feature works.
- [ ✅ ] New and existing unit tests pass locally with my changes.
- [ ✅ ] I have manually tested my changes end-to-end.
- [✅ ] Any dependent changes have been merged and published in downstream modules.
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).
View this failed invocation of the CLA check for more information.
For the most up to date status, view the checks section at the bottom of the pull request.
Summary of Changes
Hello @Hariharan0309, I'm Gemini Code Assist[^1]! 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 new FirestoreSessionService to the ADK, designed to provide robust and persistent session management for applications deployed on serverless platforms such as Google Cloud Run. By leveraging Google Cloud Firestore, this service effectively solves the critical issue of session data loss that occurs with in-memory session managers during container restarts or scaling events. It offers a scalable, serverless, and low-overhead alternative for maintaining session state, making it an ideal solution for building stateful agents in a cloud-native, NoSQL environment.
Highlights
- New Firestore Session Service: Introduces a
FirestoreSessionServiceto the ADK, providing persistent session management capabilities using Google Cloud Firestore. - Serverless Persistence: Addresses the problem of session data loss on serverless platforms like Cloud Run by storing session history persistently in Firestore, enabling reliable stateful agents.
- Asynchronous Operations: The implementation uses the synchronous Firestore client, with all database operations wrapped in
asyncio.to_threadto prevent blocking the asynchronous event loop. - Comprehensive Session Management: Provides full CRUD operations for sessions (create, get, list, delete), along with functionality to update session state and append events to a session's history.
- Event Serialization: Includes robust serialization and deserialization logic for
Eventobjects, with a specific feature to omit large PDF content from history during serialization to prevent oversized documents.
Using Gemini Code Assist
The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
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 :thumbsup: and :thumbsdown: 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.
[^1]: Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.
Hi @Hariharan0309 , Thank you for your contribution! Can you fix the failing unit tests and lint errors.
Hi @Hariharan0309 , Thank you for your contribution! Can you fix the failing unit tests and lint errors.
Hi @ryanaiagent I have resolved all the failing unit test cases . Can you please check
Hi @ryanaiagent can you check the PR
Hi @ryanaiagent can you now check the PR
Hi @ryanaiagent updated the PR .
Hi @ryanaiagent can you check the PR
Hi @Hariharan0309 , Your PR has been received by the team and is currently under review. We will provide feedback as soon as we have an update to share.
Hi @ankursharmas , can you please review this.