Commit Graph
47 Commits
Author SHA1 Message Date
GitMarco27andCopybara-Service 74959414d8 feat: full async implementation of DatabaseSessionService
Merge https://github.com/google/adk-python/pull/2889

# Implement Full async DatabaseSessionService

**Target Issue:** #1005

## Overview

This PR introduces an asynchronous implementation of the `DatabaseSessionService` with minimal breaking changes. The primary goal is to enable effective use of ADK in fully async environments and API endpoints while avoiding event loop blocking during database I/O operations.

## Changes

- Converted `DatabaseSessionService` to use async/await patterns throughout

## Testing Plan

The implementation has been tested following the project's contribution guidelines:

### Unit Tests
- All existing unit tests pass successfully
- Minor update to test requirements added to support `aiosqlite`

### Manual End-to-End Testing
- E2E tests performed using:
  - **LLM Provider:** LiteLLM
  - **Database:** PostgreSQL with `asyncpg` driver

 ```python
from google.adk.sessions.database_session_service import DatabaseSessionService

connection_string: str = (
    "postgresql+asyncpg://PG_USER:PG_PSWD@PG_HOST:5432/PG_DB"
)
session_service: DatabaseSessionService = DatabaseSessionService(
    db_url=connection_string
)

session = await session_service.create_session(
    app_name="test_app", session_id="test_session", user_id="test_user"
)
assert session is not None

sessions = await session_service.list_sessions(app_name="test_app", user_id="test_user")
assert len(sessions.sessions) > 0

session = await session_service.get_session(
    app_name="test_app", session_id="test_session", user_id="test_user"
)
assert session is not None

await session_service.delete_session(
    app_name="test_app", session_id="test_session", user_id="test_user"
)
assert (
    await session_service.get_session(
        app_name="test_app", session_id="test_session", user_id="test_user"
    )
    is None
)
```

The implementation have been also tested using the following configurations for llm provider and Runner:

```python
def get_azure_openai_model(deployment_id: str | None = None) -> LiteLlm:
    ...

    if not deployment_id:
        deployment_id = os.getenv("AZURE_OPENAI_DEPLOYMENT_ID")

    logger.info(f"Using Azure OpenAI deployment ID: {deployment_id}")

    return LiteLlm(
        model=f"azure/{os.getenv('AZURE_OPENAI_DEPLOYMENT_ID')}",
        stream=True,
    )

...

    @staticmethod
    def _get_runner(agent: Agent) -> Runner:
        storage=DatabaseSessionService(db_url=get_pg_connection_string())
        return Runner(
            agent=agent,
            app_name=APP_NAME,
            session_service=storage,
        )

...

async for event in self.runner.run_async(
    user_id=user_id,
    session_id=session_id,
    new_message=content,
    run_config=(
        RunConfig(
            streaming_mode=StreamingMode.SSE, response_modalities=["TEXT"]
        )
        if stream
        else RunConfig()
    ),
):
    last_event = event
    if stream:
        yield event

...

```

## Breaking Changes

- Database connection string format may need updates for async drivers

Co-authored-by: Shangjie Chen <deanchen@google.com>
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2889 from GitMarco27:feature/async_database_session_service e1b1b14934c1fb7975a6832cdd1549e94acab985
PiperOrigin-RevId: 830525148
2025-11-10 11:18:02 -08:00
Amy WuandCopybara-Service 7937e9ed40 fix: Use async with for API client in Vertex session management to ensure proper resource management
PiperOrigin-RevId: 830485639
2025-11-10 09:49:10 -08:00
Shangjie ChenandCopybara-Service e218254495 feat: Add SqliteSessionService and a migration script to migrate existing DB using DatabaseSessionService to SqliteSessionService
The new Sqlite version has fixed schema and use a single column to store Event data, this should avoid DB migration for future add to the Event object.

- This change introduces `SqliteSessionService`, an asynchronous session service using `aiosqlite` that stores event data as JSON within SQLite.
- A migration script, `migrate_from_sqlalchemy_sqlite.py`, is included to transition data from the older SQLAlchemy-based SQLite schema to this new format.
- The CLI service registry is updated to use SqliteSessionService for sqlite:// URIs.
- Throw error when user trying to access a legacy DB and advice the user to do the migration.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 829971174
2025-11-08 21:19:38 -08:00
hung12ctandCopybara-Service fc15c9a0c3 fix: Update DynamicPickleType to support MySQL dialect
Merge https://github.com/google/adk-python/pull/3282

