init
This commit is contained in:
raykkk
25
functionality/session_with_sqlite/README.md
Normal file
25
functionality/session_with_sqlite/README.md
Normal file
@@ -0,0 +1,25 @@
|
||||
# Session Management with Sqlite DB
|
||||
|
||||
This example demonstrates how to implement session management with a database backend. We use SQLite for simplicity,
|
||||
but the approach can be adapted for other databases.
|
||||
|
||||
Specifically, we implement a ``SqliteSession`` class that persists and retrieves session data from a SQLite table.
|
||||
The table schema includes fields for session ID, session data (stored as JSON), and timestamps for creation and last
|
||||
update.
|
||||
|
||||
We will create a simple agent and chat with it, then store the session data in the SQLite database. Then in the
|
||||
``test_load_session`` function, we will load the session data from the database and continue the chat.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Install agentscope from Pypi or source code.
|
||||
|
||||
```bash
|
||||
pip install agentscope
|
||||
```
|
||||
|
||||
Run the example by the following command
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
75
functionality/session_with_sqlite/main.py
Normal file
75
functionality/session_with_sqlite/main.py
Normal file
@@ -0,0 +1,75 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point for the session with SQLite example."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from sqlite_session import SqliteSession
|
||||
|
||||
SQLITE_PATH = "./session.db"
|
||||
|
||||
|
||||
async def main(username: str, query: str) -> None:
|
||||
"""Create an agent, load from session, chat with it, and save its state
|
||||
to SQLite.
|
||||
|
||||
Args:
|
||||
username (`str`):
|
||||
The username to identify the session.
|
||||
query (`str`):
|
||||
The user input query.
|
||||
"""
|
||||
|
||||
agent = ReActAgent(
|
||||
name="friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max",
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
)
|
||||
|
||||
# Create the SQLite session
|
||||
session = SqliteSession(SQLITE_PATH)
|
||||
|
||||
# Load the agent state by the given key "friday_of_user"
|
||||
# The load_session_state supports multiple state modules
|
||||
await session.load_session_state(
|
||||
session_id=username,
|
||||
friday_of_user=agent,
|
||||
)
|
||||
|
||||
# Chat with it to generate some state
|
||||
await agent(
|
||||
Msg("user", query, "user"),
|
||||
)
|
||||
|
||||
# Save the agent state by the given key "friday_of_user"
|
||||
# Also support multiple state modules (e.g. multiple agents)
|
||||
await session.save_session_state(
|
||||
session_id=username,
|
||||
friday_of_user=agent,
|
||||
)
|
||||
|
||||
|
||||
print("User named Alice chats with the agent ...")
|
||||
asyncio.run(main("alice", "What's the capital of America?"))
|
||||
|
||||
print("User named Bob chats with the agent ...")
|
||||
asyncio.run(main("bob", "What's the capital of China?"))
|
||||
|
||||
print(
|
||||
"\nNow, let's recover the session for Alice and ask about what the user "
|
||||
"asked before.",
|
||||
)
|
||||
asyncio.run(
|
||||
main(
|
||||
"alice",
|
||||
"What did I ask you before, what's your answer and how many "
|
||||
"questions have I asked you?",
|
||||
),
|
||||
)
|
||||
1
functionality/session_with_sqlite/requirements.txt
Normal file
1
functionality/session_with_sqlite/requirements.txt
Normal file
@@ -0,0 +1 @@
|
||||
agentscope[full]>=1.0.5
|
||||
167
functionality/session_with_sqlite/sqlite_session.py
Normal file
167
functionality/session_with_sqlite/sqlite_session.py
Normal file
@@ -0,0 +1,167 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The SQLite session class."""
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
|
||||
from agentscope import logger
|
||||
from agentscope.module import StateModule
|
||||
from agentscope.session import SessionBase
|
||||
|
||||
|
||||
class SqliteSession(SessionBase):
|
||||
"""A session that uses SQLite for storage."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sqlite_path: str,
|
||||
) -> None:
|
||||
"""Initialize the session.
|
||||
|
||||
Args:
|
||||
sqlite_path (`str`):
|
||||
The path to the SQLite database file.
|
||||
"""
|
||||
self.sqlite_path = sqlite_path
|
||||
|
||||
async def save_session_state(
|
||||
self,
|
||||
session_id: str,
|
||||
**state_modules_mapping: StateModule,
|
||||
) -> None:
|
||||
"""Save the session state to the SQLite database."""
|
||||
with sqlite3.connect(self.sqlite_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
# Prepare the session data as a dictionary
|
||||
session_data = {
|
||||
name: module.state_dict()
|
||||
for name, module in state_modules_mapping.items()
|
||||
}
|
||||
|
||||
json_data = json.dumps(session_data)
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS as_session (
|
||||
session_id TEXT,
|
||||
session_data JSON,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (session_id)
|
||||
)
|
||||
""",
|
||||
)
|
||||
|
||||
# Insert or replace the session data
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO as_session (session_id, session_data, updated_at)
|
||||
VALUES (?, json(?), CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(session_id) DO UPDATE SET
|
||||
session_data = excluded.session_data,
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(session_id, json_data),
|
||||
)
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
|
||||
async def load_session_state(
|
||||
self,
|
||||
session_id: str,
|
||||
allow_not_exist: bool = True,
|
||||
**state_modules_mapping: StateModule,
|
||||
) -> None:
|
||||
"""Get the state dictionary from the SQLite database.
|
||||
|
||||
Args:
|
||||
session_id (`str`):
|
||||
The session id.
|
||||
allow_not_exist (`bool`, defaults to `True`):
|
||||
Whether to allow the session to not exist. If `False`, raises
|
||||
an error if the session does not exist.
|
||||
**state_modules_mapping (`list[StateModule]`):
|
||||
The list of state modules to be loaded.
|
||||
"""
|
||||
if not os.path.exists(self.sqlite_path):
|
||||
if allow_not_exist:
|
||||
logger.info(
|
||||
"SQLite database %s does not exist. "
|
||||
"Skipping load for session_id %s.",
|
||||
self.sqlite_path,
|
||||
session_id,
|
||||
)
|
||||
return
|
||||
raise ValueError(
|
||||
"Failed to load session state because the SQLite database "
|
||||
f"file '{self.sqlite_path}' does not exist.",
|
||||
)
|
||||
|
||||
with sqlite3.connect(self.sqlite_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# If the table does not exist, return
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT name FROM sqlite_master WHERE type='table' AND
|
||||
name='as_session';
|
||||
""",
|
||||
)
|
||||
if cursor.fetchone() is None:
|
||||
if allow_not_exist:
|
||||
logger.info(
|
||||
"Session table does not exist in database %s. "
|
||||
"Skipping load for session_id %s.",
|
||||
self.sqlite_path,
|
||||
session_id,
|
||||
)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
"Failed to load session state because the session "
|
||||
"table 'as_session' does not exist in database "
|
||||
f"{self.sqlite_path}.",
|
||||
)
|
||||
|
||||
# Query the session data
|
||||
cursor.execute(
|
||||
"SELECT session_data FROM as_session WHERE session_id = ?",
|
||||
(session_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
if allow_not_exist:
|
||||
logger.info(
|
||||
"Session_id %s does not exist in database %s. "
|
||||
"Skip loading.",
|
||||
session_id,
|
||||
self.sqlite_path,
|
||||
)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
f"Failed to load session state for session_id "
|
||||
f"{session_id} does not exist.",
|
||||
)
|
||||
|
||||
session_data = json.loads(row[0])
|
||||
|
||||
for name, module in state_modules_mapping.items():
|
||||
if name in session_data:
|
||||
module.load_state_dict(session_data[name])
|
||||
else:
|
||||
raise ValueError(
|
||||
f"State module '{name}' not found in session "
|
||||
"data.",
|
||||
)
|
||||
logger.info(
|
||||
"Load session state for session_id %s from "
|
||||
"database %s successfully.",
|
||||
session_id,
|
||||
self.sqlite_path,
|
||||
)
|
||||
|
||||
finally:
|
||||
cursor.close()
|
||||
Reference in New Issue
Block a user