ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,11 @@
|
||||
"""
|
||||
Database implementations for memory storage.
|
||||
"""
|
||||
|
||||
from .postgres import PostgresMemoryStore
|
||||
from .sqlite import SQLiteMemoryStore
|
||||
|
||||
__all__ = [
|
||||
"PostgresMemoryStore",
|
||||
"SQLiteMemoryStore"
|
||||
]
|
||||
@@ -0,0 +1,362 @@
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
from typing import Optional
|
||||
|
||||
import pytz # Add pytz for timezone handling
|
||||
from pydantic import BaseModel
|
||||
|
||||
from aworld.core.memory import MemoryStore
|
||||
from aworld.memory.models import (
|
||||
MemoryItem, MemoryAIMessage, MemoryHumanMessage, MemorySummary, MemorySystemMessage, MemoryToolMessage,
|
||||
MessageMetadata,
|
||||
UserProfile, AgentExperience, ConversationSummary
|
||||
)
|
||||
from aworld.models.model_response import ToolCall
|
||||
|
||||
try:
|
||||
from sqlalchemy.orm import declarative_base
|
||||
|
||||
Base = declarative_base()
|
||||
except ImportError:
|
||||
print("SQLAlchemy is not installed. Please install it to use PostgresMemoryStore.")
|
||||
# Get local timezone
|
||||
LOCAL_TZ = pytz.timezone('Asia/Shanghai') # Default to China timezone
|
||||
|
||||
def to_local_time(dt: datetime) -> str:
|
||||
"""Convert UTC datetime to local timezone string."""
|
||||
if not dt:
|
||||
return None
|
||||
if dt.tzinfo is None:
|
||||
dt = pytz.utc.localize(dt)
|
||||
return dt.astimezone(LOCAL_TZ).isoformat()
|
||||
|
||||
def from_iso_time(iso_str: str) -> datetime:
|
||||
"""Convert ISO format string to UTC datetime."""
|
||||
if not iso_str:
|
||||
return datetime.now(pytz.utc)
|
||||
try:
|
||||
dt = datetime.fromisoformat(iso_str)
|
||||
if dt.tzinfo is None:
|
||||
dt = LOCAL_TZ.localize(dt)
|
||||
return dt.astimezone(pytz.utc)
|
||||
except ValueError:
|
||||
return datetime.now(pytz.utc)
|
||||
|
||||
class MemoryItemModel(Base):
|
||||
from sqlalchemy import Column, String, DateTime, Boolean, Integer, Index
|
||||
from sqlalchemy.dialects.postgresql import ARRAY, JSONB
|
||||
|
||||
"""SQLAlchemy model for memory items."""
|
||||
__tablename__ = 'aworld_memory_items'
|
||||
|
||||
id = Column(String, primary_key=True)
|
||||
content = Column(JSONB) # Using JSONB for better performance
|
||||
created_at = Column(DateTime(timezone=True))
|
||||
updated_at = Column(DateTime(timezone=True))
|
||||
memory_meta = Column(JSONB) # Renamed from metadata to memory_meta
|
||||
tags = Column(ARRAY(String))
|
||||
memory_type = Column(String)
|
||||
version = Column(Integer)
|
||||
deleted = Column(Boolean, default=False)
|
||||
|
||||
# Create indexes
|
||||
__table_args__ = (
|
||||
Index('idx_memory_items_meta', memory_meta, postgresql_using='gin'),
|
||||
Index('idx_memory_items_tags', tags, postgresql_using='gin'),
|
||||
Index('idx_memory_items_type', memory_type),
|
||||
Index('idx_memory_items_created', created_at),
|
||||
)
|
||||
|
||||
class MemoryHistoryModel(Base):
|
||||
"""SQLAlchemy model for memory history."""
|
||||
__tablename__ = 'aworld_memory_histories'
|
||||
from sqlalchemy import Column, String, DateTime, ForeignKey
|
||||
|
||||
memory_id = Column(String, ForeignKey('aworld_memory_items.id'), primary_key=True)
|
||||
history_id = Column(String, ForeignKey('aworld_memory_items.id'), primary_key=True)
|
||||
created_at = Column(DateTime(timezone=True), default=datetime.utcnow)
|
||||
|
||||
|
||||
def orm_to_memory_item(orm_item: MemoryItemModel) -> Optional[MemoryItem]:
|
||||
"""Convert ORM model to MemoryItem."""
|
||||
if not orm_item:
|
||||
return None
|
||||
|
||||
memory_meta = orm_item.memory_meta or {}
|
||||
role = memory_meta.get('role')
|
||||
message_type = orm_item.memory_type
|
||||
|
||||
base_data = {
|
||||
'id': orm_item.id,
|
||||
'created_at': to_local_time(orm_item.created_at), # Convert to local time
|
||||
'updated_at': to_local_time(orm_item.updated_at), # Convert to local time
|
||||
'tags': orm_item.tags or [],
|
||||
'version': orm_item.version,
|
||||
'deleted': orm_item.deleted
|
||||
}
|
||||
|
||||
if role == 'system':
|
||||
return MemorySystemMessage(
|
||||
content=orm_item.content,
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
**base_data
|
||||
)
|
||||
elif role == 'user':
|
||||
return MemoryHumanMessage(
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
content=orm_item.content,
|
||||
**base_data
|
||||
)
|
||||
elif role == 'assistant':
|
||||
tool_calls_jsons = memory_meta.get('tool_calls', [])
|
||||
tool_calls = []
|
||||
for tool_calls_json in tool_calls_jsons:
|
||||
tool_call = ToolCall.from_dict(tool_calls_json)
|
||||
tool_calls.append(tool_call)
|
||||
return MemoryAIMessage(
|
||||
content=orm_item.content,
|
||||
tool_calls=tool_calls,
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
**base_data
|
||||
)
|
||||
elif role == 'tool':
|
||||
return MemoryToolMessage(
|
||||
tool_call_id=memory_meta.get('tool_call_id'),
|
||||
content=orm_item.content,
|
||||
status=memory_meta.get('status', 'success'),
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
**base_data
|
||||
)
|
||||
elif message_type == 'user_profile':
|
||||
if not orm_item.content:
|
||||
return None
|
||||
if not isinstance(orm_item.content, dict):
|
||||
return None
|
||||
|
||||
|
||||
return UserProfile(
|
||||
key=orm_item.content.get('key'),
|
||||
value=orm_item.content.get('value'),
|
||||
user_id=orm_item.memory_meta.get('user_id'),
|
||||
metadata=memory_meta,
|
||||
**base_data
|
||||
)
|
||||
elif message_type == 'agent_experience':
|
||||
if not orm_item.content:
|
||||
return None
|
||||
if not isinstance(orm_item.content, dict):
|
||||
return None
|
||||
return AgentExperience(
|
||||
skill=orm_item.content.get('skill'),
|
||||
actions=orm_item.content.get('actions'),
|
||||
agent_id=orm_item.memory_meta.get('agent_id'),
|
||||
metadata=memory_meta
|
||||
)
|
||||
elif message_type == 'summary':
|
||||
if not orm_item.content:
|
||||
return None
|
||||
if not isinstance(orm_item.content, str):
|
||||
return None
|
||||
# Extract item_ids from metadata
|
||||
item_ids = memory_meta.get('item_ids', [])
|
||||
# Create MessageMetadata from memory_meta
|
||||
summary_metadata = MessageMetadata(
|
||||
agent_id=memory_meta.get('agent_id'),
|
||||
agent_name=memory_meta.get('agent_name'),
|
||||
session_id=memory_meta.get('session_id'),
|
||||
task_id=memory_meta.get('task_id'),
|
||||
user_id=memory_meta.get('user_id')
|
||||
)
|
||||
return MemorySummary(
|
||||
item_ids=item_ids,
|
||||
summary=orm_item.content,
|
||||
metadata=summary_metadata,
|
||||
**base_data
|
||||
)
|
||||
elif message_type == 'conversation_summary':
|
||||
if not orm_item.content:
|
||||
return None
|
||||
if not isinstance(orm_item.content, str):
|
||||
return None
|
||||
# Preserve all custom metadata attributes
|
||||
conversation_summary_metadata = MessageMetadata(**memory_meta)
|
||||
return ConversationSummary(
|
||||
user_id=memory_meta.get('user_id'),
|
||||
session_id=memory_meta.get('session_id'),
|
||||
summary=orm_item.content,
|
||||
metadata=conversation_summary_metadata,
|
||||
**base_data
|
||||
)
|
||||
else:
|
||||
return MemoryItem(**{
|
||||
'id': orm_item.id,
|
||||
'content': orm_item.content,
|
||||
'created_at': to_local_time(orm_item.created_at), # Convert to local time
|
||||
'updated_at': to_local_time(orm_item.updated_at), # Convert to local time
|
||||
'metadata': memory_meta, # Map back to metadata for MemoryItem
|
||||
'tags': orm_item.tags or [],
|
||||
'memory_type': orm_item.memory_type,
|
||||
'version': orm_item.version,
|
||||
'deleted': orm_item.deleted
|
||||
})
|
||||
|
||||
|
||||
def memory_item_to_orm(item: MemoryItem) -> MemoryItemModel:
|
||||
"""Convert MemoryItem to ORM model."""
|
||||
# Handle content serialization
|
||||
content = item.content
|
||||
if isinstance(content, BaseModel):
|
||||
content = content.model_dump() # Use model_dump() instead of model_dump_json() for dict conversion
|
||||
|
||||
return MemoryItemModel(
|
||||
id=item.id,
|
||||
content=content, # Use serialized content
|
||||
created_at=from_iso_time(item.created_at), # Convert to UTC
|
||||
updated_at=from_iso_time(item.updated_at), # Convert to UTC
|
||||
memory_meta=item.metadata, # Map from metadata to memory_meta
|
||||
tags=item.tags,
|
||||
memory_type=item.memory_type,
|
||||
version=item.version,
|
||||
deleted=item.deleted
|
||||
)
|
||||
|
||||
|
||||
class PostgresMemoryStore(MemoryStore):
|
||||
"""
|
||||
PostgreSQL implementation of the memory store using SQLAlchemy.
|
||||
|
||||
This class provides a PostgreSQL-based storage backend for the memory system,
|
||||
implementing all required methods from the MemoryStore interface.
|
||||
"""
|
||||
|
||||
def __init__(self, db_url: str):
|
||||
"""
|
||||
Initialize PostgreSQL memory store.
|
||||
|
||||
Args:
|
||||
db_url (str): SQLAlchemy database URL
|
||||
Format: postgresql+psycopg2://user:password@host:port/dbname
|
||||
"""
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
self.engine = create_engine(db_url, echo=False, future=True)
|
||||
Base.metadata.create_all(self.engine)
|
||||
self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
|
||||
|
||||
def _build_filters(self, query, filters: dict = None):
|
||||
"""Build SQLAlchemy query filters."""
|
||||
if not filters:
|
||||
return query.filter(MemoryItemModel.deleted == False)
|
||||
|
||||
query = query.filter(MemoryItemModel.deleted == False)
|
||||
for key, value in filters.items():
|
||||
if value is not None:
|
||||
if key in ['user_id', 'agent_id', 'session_id', 'task_id', 'agent_name', 'tool_call_id']:
|
||||
query = query.filter(MemoryItemModel.memory_meta[key].astext == value)
|
||||
elif key == 'memory_type':
|
||||
# Handle memory_type as a list or single value
|
||||
if isinstance(value, list):
|
||||
query = query.filter(MemoryItemModel.memory_type.in_(value))
|
||||
else:
|
||||
query = query.filter(MemoryItemModel.memory_type == value)
|
||||
return query
|
||||
|
||||
def add(self, memory_item: MemoryItem):
|
||||
"""Add a new memory item to the store."""
|
||||
with self.Session() as session:
|
||||
orm_item = memory_item_to_orm(memory_item)
|
||||
session.add(orm_item)
|
||||
session.commit()
|
||||
|
||||
def get(self, memory_id) -> Optional[MemoryItem]:
|
||||
"""Get a memory item by ID."""
|
||||
with self.Session() as session:
|
||||
orm_item = session.query(MemoryItemModel).filter_by(
|
||||
id=memory_id, deleted=False
|
||||
).first()
|
||||
return orm_to_memory_item(orm_item)
|
||||
|
||||
def get_first(self, filters: dict = None) -> Optional[MemoryItem]:
|
||||
"""Get the first memory item matching the filters."""
|
||||
with self.Session() as session:
|
||||
query = session.query(MemoryItemModel)
|
||||
query = self._build_filters(query, filters)
|
||||
orm_item = query.order_by(MemoryItemModel.created_at.asc()).first()
|
||||
return orm_to_memory_item(orm_item)
|
||||
|
||||
def total_rounds(self, filters: dict = None) -> int:
|
||||
"""Get total number of memory rounds matching the filters."""
|
||||
with self.Session() as session:
|
||||
query = session.query(MemoryItemModel)
|
||||
query = self._build_filters(query, filters)
|
||||
return query.count()
|
||||
|
||||
def get_all(self, filters: dict = None) -> list[MemoryItem]:
|
||||
"""Get all memory items matching the filters."""
|
||||
with self.Session() as session:
|
||||
query = session.query(MemoryItemModel)
|
||||
query = self._build_filters(query, filters)
|
||||
orm_items = query.order_by(MemoryItemModel.created_at.asc()).all()
|
||||
return [orm_to_memory_item(item) for item in orm_items]
|
||||
|
||||
def get_last_n(self, last_rounds: int, filters: dict = None) -> list[MemoryItem]:
|
||||
"""Get the last N memory rounds matching the filters."""
|
||||
with self.Session() as session:
|
||||
query = session.query(MemoryItemModel)
|
||||
query = self._build_filters(query, filters)
|
||||
orm_items = query.order_by(MemoryItemModel.created_at.desc()).limit(last_rounds).all()
|
||||
return [orm_to_memory_item(item) for item in reversed(orm_items)]
|
||||
|
||||
def update(self, memory_item: MemoryItem):
|
||||
"""Update a memory item."""
|
||||
with self.Session() as session:
|
||||
orm_item = session.query(MemoryItemModel).filter_by(id=memory_item.id).first()
|
||||
if orm_item:
|
||||
orm_item.content = memory_item.content
|
||||
orm_item.created_at = from_iso_time(memory_item.created_at)
|
||||
orm_item.updated_at = from_iso_time(memory_item.updated_at) # Convert to UTC
|
||||
orm_item.memory_meta = memory_item.metadata
|
||||
orm_item.tags = memory_item.tags
|
||||
orm_item.memory_type = memory_item.memory_type
|
||||
orm_item.version = memory_item.version
|
||||
orm_item.deleted = memory_item.deleted
|
||||
session.commit()
|
||||
|
||||
def delete(self, memory_id):
|
||||
"""Soft delete a memory item."""
|
||||
with self.Session() as session:
|
||||
orm_item = session.query(MemoryItemModel).filter_by(id=memory_id).first()
|
||||
if orm_item:
|
||||
orm_item.deleted = True
|
||||
orm_item.updated_at = datetime.now(pytz.utc) # Use UTC time
|
||||
session.commit()
|
||||
|
||||
def delete_items(self, message_types: list[str], session_id: str, task_id: str, filters: dict = None):
|
||||
filters = filters or {}
|
||||
filters['memory_type'] = message_types
|
||||
filters['session_id'] = session_id
|
||||
filters['task_id'] = task_id
|
||||
with self.Session() as session:
|
||||
query = session.query(MemoryItemModel)
|
||||
query = self._build_filters(query, filters)
|
||||
query.update({
|
||||
MemoryItemModel.deleted: True,
|
||||
MemoryItemModel.updated_at: datetime.now(pytz.utc) # Use UTC time
|
||||
})
|
||||
session.commit()
|
||||
|
||||
def history(self, memory_id) -> list[MemoryItem] | None:
|
||||
"""Get the history of a memory item."""
|
||||
with self.Session() as session:
|
||||
history_items = session.query(MemoryItemModel).join(
|
||||
MemoryHistoryModel,
|
||||
MemoryHistoryModel.history_id == MemoryItemModel.id
|
||||
).filter(
|
||||
MemoryHistoryModel.memory_id == memory_id
|
||||
).order_by(MemoryItemModel.created_at.asc()).all()
|
||||
|
||||
if not history_items:
|
||||
return None
|
||||
|
||||
return [orm_to_memory_item(item) for item in history_items]
|
||||
@@ -0,0 +1,423 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
from typing import Optional, List, Dict, Any
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from aworld.core.memory import MemoryStore
|
||||
from aworld.memory.models import (
|
||||
MemoryItem, MemoryAIMessage, MemoryHumanMessage, MemorySummary,
|
||||
MemorySystemMessage, MemoryToolMessage, MessageMetadata,
|
||||
UserProfile, AgentExperience, ConversationSummary
|
||||
)
|
||||
from aworld.models.model_response import ToolCall
|
||||
|
||||
|
||||
class SQLiteMemoryStore(MemoryStore):
|
||||
"""
|
||||
SQLite implementation of the memory store.
|
||||
|
||||
This class provides a SQLite-based storage backend for the memory system,
|
||||
implementing all required methods from the MemoryStore interface.
|
||||
"""
|
||||
|
||||
def __init__(self, db_path: str = "./data/aworld_memory.db"):
|
||||
"""
|
||||
Initialize SQLite memory store.
|
||||
|
||||
Args:
|
||||
db_path (str): Path to SQLite database file
|
||||
"""
|
||||
self.db_path = Path(db_path)
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._init_database()
|
||||
|
||||
def _init_database(self) -> None:
|
||||
"""Initialize database tables and indexes."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS aworld_memory_items (
|
||||
id TEXT PRIMARY KEY,
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
memory_meta TEXT NOT NULL,
|
||||
tags TEXT NOT NULL,
|
||||
memory_type TEXT NOT NULL,
|
||||
version INTEGER NOT NULL DEFAULT 1,
|
||||
deleted BOOLEAN NOT NULL DEFAULT FALSE
|
||||
)
|
||||
""")
|
||||
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS aworld_memory_histories (
|
||||
memory_id TEXT NOT NULL,
|
||||
history_id TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
PRIMARY KEY (memory_id, history_id),
|
||||
FOREIGN KEY (memory_id) REFERENCES aworld_memory_items (id),
|
||||
FOREIGN KEY (history_id) REFERENCES aworld_memory_items (id)
|
||||
)
|
||||
""")
|
||||
|
||||
# Create indexes for better performance
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_memory_items_type ON aworld_memory_items (memory_type)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_memory_items_created ON aworld_memory_items (created_at)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_memory_items_deleted ON aworld_memory_items (deleted)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_memory_items_meta_user_id ON aworld_memory_items (json_extract(memory_meta, '$.user_id'))")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_memory_items_meta_agent_id ON aworld_memory_items (json_extract(memory_meta, '$.agent_id'))")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_memory_items_meta_session_id ON aworld_memory_items (json_extract(memory_meta, '$.session_id'))")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_memory_items_meta_task_id ON aworld_memory_items (json_extract(memory_meta, '$.task_id'))")
|
||||
|
||||
conn.commit()
|
||||
|
||||
def _serialize_content(self, content: Any) -> str:
|
||||
"""Serialize content to JSON string."""
|
||||
if content is None:
|
||||
return ""
|
||||
if isinstance(content, (dict, list, str, int, float, bool)):
|
||||
return json.dumps(content, ensure_ascii=False)
|
||||
if isinstance(content, BaseModel):
|
||||
return content.model_dump_json()
|
||||
return json.dumps(content, ensure_ascii=False, default=str)
|
||||
|
||||
def _deserialize_content(self, content_str: str) -> Any:
|
||||
"""Deserialize content from JSON string."""
|
||||
if not content_str:
|
||||
return None
|
||||
try:
|
||||
return json.loads(content_str)
|
||||
except json.JSONDecodeError:
|
||||
return content_str
|
||||
|
||||
def _serialize_metadata(self, metadata: Dict[str, Any]) -> str:
|
||||
"""Serialize metadata to JSON string."""
|
||||
if not metadata:
|
||||
return "{}"
|
||||
return json.dumps(metadata, ensure_ascii=False)
|
||||
|
||||
def _deserialize_metadata(self, metadata_str: str) -> Dict[str, Any]:
|
||||
"""Deserialize metadata from JSON string."""
|
||||
if not metadata_str:
|
||||
return {}
|
||||
try:
|
||||
return json.loads(metadata_str)
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
|
||||
def _serialize_tags(self, tags: List[str]) -> str:
|
||||
"""Serialize tags list to JSON string."""
|
||||
if not tags:
|
||||
return "[]"
|
||||
return json.dumps(tags, ensure_ascii=False)
|
||||
|
||||
def _deserialize_tags(self, tags_str: str) -> List[str]:
|
||||
"""Deserialize tags from JSON string."""
|
||||
if not tags_str:
|
||||
return []
|
||||
try:
|
||||
return json.loads(tags_str)
|
||||
except json.JSONDecodeError:
|
||||
return []
|
||||
|
||||
def _memory_item_to_row(self, item: MemoryItem) -> tuple:
|
||||
"""Convert MemoryItem to database row tuple."""
|
||||
content = self._serialize_content(item.content)
|
||||
metadata = self._serialize_metadata(item.metadata)
|
||||
tags = self._serialize_tags(item.tags)
|
||||
|
||||
return (
|
||||
item.id,
|
||||
content,
|
||||
item.created_at or datetime.now().isoformat(),
|
||||
item.updated_at or datetime.now().isoformat(),
|
||||
metadata,
|
||||
tags,
|
||||
item.memory_type,
|
||||
item.version,
|
||||
item.deleted
|
||||
)
|
||||
|
||||
def _row_to_memory_item(self, row: tuple) -> Optional[MemoryItem]:
|
||||
"""Convert database row to MemoryItem."""
|
||||
if not row:
|
||||
return None
|
||||
|
||||
(id_, content, created_at, updated_at, metadata, tags, memory_type, version, deleted) = row
|
||||
|
||||
memory_meta = self._deserialize_metadata(metadata)
|
||||
role = memory_meta.get('role')
|
||||
|
||||
base_data = {
|
||||
'id': id_,
|
||||
'created_at': created_at,
|
||||
'updated_at': updated_at,
|
||||
'tags': self._deserialize_tags(tags),
|
||||
'version': version,
|
||||
'deleted': bool(deleted)
|
||||
}
|
||||
|
||||
# Handle different message types
|
||||
if role == 'system':
|
||||
return MemorySystemMessage(
|
||||
content=self._deserialize_content(content),
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
**base_data
|
||||
)
|
||||
elif role == 'user':
|
||||
return MemoryHumanMessage(
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
content=self._deserialize_content(content),
|
||||
**base_data
|
||||
)
|
||||
elif role == 'assistant':
|
||||
tool_calls_jsons = memory_meta.get('tool_calls', [])
|
||||
tool_calls = []
|
||||
for tool_call_json in tool_calls_jsons:
|
||||
tool_call = ToolCall.from_dict(tool_call_json)
|
||||
tool_calls.append(tool_call)
|
||||
return MemoryAIMessage(
|
||||
content=self._deserialize_content(content),
|
||||
tool_calls=tool_calls,
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
**base_data
|
||||
)
|
||||
elif role == 'tool':
|
||||
return MemoryToolMessage(
|
||||
tool_call_id=memory_meta.get('tool_call_id'),
|
||||
content=self._deserialize_content(content),
|
||||
status=memory_meta.get('status', 'success'),
|
||||
metadata=MessageMetadata(**memory_meta),
|
||||
**base_data
|
||||
)
|
||||
elif memory_type == 'user_profile':
|
||||
content_data = self._deserialize_content(content)
|
||||
if not content_data or not isinstance(content_data, dict):
|
||||
return None
|
||||
return UserProfile(
|
||||
key=content_data.get('key'),
|
||||
value=content_data.get('value'),
|
||||
user_id=memory_meta.get('user_id'),
|
||||
metadata=memory_meta,
|
||||
**base_data
|
||||
)
|
||||
elif memory_type == 'agent_experience':
|
||||
content_data = self._deserialize_content(content)
|
||||
if not content_data or not isinstance(content_data, dict):
|
||||
return None
|
||||
return AgentExperience(
|
||||
skill=content_data.get('skill'),
|
||||
actions=content_data.get('actions'),
|
||||
agent_id=memory_meta.get('agent_id'),
|
||||
metadata=memory_meta,
|
||||
**base_data
|
||||
)
|
||||
elif memory_type == 'summary':
|
||||
content_data = self._deserialize_content(content)
|
||||
if not content_data or not isinstance(content_data, str):
|
||||
return None
|
||||
item_ids = memory_meta.get('item_ids', [])
|
||||
summary_metadata = MessageMetadata(
|
||||
agent_id=memory_meta.get('agent_id'),
|
||||
agent_name=memory_meta.get('agent_name'),
|
||||
session_id=memory_meta.get('session_id'),
|
||||
task_id=memory_meta.get('task_id'),
|
||||
user_id=memory_meta.get('user_id')
|
||||
)
|
||||
return MemorySummary(
|
||||
item_ids=item_ids,
|
||||
summary=content_data,
|
||||
metadata=summary_metadata,
|
||||
**base_data
|
||||
)
|
||||
elif memory_type == 'conversation_summary':
|
||||
content_data = self._deserialize_content(content)
|
||||
if not content_data or not isinstance(content_data, str):
|
||||
return None
|
||||
# Preserve all custom metadata attributes
|
||||
conversation_summary_metadata = MessageMetadata(**memory_meta)
|
||||
return ConversationSummary(
|
||||
user_id=memory_meta.get('user_id'),
|
||||
session_id=memory_meta.get('session_id'),
|
||||
summary=content_data,
|
||||
metadata=conversation_summary_metadata,
|
||||
**base_data
|
||||
)
|
||||
else:
|
||||
return MemoryItem(
|
||||
content=self._deserialize_content(content),
|
||||
metadata=memory_meta,
|
||||
memory_type=memory_type,
|
||||
**base_data
|
||||
)
|
||||
|
||||
def _build_filters(self, filters: Dict[str, Any] = None) -> tuple[str, tuple]:
|
||||
"""Build SQL WHERE clause and parameters from filters."""
|
||||
if not filters:
|
||||
return "WHERE deleted = FALSE", ()
|
||||
|
||||
conditions = ["deleted = FALSE"]
|
||||
params = []
|
||||
|
||||
for key, value in filters.items():
|
||||
if value is not None:
|
||||
if key in ['user_id', 'agent_id', 'session_id', 'task_id', 'agent_name', 'tool_call_id']:
|
||||
conditions.append(f"json_extract(memory_meta, '$.{key}') = ?")
|
||||
params.append(value)
|
||||
elif key == 'memory_type':
|
||||
if isinstance(value, list):
|
||||
placeholders = ','.join(['?' for _ in value])
|
||||
conditions.append(f"memory_type IN ({placeholders})")
|
||||
params.extend(value)
|
||||
else:
|
||||
conditions.append("memory_type = ?")
|
||||
params.append(value)
|
||||
|
||||
where_clause = "WHERE " + " AND ".join(conditions)
|
||||
return where_clause, tuple(params)
|
||||
|
||||
def add(self, memory_item: MemoryItem) -> None:
|
||||
"""Add a new memory item to the store."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
row = self._memory_item_to_row(memory_item)
|
||||
conn.execute("""
|
||||
INSERT INTO aworld_memory_items
|
||||
(id, content, created_at, updated_at, memory_meta, tags, memory_type, version, deleted)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", row)
|
||||
conn.commit()
|
||||
|
||||
def get(self, memory_id: str) -> Optional[MemoryItem]:
|
||||
"""Get a memory item by ID."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.execute("""
|
||||
SELECT id, content, created_at, updated_at, memory_meta, tags, memory_type, version, deleted
|
||||
FROM aworld_memory_items
|
||||
WHERE id = ? AND deleted = FALSE
|
||||
""", (memory_id,))
|
||||
row = cursor.fetchone()
|
||||
return self._row_to_memory_item(row)
|
||||
|
||||
def get_first(self, filters: Dict[str, Any] = None) -> Optional[MemoryItem]:
|
||||
"""Get the first memory item matching the filters."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
where_clause, params = self._build_filters(filters)
|
||||
cursor = conn.execute(f"""
|
||||
SELECT id, content, created_at, updated_at, memory_meta, tags, memory_type, version, deleted
|
||||
FROM aworld_memory_items
|
||||
{where_clause}
|
||||
ORDER BY created_at ASC
|
||||
LIMIT 1
|
||||
""", params)
|
||||
row = cursor.fetchone()
|
||||
return self._row_to_memory_item(row)
|
||||
|
||||
def total_rounds(self, filters: Dict[str, Any] = None) -> int:
|
||||
"""Get total number of memory rounds matching the filters."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
where_clause, params = self._build_filters(filters)
|
||||
cursor = conn.execute(f"""
|
||||
SELECT COUNT(*) FROM aworld_memory_items {where_clause}
|
||||
""", params)
|
||||
return cursor.fetchone()[0]
|
||||
|
||||
def get_all(self, filters: Dict[str, Any] = None) -> List[MemoryItem]:
|
||||
"""Get all memory items matching the filters."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
where_clause, params = self._build_filters(filters)
|
||||
cursor = conn.execute(f"""
|
||||
SELECT id, content, created_at, updated_at, memory_meta, tags, memory_type, version, deleted
|
||||
FROM aworld_memory_items
|
||||
{where_clause}
|
||||
ORDER BY created_at ASC
|
||||
""", params)
|
||||
rows = cursor.fetchall()
|
||||
return [self._row_to_memory_item(row) for row in rows if row]
|
||||
|
||||
def get_last_n(self, last_rounds: int, filters: Dict[str, Any] = None) -> List[MemoryItem]:
|
||||
"""Get the last N memory rounds matching the filters."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
where_clause, params = self._build_filters(filters)
|
||||
cursor = conn.execute(f"""
|
||||
SELECT id, content, created_at, updated_at, memory_meta, tags, memory_type, version, deleted
|
||||
FROM aworld_memory_items
|
||||
{where_clause}
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
""", params + (last_rounds,))
|
||||
rows = cursor.fetchall()
|
||||
# Reverse to maintain chronological order
|
||||
return [self._row_to_memory_item(row) for row in reversed(rows) if row]
|
||||
|
||||
def update(self, memory_item: MemoryItem) -> None:
|
||||
"""Update a memory item."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
row = self._memory_item_to_row(memory_item)
|
||||
conn.execute("""
|
||||
UPDATE aworld_memory_items
|
||||
SET content = ?, created_at = ?, updated_at = ?, memory_meta = ?,
|
||||
tags = ?, memory_type = ?, version = ?, deleted = ?
|
||||
WHERE id = ?
|
||||
""", row[1:] + (memory_item.id,))
|
||||
conn.commit()
|
||||
|
||||
def delete(self, memory_id: str) -> None:
|
||||
"""Soft delete a memory item."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.execute("""
|
||||
UPDATE aworld_memory_items
|
||||
SET deleted = TRUE, updated_at = ?
|
||||
WHERE id = ?
|
||||
""", (datetime.now().isoformat(), memory_id))
|
||||
conn.commit()
|
||||
|
||||
def delete_items(self, message_types: List[str], session_id: str, task_id: str, filters: Dict[str, Any] = None) -> None:
|
||||
"""Delete multiple memory items by message types, session_id, and task_id."""
|
||||
filters = filters or {}
|
||||
filters['memory_type'] = message_types
|
||||
filters['session_id'] = session_id
|
||||
filters['task_id'] = task_id
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
where_clause, params = self._build_filters(filters)
|
||||
# Remove the "WHERE" keyword and convert to proper WHERE clause for UPDATE
|
||||
where_conditions = where_clause.replace('WHERE ', '')
|
||||
conn.execute(f"""
|
||||
UPDATE aworld_memory_items
|
||||
SET deleted = TRUE, updated_at = ?
|
||||
WHERE {where_conditions}
|
||||
""", (datetime.now().isoformat(),) + params)
|
||||
conn.commit()
|
||||
|
||||
def history(self, memory_id: str) -> Optional[List[MemoryItem]]:
|
||||
"""Get the history of a memory item."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.execute("""
|
||||
SELECT m.id, m.content, m.created_at, m.updated_at, m.memory_meta,
|
||||
m.tags, m.memory_type, m.version, m.deleted
|
||||
FROM aworld_memory_items m
|
||||
JOIN aworld_memory_histories h ON h.history_id = m.id
|
||||
WHERE h.memory_id = ? AND m.deleted = FALSE
|
||||
ORDER BY m.created_at ASC
|
||||
""", (memory_id,))
|
||||
rows = cursor.fetchall()
|
||||
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
return [self._row_to_memory_item(row) for row in rows if row]
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close database connections."""
|
||||
# SQLite connections are automatically closed when exiting context managers
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry."""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit."""
|
||||
self.close()
|
||||
Reference in New Issue
Block a user