The `process_bind_param` and `process_result_value` methods in the `DynamicPickleType` class have been modified to handle MySQL dialect in addition to Spanner. This change ensures that pickled values are correctly processed for both database types.

**Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.**

### Link to Issue or Description of Change

**1. Link to an existing issue (if applicable):**

- Closes: #3283

**2. Or, if no issue exists, describe the change:**

_If applicable, please follow the issue templates to provide as much detail as
possible._

**Problem:**
When using `DatabaseSessionService` with MySQL backend in google-adk v1.17.0, the application crashes with the following error: app.resources.runner:event_generator:260 - Error in event_generator: (builtins.TypeError) 'tuple' object cannot be interpreted as an integer
<img width="1237" height="129" alt="image" src="https://github.com/user-attachments/assets/0a5fc223-600a-4a92-8443-4d37fb1267f6" />

Root cause: The `DynamicPickleType` class in `database_session_service.py` configures MySQL dialect to use `LONGBLOB` for storing pickled data (line 117-118), but the `process_bind_param` and `process_result_value` methods only handle pickle serialization/deserialization for Spanner dialect, not MySQL. This causes MySQL to attempt storing raw Python objects instead of pickled bytes, leading to serialization errors and potential data corruption.

**Solution:**
Added MySQL to the pickle serialization logic in both `process_bind_param` and `process_result_value` methods, treating it the same way as Spanner dialect. This ensures that:
- Data is properly pickled to bytes before being stored in MySQL's LONGBLOB column
- Data is properly unpickled when retrieved from the database
- No breaking changes to existing functionality for other dialects (SQLite, PostgreSQL)

### 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._

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

**Summary of `pytest` results:**
<img width="929" height="306" alt="image" src="https://github.com/user-attachments/assets/3d548b96-ac49-4101-8405-a289a722293c" />

**Manual End-to-End (E2E) Tests:**

_Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix._

### Checklist

