You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
docs: Add PostgreSQL session storage sample and documentation
Merge https://github.com/google/adk-python/pull/3926 ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** - Related: #3916 **2. Or, if no issue exists, describe the change:** **Problem:** While `DatabaseSessionService` already supports PostgreSQL through SQLAlchemy, there is no documentation or sample code showing users how to configure and use it. **Solution:** Add a comprehensive sample under `contributing/samples/postgres_session_service/` that demonstrates: - How to configure `DatabaseSessionService` with PostgreSQL - The auto-generated database schema (sessions, events, app_states, user_states tables) - Connection URL format and configuration options - A working sample agent with session persistence ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. This is a documentation-only change (new sample), so no new unit tests are required. Existing tests continue to pass. **Manual End-to-End (E2E) Tests:** Tested locally with the following steps: 1. Started PostgreSQL using `docker compose up -d` 2. Set environment variables: ```bash export POSTGRES_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/adk_sessions export GOOGLE_CLOUD_PROJECT=$(gcloud config get-value project) export GOOGLE_CLOUD_LOCATION=us-central1 export GOOGLE_GENAI_USE_VERTEXAI=true ``` 3. Ran `pip install google-adk asyncpg greenlet` to install the required packages 4. Ran `python main.py` - session created successfully 5. Ran `python main.py` again - previous session resumed with event history 6. Verified tables and rows created in PostgreSQL (sessions, events, app_states, user_states) ### Checklist - [x] I have read the https://github.com/google/adk-python/blob/main/CONTRIBUTING.md document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [x] Any dependent changes have been merged and published in downstream modules. ### Additional context This PR adds documentation and a working sample for an already-supported feature. The DatabaseSessionService class already handles PostgreSQL through its DynamicJSON type decorator which uses JSONB for PostgreSQL. Files added: - contributing/samples/postgres_session_service/README.md - Comprehensive guide - contributing/samples/postgres_session_service/agent.py - Sample agent - contributing/samples/postgres_session_service/main.py - Usage example - contributing/samples/postgres_session_service/compose.yml - Local PostgreSQL setup - contributing/samples/postgres_session_service/\_\_init\_\_.py - Package init Co-authored-by: Liang Wu <wuliang@google.com> COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3926 from hiroakis:feat-support-pg-for-conversation a5279d4fb2a63699f384c670928d77d882c25a25 PiperOrigin-RevId: 868816317
This commit is contained in:
committed by
Copybara-Service
parent
61c329f8ce
commit
657acfadbb
@@ -0,0 +1,198 @@
|
||||
# Using PostgreSQL with DatabaseSessionService
|
||||
|
||||
This sample demonstrates how to configure `DatabaseSessionService` to use PostgreSQL for persisting sessions, events, and state.
|
||||
|
||||
## Overview
|
||||
|
||||
ADK's `DatabaseSessionService` supports multiple database backends through SQLAlchemy. This guide shows how to:
|
||||
|
||||
- Set up PostgreSQL as the session storage backend
|
||||
- Configure async connections with `asyncpg`
|
||||
- Understand the auto-generated schema
|
||||
- Run the sample agent with persistent sessions
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- **PostgreSQL Database**: A running PostgreSQL instance (local or cloud)
|
||||
- **asyncpg**: Async PostgreSQL driver for Python
|
||||
|
||||
## Installation
|
||||
|
||||
Install the required Python packages:
|
||||
|
||||
```bash
|
||||
pip install google-adk asyncpg greenlet
|
||||
```
|
||||
|
||||
## Database Schema
|
||||
|
||||
`DatabaseSessionService` automatically creates the following tables on first use:
|
||||
|
||||
### sessions
|
||||
|
||||
| Column | Type | Description |
|
||||
| ----------- | ------------ | ------------------------------ |
|
||||
| app_name | VARCHAR(128) | Application identifier (PK) |
|
||||
| user_id | VARCHAR(128) | User identifier (PK) |
|
||||
| id | VARCHAR(128) | Session UUID (PK) |
|
||||
| state | JSONB | Session state as JSON |
|
||||
| create_time | TIMESTAMP | Creation timestamp |
|
||||
| update_time | TIMESTAMP | Last update timestamp |
|
||||
|
||||
### events
|
||||
|
||||
| Column | Type | Description |
|
||||
| ------------------ | ------------ | ------------------------------ |
|
||||
| id | VARCHAR(256) | Event UUID (PK) |
|
||||
| app_name | VARCHAR(128) | Application identifier (PK) |
|
||||
| user_id | VARCHAR(128) | User identifier (PK) |
|
||||
| session_id | VARCHAR(128) | Session reference (PK, FK) |
|
||||
| invocation_id | VARCHAR(256) | Invocation identifier |
|
||||
| timestamp | TIMESTAMP | Event timestamp |
|
||||
| event_data | JSONB | Event content as JSON |
|
||||
|
||||
### app_states
|
||||
|
||||
| Column | Type | Description |
|
||||
| ----------- | ------------ | ------------------------------ |
|
||||
| app_name | VARCHAR(128) | Application identifier (PK) |
|
||||
| state | JSONB | Application-level state |
|
||||
| update_time | TIMESTAMP | Last update timestamp |
|
||||
|
||||
### user_states
|
||||
|
||||
| Column | Type | Description |
|
||||
| ----------- | ------------ | ------------------------------ |
|
||||
| app_name | VARCHAR(128) | Application identifier (PK) |
|
||||
| user_id | VARCHAR(128) | User identifier (PK) |
|
||||
| state | JSONB | User-level state |
|
||||
| update_time | TIMESTAMP | Last update timestamp |
|
||||
|
||||
### adk_internal_metadata
|
||||
|
||||
| Column | Type | Description |
|
||||
| ----------- | ------------ | ------------------------------ |
|
||||
| key | VARCHAR(128) | Metadata key |
|
||||
| value | VARCHAR(256) | Metadata value |
|
||||
|
||||
|
||||
## Configuration
|
||||
|
||||
### Connection URL Format
|
||||
|
||||
```python
|
||||
postgresql+asyncpg://username:password@host:port/database
|
||||
```
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```python
|
||||
from google.adk.sessions.database_session_service import DatabaseSessionService
|
||||
from google.adk.runners import Runner
|
||||
|
||||
# Initialize with PostgreSQL URL
|
||||
session_service = DatabaseSessionService(
|
||||
"postgresql+asyncpg://postgres:postgres@localhost:5432/adk_sessions"
|
||||
)
|
||||
|
||||
# Use with Runner
|
||||
runner = Runner(
|
||||
app_name="my_app",
|
||||
agent=my_agent,
|
||||
session_service=session_service,
|
||||
)
|
||||
```
|
||||
|
||||
### Advanced Configuration
|
||||
|
||||
Pass additional SQLAlchemy engine options:
|
||||
|
||||
```python
|
||||
session_service = DatabaseSessionService(
|
||||
"postgresql+asyncpg://postgres:postgres@localhost:5432/adk_sessions",
|
||||
pool_size=10,
|
||||
max_overflow=20,
|
||||
pool_timeout=30,
|
||||
pool_recycle=1800,
|
||||
)
|
||||
```
|
||||
|
||||
## Running the Sample
|
||||
|
||||
### 1. Start PostgreSQL
|
||||
|
||||
Using Docker:
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
Or use an existing PostgreSQL instance.
|
||||
|
||||
### 2. Configure Connection
|
||||
|
||||
Create a `.env` file:
|
||||
|
||||
```bash
|
||||
POSTGRES_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/adk_sessions
|
||||
GOOGLE_CLOUD_PROJECT=<your-gcp-project-id>
|
||||
GOOGLE_CLOUD_LOCATION=us-central1
|
||||
GOOGLE_GENAI_USE_VERTEXAI=true
|
||||
```
|
||||
|
||||
Or run export command.
|
||||
|
||||
```bash
|
||||
export POSTGRES_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/adk_sessions
|
||||
export GOOGLE_CLOUD_PROJECT=$(gcloud config get-value project)
|
||||
export GOOGLE_CLOUD_LOCATION=us-central1
|
||||
export GOOGLE_GENAI_USE_VERTEXAI=true
|
||||
```
|
||||
|
||||
### 3. Run the Agent
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
Or use the ADK:
|
||||
|
||||
```bash
|
||||
adk run .
|
||||
```
|
||||
|
||||
## Session Persistence
|
||||
|
||||
Sessions and events are persisted across application restarts:
|
||||
|
||||
```python
|
||||
# First run - creates a new session
|
||||
session = await session_service.create_session(
|
||||
app_name="my_app",
|
||||
user_id="user1",
|
||||
session_id="persistent-session-123",
|
||||
)
|
||||
|
||||
# Later run - retrieves the existing session
|
||||
session = await session_service.get_session(
|
||||
app_name="my_app",
|
||||
user_id="user1",
|
||||
session_id="persistent-session-123",
|
||||
)
|
||||
```
|
||||
|
||||
## State Management
|
||||
|
||||
PostgreSQL's JSONB type provides efficient storage for state data:
|
||||
|
||||
- **Session state**: Stored in `sessions.state`
|
||||
- **User state**: Stored in `user_states.state`
|
||||
- **App state**: Stored in `app_states.state`
|
||||
|
||||
## Production Considerations
|
||||
|
||||
1. **Connection Pooling**: Use `pool_size` and `max_overflow` for high-traffic applications
|
||||
2. **SSL/TLS**: Always use encrypted connections in production
|
||||
3. **Backups**: Implement regular backup strategies for session data
|
||||
4. **Indexing**: The default schema includes primary key indexes; add additional indexes based on query patterns
|
||||
5. **Monitoring**: Monitor connection pool usage and query performance
|
||||
@@ -0,0 +1,16 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
from . import agent
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Sample agent demonstrating PostgreSQL session persistence."""
|
||||
|
||||
from datetime import datetime
|
||||
from datetime import timezone
|
||||
|
||||
from google.adk.agents.llm_agent import Agent
|
||||
|
||||
|
||||
def get_current_time() -> str:
|
||||
"""Get the current time.
|
||||
|
||||
Returns:
|
||||
A string with the current time in ISO 8601 format.
|
||||
"""
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
root_agent = Agent(
|
||||
model="gemini-2.0-flash",
|
||||
name="postgres_session_agent",
|
||||
description="A sample agent demonstrating PostgreSQL session persistence.",
|
||||
instruction="""
|
||||
You are a helpful assistant that demonstrates session persistence.
|
||||
You can remember previous conversations within the same session.
|
||||
Use the get_current_time tool when asked about the time.
|
||||
When the user asks what you remember, summarize the previous conversation.
|
||||
""",
|
||||
tools=[get_current_time],
|
||||
)
|
||||
@@ -0,0 +1,24 @@
|
||||
# Docker Compose configuration for the postgres_session_service sample.
|
||||
#
|
||||
# This file defines a PostgreSQL service used to demonstrate ADK's
|
||||
# DatabaseSessionService with a persistent backend. It sets up a
|
||||
# postgres:16-alpine container with:
|
||||
# - Default credentials (user: postgres, password: postgres)
|
||||
# - A pre-created database named 'adk_sessions'
|
||||
# - Port 5432 exposed for local access
|
||||
# - A named volume 'postgres_data' for data persistence
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16-alpine
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: adk_sessions
|
||||
ports:
|
||||
- "5432:5432"
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Example demonstrating PostgreSQL session persistence with DatabaseSessionService."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import agent
|
||||
from dotenv import load_dotenv
|
||||
from google.adk.runners import Runner
|
||||
from google.adk.sessions.database_session_service import DatabaseSessionService
|
||||
from google.adk.sessions.session import Session
|
||||
from google.genai import types
|
||||
|
||||
load_dotenv(override=True)
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main function demonstrating PostgreSQL session persistence."""
|
||||
postgres_url = os.environ.get("POSTGRES_URL")
|
||||
if not postgres_url:
|
||||
raise ValueError(
|
||||
"POSTGRES_URL environment variable not set. "
|
||||
"Please create a .env file with"
|
||||
" POSTGRES_URL=postgresql+asyncpg://user:password@localhost:5432/adk_sessions"
|
||||
)
|
||||
|
||||
app_name = "postgres_session_demo"
|
||||
user_id = "demo_user"
|
||||
session_id = "persistent-session"
|
||||
|
||||
# Initialize PostgreSQL-backed session service
|
||||
session_service = DatabaseSessionService(postgres_url)
|
||||
|
||||
runner = Runner(
|
||||
app_name=app_name,
|
||||
agent=agent.root_agent,
|
||||
session_service=session_service,
|
||||
)
|
||||
|
||||
# Try to get existing session or create new one
|
||||
session = await session_service.get_session(
|
||||
app_name=app_name,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
)
|
||||
|
||||
if session:
|
||||
print(f"Resuming existing session: {session.id}")
|
||||
print(f"Previous events count: {len(session.events)}")
|
||||
else:
|
||||
session = await session_service.create_session(
|
||||
app_name=app_name,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
)
|
||||
print(f"Created new session: {session.id}")
|
||||
|
||||
async def run_prompt(session: Session, new_message: str):
|
||||
"""Send a prompt to the agent and print the response."""
|
||||
content = types.Content(
|
||||
role="user", parts=[types.Part.from_text(text=new_message)]
|
||||
)
|
||||
print(f"User: {new_message}")
|
||||
async for event in runner.run_async(
|
||||
user_id=user_id,
|
||||
session_id=session.id,
|
||||
new_message=content,
|
||||
):
|
||||
if event.content and event.content.parts and event.content.parts[0].text:
|
||||
print(f"{event.author}: {event.content.parts[0].text}")
|
||||
|
||||
print("------------------------------------")
|
||||
await run_prompt(session, "What time is it? Please remember this.")
|
||||
print("------------------------------------")
|
||||
await run_prompt(session, "What did I just ask you?")
|
||||
print("------------------------------------")
|
||||
|
||||
print("\nSession persisted to PostgreSQL. Run again to see event history.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user