You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
Shangjie Chen
a501c59ac4
This change introduces a mechanism for users to register their own custom backend services for sessions, memory, and artifacts without modifying the ADK framework. This enhances the extensibility of ADK.
Two methods of registration are supported, both by placing a file in the parent directory of the agents.
**YAML Configuration (services.yaml or .yml)**
This is the recommended approach for simple services that can be instantiated with a constructor like MyService(uri="...", **kwargs).
Example services.yaml:
```
services:
- scheme: mysession
type: session
class: my_package.my_module.MyCustomSessionService
```
**Python Registration (services.py)**
For services requiring more complex initialization logic, users can define factory functions in a services.py file.
Example services.py
```
from google.adk.cli.service_registry import get_service_registry
from my_package.my_module import MyCustomSessionService
def my_session_factory(uri: str, **kwargs):
# custom initialization logic
return MyCustomSessionService(...)
get_service_registry().register_session_service("mysession", my_session_factory)
```
ADK will load services from services.yaml/.yml first, and then from services.py. If the same service scheme is defined in both, the registration in services.py will take precedence.
To use a registered service, specify its URI via the corresponding command-line flag, e.g., `--session_service_uri=mysession://....`
Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 831211371
97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
# Copyright 2025 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.
|
|
"""Dummy service implementations for testing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING
|
|
|
|
from google.adk.memory.base_memory_service import BaseMemoryService
|
|
from google.adk.memory.base_memory_service import SearchMemoryResponse
|
|
from google.adk.memory.memory_entry import MemoryEntry
|
|
from google.genai import types
|
|
from typing_extensions import override
|
|
|
|
if TYPE_CHECKING:
|
|
from google.adk.sessions.session import Session
|
|
|
|
|
|
class FooMemoryService(BaseMemoryService):
|
|
"""A dummy memory service that returns a fixed response."""
|
|
|
|
def __init__(self, uri: str | None = None, **kwargs):
|
|
"""Initializes the foo memory service.
|
|
|
|
Args:
|
|
uri: The service URI.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
del uri, kwargs # Unused in this dummy implementation.
|
|
|
|
@override
|
|
async def add_session_to_memory(self, session: Session):
|
|
print('FooMemoryService.add_session_to_memory')
|
|
|
|
@override
|
|
async def search_memory(
|
|
self, *, app_name: str, user_id: str, query: str
|
|
) -> SearchMemoryResponse:
|
|
print('FooMemoryService.search_memory')
|
|
return SearchMemoryResponse(
|
|
memories=[
|
|
MemoryEntry(
|
|
content=types.Content(
|
|
parts=[types.Part(text='I love ADK from Foo')]
|
|
),
|
|
author='bot',
|
|
timestamp=datetime.now().isoformat(),
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
class BarMemoryService(BaseMemoryService):
|
|
"""A dummy memory service that returns a fixed response."""
|
|
|
|
def __init__(self, uri: str | None = None, **kwargs):
|
|
"""Initializes the bar memory service.
|
|
|
|
Args:
|
|
uri: The service URI.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
del uri, kwargs # Unused in this dummy implementation.
|
|
|
|
@override
|
|
async def add_session_to_memory(self, session: Session):
|
|
print('BarMemoryService.add_session_to_memory')
|
|
|
|
@override
|
|
async def search_memory(
|
|
self, *, app_name: str, user_id: str, query: str
|
|
) -> SearchMemoryResponse:
|
|
print('BarMemoryService.search_memory')
|
|
return SearchMemoryResponse(
|
|
memories=[
|
|
MemoryEntry(
|
|
content=types.Content(
|
|
parts=[types.Part(text='I love ADK from Bar')]
|
|
),
|
|
author='bot',
|
|
timestamp=datetime.now().isoformat(),
|
|
)
|
|
]
|
|
)
|