- [x]  I have read the [CONTRIBUTING.md](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

_Add any other context or screenshots about the feature request here._

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3282 from hung12ct:fix/mysql-pickle-serialization d9df37adb7dfbbfd8502a0fe65c4f8bca3d0d978
PiperOrigin-RevId: 825834360
2025-10-29 20:39:58 -07:00
Google Team MemberandCopybara-Service 9014a849ea feat: Add api key argument to Vertex Session and Memory services for Express Mode support
We also change VertexAiSessionService and VertexAiMemoryBankService to both use keyword arguments for project, location, agent engine id, and express mode api key

PiperOrigin-RevId: 825719331
2025-10-29 14:54:12 -07:00
Shangjie ChenandCopybara-Service f7e2a7a40e chore: Make VertexAiSessionService fully asynchronous
This CL refactors VertexAiSessionService to use the asynchronous aio client for all Vertex AI API calls. This ensures that the service methods are non-blocking and can be used effectively in an asyncio environment.

PiperOrigin-RevId: 824573356
2025-10-27 10:19:25 -07:00
Shangjie ChenandCopybara-Service 2a901d12f4 chore: Raise AlreadyExistsError when trying to create a resource with same ID
Move the dedupe logic into session service so that the internal error can be surfaced to client

PiperOrigin-RevId: 822294430
2025-10-21 15:11:33 -07:00
Shangjie ChenandCopybara-Service 6a5eac0bdc feat: Allow passing extra kwargs to create_session of VertexAiSessionService
This can be used to set ttl and other configs.

PiperOrigin-RevId: 821782343
2025-10-20 13:34:07 -07:00
Shangjie ChenandCopybara-Service 8b3ed059c2 chore: Refactor and fix state management in the session service
Also refactoring the test cases to focus on the expected behaviors

PiperOrigin-RevId: 820734484
2025-10-17 10:04:36 -07:00
Shangjie ChenandCopybara-Service 6ab1498aa0 fix: Add usage_metadata and citation_metadata to DatabaseSessionService
PiperOrigin-RevId: 819900773
2025-10-16 13:05:52 -07:00
Shangjie ChenandCopybara-Service a985cc38ec feat: Support return all sessions when user id is none
PiperOrigin-RevId: 819884236
2025-10-15 13:14:24 -07:00
Google Team MemberandCopybara-Service f9c09ef075 feat: Support returning all sessions when user_id is none in the request
resolves https://github.com/google/adk-python/issues/3154

PiperOrigin-RevId: 819417330
2025-10-14 15:10:41 -07:00
Shangjie ChenandCopybara-Service 141318f775 feat: Support returning all sessions when user_id is none in the request
resolves https://github.com/google/adk-python/issues/3154

PiperOrigin-RevId: 819401023
2025-10-14 14:30:47 -07:00
Wei Sun (Jack)andCopybara-Service 24342e95f8 chore: Remove temp state deltas before appending an event
PiperOrigin-RevId: 816902208
2025-10-08 16:13:52 -07:00
Shangjie ChenandCopybara-Service 636def3687 fix: Add AuthConfig json serialization in vertex ai session service
Resolves https://github.com/google/adk-python/issues/2035

PiperOrigin-RevId: 816387628
2025-10-07 14:47:28 -07:00
Shangjie ChenandCopybara-Service 90d4c19c51 feat: Migrate vertex ai session service to use agent engine sdk
PiperOrigin-RevId: 816259798
2025-10-07 09:58:27 -07:00
Google Team MemberandCopybara-Service d1efc8461e feat: Migrate vertex_ai_session_service to use Agent Engine SDK
PiperOrigin-RevId: 814967790
2025-10-03 22:14:23 -07:00
Shangjie ChenandCopybara-Service 97b950b36b feat: Migrate vertex_ai_session_service to use Agent Engine SDK
PiperOrigin-RevId: 814948921
2025-10-03 21:00:15 -07:00
Shangjie ChenandCopybara-Service e63fe0c0eb fix: Fix pagination of list_sessions in VertexAiSessionService
Resolves https://github.com/google/adk-python/issues/2860

PiperOrigin-RevId: 804511401
2025-09-08 11:14:32 -07:00
Shangjie ChenandCopybara-Service fb009d8ea6 fix: Add custom_metadata to DatabaseSessionService
Resolve https://github.com/google/adk-python/issues/2677

PiperOrigin-RevId: 802375768
2025-09-02 20:09:46 -07:00
Shangjie ChenandCopybara-Service 314d6a4f95 fix: Return session state in list_session API endpoint
Resolves https://github.com/google/adk-python/issues/2193
Resolves https://github.com/google/adk-python/issues/781

PiperOrigin-RevId: 789143973
2025-07-30 19:53:46 -07:00
Xiang (Sean) ZhouandCopybara-Service 927c75f0ee chore: Replace imports by importing from actual module instead of from package (__init__.py)
PiperOrigin-RevId: 786342250
2025-07-23 10:49:43 -07:00
Max ThorméandCopybara-Service d4f01afc15 fix(comment): fix comment to malicious user cannot obtain another user’s session
Merge https://github.com/google/adk-python/pull/1959

### What

Fix misleading comment.

```diff
- # Make sure a malicious user can obtain a session and events not belonging to them
+ # Make sure a malicious user **cannot** obtain a session or events not belonging to them
```

### Why

The previous wording contradicted the assertion `assert len(session_mismatch.events) == 0`, which verifies that a malicious user **cannot** access another user’s session or events.

### Testing plan

Docs-only change.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/1959 from mthorme:fix-comment-session-mismatch b1f139af340bd240d40ed58f5dea3274c3a3bd83
PiperOrigin-RevId: 785908548
2025-07-22 10:33:57 -07:00
Google Team MemberandCopybara-Service a50419961c fix: Process ListEvents response correctly when getSession is called
When get session is being called on a session with more than 1 page of events (100+), the response of the subsequent listevents calls fails due to missing parsing logic. This change fixes the processing of the listevents responses.

PiperOrigin-RevId: 783166959
2025-07-14 21:21:51 -07:00
Sindri Snær GunnarssonandSindri Snær Gunnarsson 362fb3f2b7 fix: scenario where a user can access another users events given the same session id
test: refine test_session_state in test_session_state to catch event leakage

fix: revert app name to my_app for test_session_state test

style: fix pyink style warnings

fix: add app_name to filter as per suggestion from rpedela-recurly
2025-07-02 18:56:08 +00:00