Add gfa_core/mas/memory module
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
potaser
Executable
+21
@@ -0,0 +1,21 @@
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry, SearchResult
|
||||
from gfa_core.mas.memory.case import AnalysisCase, CaseMemory
|
||||
from gfa_core.mas.memory.knowledge import KnowledgeEntry, KnowledgeMemory
|
||||
from gfa_core.mas.memory.manager import MemoryManager, MemoryStats
|
||||
from gfa_core.mas.memory.session import SessionData, SessionMemory
|
||||
from gfa_core.mas.memory.persistent import PersistentMemory
|
||||
|
||||
__all__ = [
|
||||
"BaseMemory",
|
||||
"MemoryEntry",
|
||||
"SearchResult",
|
||||
"SessionData",
|
||||
"SessionMemory",
|
||||
"AnalysisCase",
|
||||
"CaseMemory",
|
||||
"KnowledgeEntry",
|
||||
"KnowledgeMemory",
|
||||
"MemoryManager",
|
||||
"MemoryStats",
|
||||
"PersistentMemory",
|
||||
]
|
||||
Executable
+46
@@ -0,0 +1,46 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryEntry(Generic[T]):
|
||||
key: str
|
||||
value: T
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
ttl: int | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
key: str
|
||||
value: Any
|
||||
score: float
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class BaseMemory(ABC):
|
||||
@abstractmethod
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, key: str) -> bool:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def exists(self, key: str) -> bool:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def clear(self) -> None:
|
||||
pass
|
||||
Executable
+181
@@ -0,0 +1,181 @@
|
||||
import json
|
||||
import sqlite3
|
||||
import os
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, SearchResult
|
||||
from gfa_core.project_paths import MEMORY_ROOT
|
||||
|
||||
|
||||
@dataclass
|
||||
class AnalysisCase:
|
||||
case_id: str
|
||||
user_id: str | None = None
|
||||
gene_family: str | None = None
|
||||
species: str | None = None
|
||||
workflow: list[dict] = field(default_factory=list)
|
||||
results: dict[str, Any] = field(default_factory=dict)
|
||||
status: str = "pending"
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
completed_at: datetime | None = None
|
||||
feedback: dict[str, Any] = field(default_factory=dict)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class CaseMemory(BaseMemory):
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = os.environ.get("BBAV5_CASE_DB") or MEMORY_ROOT / "cases.db"
|
||||
self._db_path = Path(db_path)
|
||||
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._init_db()
|
||||
print(f"[CaseMemory] SQLite storage: {self._db_path}")
|
||||
|
||||
def _init_db(self) -> None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS cases (
|
||||
case_id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
gene_family TEXT,
|
||||
species TEXT,
|
||||
workflow TEXT,
|
||||
results TEXT,
|
||||
status TEXT,
|
||||
created_at TEXT,
|
||||
updated_at TEXT,
|
||||
completed_at TEXT,
|
||||
feedback TEXT,
|
||||
metadata TEXT
|
||||
)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_user_id ON cases(user_id)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_gene_family ON cases(gene_family)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_status ON cases(status)
|
||||
""")
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
if isinstance(value, AnalysisCase):
|
||||
case = value
|
||||
else:
|
||||
case = AnalysisCase(case_id=key, **value)
|
||||
|
||||
case.updated_at = datetime.now()
|
||||
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("""
|
||||
INSERT OR REPLACE INTO cases
|
||||
(case_id, user_id, gene_family, species, workflow, results, status,
|
||||
created_at, updated_at, completed_at, feedback, metadata)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
case.case_id,
|
||||
case.user_id,
|
||||
case.gene_family,
|
||||
case.species,
|
||||
json.dumps(case.workflow),
|
||||
json.dumps(case.results),
|
||||
case.status,
|
||||
case.created_at.isoformat(),
|
||||
case.updated_at.isoformat(),
|
||||
case.completed_at.isoformat() if case.completed_at else None,
|
||||
json.dumps(case.feedback),
|
||||
json.dumps(case.metadata),
|
||||
))
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
return self._row_to_case(row)
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"DELETE FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
return cursor.rowcount > 0
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT 1 FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
return cursor.fetchone() is not None
|
||||
|
||||
async def clear(self) -> None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("DELETE FROM cases")
|
||||
|
||||
def _row_to_case(self, row: sqlite3.Row) -> AnalysisCase:
|
||||
return AnalysisCase(
|
||||
case_id=row["case_id"],
|
||||
user_id=row["user_id"],
|
||||
gene_family=row["gene_family"],
|
||||
species=row["species"],
|
||||
workflow=json.loads(row["workflow"]) if row["workflow"] else [],
|
||||
results=json.loads(row["results"]) if row["results"] else {},
|
||||
status=row["status"],
|
||||
created_at=datetime.fromisoformat(row["created_at"]),
|
||||
updated_at=datetime.fromisoformat(row["updated_at"]),
|
||||
completed_at=datetime.fromisoformat(row["completed_at"]) if row["completed_at"] else None,
|
||||
feedback=json.loads(row["feedback"]) if row["feedback"] else {},
|
||||
metadata=json.loads(row["metadata"]) if row["metadata"] else {},
|
||||
)
|
||||
|
||||
async def list_by_user(self, user_id: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE user_id = ? ORDER BY created_at DESC",
|
||||
(user_id,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def list_by_gene_family(self, gene_family: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE gene_family = ? ORDER BY created_at DESC",
|
||||
(gene_family,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def list_by_status(self, status: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE status = ? ORDER BY created_at DESC",
|
||||
(status,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def update_status(self, case_id: str, status: str) -> None:
|
||||
case = await self.retrieve(case_id)
|
||||
if case:
|
||||
case.status = status
|
||||
if status == "completed":
|
||||
case.completed_at = datetime.now()
|
||||
await self.store(case_id, case)
|
||||
|
||||
async def add_feedback(self, case_id: str, feedback: dict[str, Any]) -> None:
|
||||
case = await self.retrieve(case_id)
|
||||
if case:
|
||||
case.feedback.update(feedback)
|
||||
await self.store(case_id, case)
|
||||
Executable
+338
@@ -0,0 +1,338 @@
|
||||
"""
|
||||
Conversation Manager - 类似ChatGPT的多会话管理
|
||||
支持用户维度、会话列表、历史消息
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional, List
|
||||
from enum import Enum
|
||||
|
||||
from gfa_core.project_paths import MEMORY_ROOT
|
||||
|
||||
|
||||
class ConversationStatus(str, Enum):
|
||||
"""会话状态 - 使用str继承以便JSON序列化"""
|
||||
ACTIVE = "active"
|
||||
ARCHIVED = "archived"
|
||||
DELETED = "deleted"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Conversation:
|
||||
"""会话（类似ChatGPT的一个聊天窗口）"""
|
||||
conversation_id: str
|
||||
user_id: str
|
||||
title: str
|
||||
status: str = "active" # 使用str而不是Enum
|
||||
|
||||
# 关联信息
|
||||
gene_family: Optional[str] = None
|
||||
species: Optional[str] = None
|
||||
job_id: Optional[str] = None
|
||||
input_dir: Optional[str] = None
|
||||
output_dir: Optional[str] = None
|
||||
working_dir: Optional[str] = None
|
||||
|
||||
# 消息历史
|
||||
messages: List[dict] = field(default_factory=list)
|
||||
|
||||
# 元数据
|
||||
created_at: str = field(default_factory=lambda: datetime.now().isoformat())
|
||||
updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
|
||||
message_count: int = 0
|
||||
|
||||
# 设置
|
||||
model: str = "glm-5"
|
||||
system_prompt: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversationSummary:
|
||||
"""会话摘要（用于列表展示）"""
|
||||
conversation_id: str
|
||||
title: str
|
||||
last_message: str
|
||||
updated_at: str
|
||||
message_count: int
|
||||
gene_family: Optional[str] = None
|
||||
job_id: Optional[str] = None
|
||||
input_dir: Optional[str] = None
|
||||
output_dir: Optional[str] = None
|
||||
working_dir: Optional[str] = None
|
||||
|
||||
|
||||
class ConversationManager:
|
||||
"""
|
||||
会话管理器
|
||||
|
||||
功能：
|
||||
1. 创建/删除/归档会话
|
||||
2. 按用户列出所有会话
|
||||
3. 获取会话消息历史
|
||||
4. 自动更新会话标题
|
||||
"""
|
||||
|
||||
def __init__(self, storage_dir: Path = None):
|
||||
if storage_dir is None:
|
||||
storage_dir = Path(os.environ.get("BBAV5_CONVERSATION_DIR") or MEMORY_ROOT / "conversations")
|
||||
self.storage_dir = Path(storage_dir)
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 用户索引：user_id -> list of conversation_ids
|
||||
self._user_index_file = self.storage_dir / "user_index.json"
|
||||
self._user_index: dict = {}
|
||||
self._load_user_index()
|
||||
|
||||
self._lock = asyncio.Lock()
|
||||
print(f"[ConversationManager] Storage: {self.storage_dir}")
|
||||
|
||||
def _load_user_index(self):
|
||||
"""加载用户索引"""
|
||||
if self._user_index_file.exists():
|
||||
try:
|
||||
with open(self._user_index_file, 'r', encoding='utf-8') as f:
|
||||
self._user_index = json.load(f)
|
||||
except Exception as e:
|
||||
print(f"[ConversationManager] Failed to load index: {e}")
|
||||
self._user_index = {}
|
||||
|
||||
async def _save_user_index(self):
|
||||
"""保存用户索引"""
|
||||
async with self._lock:
|
||||
with open(self._user_index_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self._user_index, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def _get_conv_file(self, conversation_id: str) -> Path:
|
||||
"""获取会话文件路径。
|
||||
|
||||
Existing conversation ids contain their creation date, e.g.
|
||||
conv_20260524_172646_default. Resolve that month first so old chats
|
||||
remain readable after the calendar month changes.
|
||||
"""
|
||||
match = None
|
||||
try:
|
||||
import re
|
||||
match = re.match(r"^conv_(\d{6})\d{2}_", str(conversation_id or ""))
|
||||
except Exception:
|
||||
match = None
|
||||
candidate_months = []
|
||||
if match:
|
||||
candidate_months.append(match.group(1))
|
||||
candidate_months.append(datetime.now().strftime("%Y%m"))
|
||||
seen = set()
|
||||
for date_prefix in candidate_months:
|
||||
if date_prefix in seen:
|
||||
continue
|
||||
seen.add(date_prefix)
|
||||
dir_path = self.storage_dir / date_prefix
|
||||
path = dir_path / f"{conversation_id}.json"
|
||||
if path.exists():
|
||||
return path
|
||||
dir_path = self.storage_dir / (candidate_months[0] if candidate_months else datetime.now().strftime("%Y%m"))
|
||||
dir_path.mkdir(exist_ok=True)
|
||||
return dir_path / f"{conversation_id}.json"
|
||||
|
||||
async def create_conversation(
|
||||
self,
|
||||
user_id: str,
|
||||
title: str = None,
|
||||
gene_family: str = None,
|
||||
species: str = None,
|
||||
system_prompt: str = None,
|
||||
job_id: str = None,
|
||||
input_dir: str = None,
|
||||
output_dir: str = None,
|
||||
working_dir: str = None,
|
||||
) -> Conversation:
|
||||
"""创建新会话"""
|
||||
conversation_id = f"conv_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{user_id[:8]}"
|
||||
|
||||
if title is None:
|
||||
title = f"新会话 {datetime.now().strftime('%m-%d %H:%M')}"
|
||||
|
||||
conversation = Conversation(
|
||||
conversation_id=conversation_id,
|
||||
user_id=user_id,
|
||||
title=title,
|
||||
gene_family=gene_family,
|
||||
species=species,
|
||||
system_prompt=system_prompt,
|
||||
job_id=job_id,
|
||||
input_dir=input_dir,
|
||||
output_dir=output_dir,
|
||||
working_dir=working_dir,
|
||||
)
|
||||
|
||||
# 保存会话
|
||||
await self._save_conversation(conversation)
|
||||
|
||||
# 更新用户索引
|
||||
if user_id not in self._user_index:
|
||||
self._user_index[user_id] = []
|
||||
self._user_index[user_id].insert(0, conversation_id)
|
||||
await self._save_user_index()
|
||||
|
||||
print(f"[ConversationManager] Created: {conversation_id} for user {user_id}")
|
||||
return conversation
|
||||
|
||||
async def _save_conversation(self, conversation: Conversation):
|
||||
"""保存会话到文件"""
|
||||
file_path = self._get_conv_file(conversation.conversation_id)
|
||||
async with self._lock:
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(asdict(conversation), f, ensure_ascii=False, indent=2)
|
||||
|
||||
async def get_conversation(self, conversation_id: str) -> Optional[Conversation]:
|
||||
"""获取会话详情"""
|
||||
file_path = self._get_conv_file(conversation_id)
|
||||
if not file_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
async with self._lock:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
return Conversation(**data)
|
||||
except Exception as e:
|
||||
print(f"[ConversationManager] Failed to load conversation: {e}")
|
||||
return None
|
||||
|
||||
async def list_user_conversations(
|
||||
self,
|
||||
user_id: str,
|
||||
status: str = "active",
|
||||
limit: int = 20,
|
||||
offset: int = 0,
|
||||
) -> List[ConversationSummary]:
|
||||
"""
|
||||
列出用户的所有会话（类似ChatGPT左侧边栏）
|
||||
|
||||
Returns:
|
||||
List[ConversationSummary]: 会话摘要列表，按更新时间倒序
|
||||
"""
|
||||
conversation_ids = self._user_index.get(user_id, [])
|
||||
|
||||
summaries = []
|
||||
for conv_id in conversation_ids:
|
||||
conv = await self.get_conversation(conv_id)
|
||||
if conv and conv.status == status:
|
||||
# 获取最后一条消息
|
||||
last_msg = ""
|
||||
if conv.messages:
|
||||
last_msg_content = conv.messages[-1].get('content', '')
|
||||
last_msg = last_msg_content[:50] + "..." if len(last_msg_content) > 50 else last_msg_content
|
||||
|
||||
summaries.append(ConversationSummary(
|
||||
conversation_id=conv.conversation_id,
|
||||
title=conv.title,
|
||||
last_message=last_msg,
|
||||
updated_at=conv.updated_at,
|
||||
message_count=conv.message_count,
|
||||
gene_family=conv.gene_family,
|
||||
job_id=conv.job_id,
|
||||
input_dir=conv.input_dir,
|
||||
output_dir=conv.output_dir,
|
||||
working_dir=conv.working_dir,
|
||||
))
|
||||
|
||||
# 按更新时间倒序
|
||||
summaries.sort(key=lambda x: x.updated_at, reverse=True)
|
||||
|
||||
return summaries[offset:offset + limit]
|
||||
|
||||
async def add_message(
|
||||
self,
|
||||
conversation_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
metadata: dict = None,
|
||||
) -> bool:
|
||||
"""添加消息到会话"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
message = {
|
||||
"role": role,
|
||||
"content": content,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"metadata": metadata or {},
|
||||
}
|
||||
|
||||
conv.messages.append(message)
|
||||
conv.message_count += 1
|
||||
conv.updated_at = datetime.now().isoformat()
|
||||
|
||||
# 自动更新标题（第一条用户消息）
|
||||
if conv.message_count == 1 and role == "user":
|
||||
conv.title = content[:30] + "..." if len(content) > 30 else content
|
||||
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
|
||||
async def get_messages(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> List[dict]:
|
||||
"""获取会话消息历史"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return []
|
||||
|
||||
return conv.messages[offset:offset + limit]
|
||||
|
||||
async def update_title(self, conversation_id: str, title: str) -> bool:
|
||||
"""更新会话标题"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.title = title
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
|
||||
async def delete_conversation(self, conversation_id: str) -> bool:
|
||||
"""删除会话"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.status = "deleted"
|
||||
await self._save_conversation(conv)
|
||||
|
||||
# 从用户索引中移除
|
||||
if conv.user_id in self._user_index:
|
||||
if conversation_id in self._user_index[conv.user_id]:
|
||||
self._user_index[conv.user_id].remove(conversation_id)
|
||||
await self._save_user_index()
|
||||
|
||||
return True
|
||||
|
||||
async def archive_conversation(self, conversation_id: str) -> bool:
|
||||
"""归档会话"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.status = "archived"
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
|
||||
async def link_to_job(self, conversation_id: str, job_id: str) -> bool:
|
||||
"""关联会话到分析任务"""
|
||||
conv = await self.get_conversation(conversation_id)
|
||||
if not conv:
|
||||
return False
|
||||
|
||||
conv.job_id = job_id
|
||||
await self._save_conversation(conv)
|
||||
return True
|
||||
Executable
+151
@@ -0,0 +1,151 @@
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry, SearchResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class KnowledgeEntry:
|
||||
entry_id: str
|
||||
content: str
|
||||
embedding: list[float] | None = None
|
||||
source: str | None = None
|
||||
category: str | None = None
|
||||
tags: list[str] = field(default_factory=list)
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class KnowledgeMemory(BaseMemory):
|
||||
def __init__(self, embedding_dim: int = 1536):
|
||||
self._embedding_dim = embedding_dim
|
||||
self._entries: dict[str, MemoryEntry[KnowledgeEntry]] = {}
|
||||
self._embeddings: dict[str, list[float]] = {}
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
if isinstance(value, KnowledgeEntry):
|
||||
entry = value
|
||||
else:
|
||||
entry = KnowledgeEntry(entry_id=key, **value)
|
||||
|
||||
memory_entry = MemoryEntry(key=key, value=entry)
|
||||
self._entries[key] = memory_entry
|
||||
|
||||
if entry.embedding:
|
||||
self._embeddings[key] = entry.embedding
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
entry = self._entries.get(key)
|
||||
return entry.value if entry else None
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
if key in self._entries:
|
||||
del self._entries[key]
|
||||
self._embeddings.pop(key, None)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
return key in self._entries
|
||||
|
||||
async def clear(self) -> None:
|
||||
self._entries.clear()
|
||||
self._embeddings.clear()
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 10,
|
||||
threshold: float = 0.0,
|
||||
) -> list[SearchResult]:
|
||||
if not self._embeddings:
|
||||
return []
|
||||
|
||||
scores = []
|
||||
for key, embedding in self._embeddings.items():
|
||||
score = self._cosine_similarity(query_embedding, embedding)
|
||||
if score >= threshold:
|
||||
scores.append((key, score))
|
||||
|
||||
scores.sort(key=lambda x: x[1], reverse=True)
|
||||
top_results = scores[:top_k]
|
||||
|
||||
results = []
|
||||
for key, score in top_results:
|
||||
entry = self._entries.get(key)
|
||||
if entry:
|
||||
results.append(SearchResult(
|
||||
key=key,
|
||||
value=entry.value,
|
||||
score=score,
|
||||
metadata=entry.metadata,
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
|
||||
if len(a) != len(b):
|
||||
return 0.0
|
||||
|
||||
dot_product = sum(x * y for x, y in zip(a, b))
|
||||
norm_a = sum(x * x for x in a) ** 0.5
|
||||
norm_b = sum(x * x for x in b) ** 0.5
|
||||
|
||||
if norm_a == 0 or norm_b == 0:
|
||||
return 0.0
|
||||
|
||||
return dot_product / (norm_a * norm_b)
|
||||
|
||||
async def search_by_tags(self, tags: list[str], top_k: int = 10) -> list[SearchResult]:
|
||||
results = []
|
||||
for key, entry in self._entries.items():
|
||||
if entry.value.tags and any(tag in entry.value.tags for tag in tags):
|
||||
results.append(SearchResult(
|
||||
key=key,
|
||||
value=entry.value,
|
||||
score=1.0,
|
||||
metadata=entry.metadata,
|
||||
))
|
||||
if len(results) >= top_k:
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
async def search_by_category(self, category: str, top_k: int = 10) -> list[SearchResult]:
|
||||
results = []
|
||||
for key, entry in self._entries.items():
|
||||
if entry.value.category == category:
|
||||
results.append(SearchResult(
|
||||
key=key,
|
||||
value=entry.value,
|
||||
score=1.0,
|
||||
metadata=entry.metadata,
|
||||
))
|
||||
if len(results) >= top_k:
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
async def add_entry(
|
||||
self,
|
||||
entry_id: str,
|
||||
content: str,
|
||||
embedding: list[float] | None = None,
|
||||
source: str | None = None,
|
||||
category: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
entry = KnowledgeEntry(
|
||||
entry_id=entry_id,
|
||||
content=content,
|
||||
embedding=embedding,
|
||||
source=source,
|
||||
category=category,
|
||||
tags=tags or [],
|
||||
metadata=metadata or {},
|
||||
)
|
||||
await self.store(entry_id, entry)
|
||||
Executable
+154
@@ -0,0 +1,154 @@
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from gfa_core.mas.memory.base import SearchResult
|
||||
from gfa_core.mas.memory.case import AnalysisCase, CaseMemory
|
||||
from gfa_core.mas.memory.knowledge import KnowledgeEntry, KnowledgeMemory
|
||||
from gfa_core.mas.memory.session import SessionData, SessionMemory
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryStats:
|
||||
session_count: int = 0
|
||||
case_count: int = 0
|
||||
knowledge_count: int = 0
|
||||
last_updated: datetime = field(default_factory=datetime.now)
|
||||
|
||||
|
||||
class MemoryManager:
|
||||
def __init__(
|
||||
self,
|
||||
session_memory: SessionMemory | None = None,
|
||||
case_memory: CaseMemory | None = None,
|
||||
knowledge_memory: KnowledgeMemory | None = None,
|
||||
):
|
||||
self.session = session_memory or SessionMemory()
|
||||
self.case = case_memory or CaseMemory()
|
||||
self.knowledge = knowledge_memory or KnowledgeMemory()
|
||||
|
||||
async def create_session(self, session_id: str, initial_state: dict[str, Any] | None = None) -> SessionData:
|
||||
session = SessionData(
|
||||
session_id=session_id,
|
||||
state=initial_state or {},
|
||||
)
|
||||
await self.session.store(session_id, session)
|
||||
return session
|
||||
|
||||
async def get_session(self, session_id: str) -> SessionData | None:
|
||||
return await self.session.retrieve(session_id)
|
||||
|
||||
async def add_message_to_session(self, session_id: str, role: str, content: str) -> None:
|
||||
await self.session.add_message(session_id, {"role": role, "content": content})
|
||||
|
||||
async def update_session_state(self, session_id: str, state: dict[str, Any]) -> None:
|
||||
await self.session.update_state(session_id, state)
|
||||
|
||||
async def create_case(
|
||||
self,
|
||||
case_id: str,
|
||||
user_id: str | None = None,
|
||||
gene_family: str | None = None,
|
||||
species: str | None = None,
|
||||
workflow: list[dict] | None = None,
|
||||
) -> AnalysisCase:
|
||||
case = AnalysisCase(
|
||||
case_id=case_id,
|
||||
user_id=user_id,
|
||||
gene_family=gene_family,
|
||||
species=species,
|
||||
workflow=workflow or [],
|
||||
)
|
||||
await self.case.store(case_id, case)
|
||||
return case
|
||||
|
||||
async def get_case(self, case_id: str) -> AnalysisCase | None:
|
||||
return await self.case.retrieve(case_id)
|
||||
|
||||
async def update_case_results(self, case_id: str, results: dict[str, Any]) -> None:
|
||||
case = await self.case.retrieve(case_id)
|
||||
if case:
|
||||
case.results.update(results)
|
||||
await self.case.store(case_id, case)
|
||||
|
||||
async def complete_case(self, case_id: str) -> None:
|
||||
await self.case.update_status(case_id, "completed")
|
||||
|
||||
async def add_case_feedback(self, case_id: str, feedback: dict[str, Any]) -> None:
|
||||
await self.case.add_feedback(case_id, feedback)
|
||||
|
||||
async def add_knowledge(
|
||||
self,
|
||||
entry_id: str,
|
||||
content: str,
|
||||
embedding: list[float] | None = None,
|
||||
source: str | None = None,
|
||||
category: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> None:
|
||||
await self.knowledge.add_entry(
|
||||
entry_id=entry_id,
|
||||
content=content,
|
||||
embedding=embedding,
|
||||
source=source,
|
||||
category=category,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
async def search_knowledge(
|
||||
self,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 10,
|
||||
) -> list[SearchResult]:
|
||||
return await self.knowledge.search(query_embedding, top_k)
|
||||
|
||||
async def search_knowledge_by_tags(self, tags: list[str], top_k: int = 10) -> list[SearchResult]:
|
||||
return await self.knowledge.search_by_tags(tags, top_k)
|
||||
|
||||
async def extract_case_to_knowledge(self, case_id: str) -> KnowledgeEntry | None:
|
||||
case = await self.case.retrieve(case_id)
|
||||
if not case or case.status != "completed":
|
||||
return None
|
||||
|
||||
entry_id = f"case_{case_id}"
|
||||
content = self._extract_knowledge_content(case)
|
||||
|
||||
entry = KnowledgeEntry(
|
||||
entry_id=entry_id,
|
||||
content=content,
|
||||
source=f"case:{case_id}",
|
||||
category="workflow_template",
|
||||
tags=[case.gene_family, case.species] if case.gene_family and case.species else [],
|
||||
metadata={
|
||||
"gene_family": case.gene_family,
|
||||
"species": case.species,
|
||||
"workflow_steps": len(case.workflow),
|
||||
},
|
||||
)
|
||||
|
||||
await self.knowledge.store(entry_id, entry)
|
||||
return entry
|
||||
|
||||
def _extract_knowledge_content(self, case: AnalysisCase) -> str:
|
||||
parts = []
|
||||
if case.gene_family:
|
||||
parts.append(f"Gene Family: {case.gene_family}")
|
||||
if case.species:
|
||||
parts.append(f"Species: {case.species}")
|
||||
if case.workflow:
|
||||
parts.append("Workflow:")
|
||||
for i, step in enumerate(case.workflow, 1):
|
||||
parts.append(f" {i}. {step.get('tool', 'unknown')}: {step.get('description', '')}")
|
||||
return "\n".join(parts)
|
||||
|
||||
async def get_stats(self) -> MemoryStats:
|
||||
return MemoryStats(
|
||||
session_count=len(self.session._sessions),
|
||||
case_count=len(await self.case.list_by_status("completed")),
|
||||
knowledge_count=len(self.knowledge._entries),
|
||||
)
|
||||
|
||||
async def clear_all(self) -> None:
|
||||
await self.session.clear()
|
||||
await self.case.clear()
|
||||
await self.knowledge.clear()
|
||||
Executable
+142
@@ -0,0 +1,142 @@
|
||||
"""
|
||||
持久化记忆存储 - 将记忆保存到文件
|
||||
"""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
from dataclasses import asdict
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry
|
||||
|
||||
|
||||
class PersistentMemory(BaseMemory):
|
||||
"""基于文件的持久化记忆存储"""
|
||||
|
||||
def __init__(self, storage_dir: Path, max_entries: int = 10000):
|
||||
self.storage_dir = Path(storage_dir)
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.max_entries = max_entries
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
def _get_file_path(self, key: str) -> Path:
|
||||
"""获取存储文件路径"""
|
||||
# 使用key的前2个字符作为子目录，避免单目录文件过多
|
||||
subdir = key[:2] if len(key) >= 2 else "xx"
|
||||
dir_path = self.storage_dir / subdir
|
||||
dir_path.mkdir(exist_ok=True)
|
||||
return dir_path / f"{key}.json"
|
||||
|
||||
def _serialize(self, value: Any) -> dict:
|
||||
"""序列化数据"""
|
||||
if hasattr(value, '__dataclass_fields__'):
|
||||
return {
|
||||
"_type": "dataclass",
|
||||
"_class": value.__class__.__name__,
|
||||
"data": asdict(value)
|
||||
}
|
||||
return {"_type": "raw", "data": value}
|
||||
|
||||
def _deserialize(self, data: dict) -> Any:
|
||||
"""反序列化数据"""
|
||||
from gfa_core.mas.memory.session import SessionData
|
||||
from gfa_core.mas.memory.case import AnalysisCase
|
||||
from gfa_core.mas.memory.knowledge import KnowledgeEntry
|
||||
|
||||
if data.get("_type") == "dataclass":
|
||||
class_name = data.get("_class")
|
||||
class_map = {
|
||||
"SessionData": SessionData,
|
||||
"AnalysisCase": AnalysisCase,
|
||||
"KnowledgeEntry": KnowledgeEntry,
|
||||
}
|
||||
cls = class_map.get(class_name)
|
||||
if cls:
|
||||
return cls(**data["data"])
|
||||
return data.get("data")
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
"""存储数据到文件"""
|
||||
async with self._lock:
|
||||
file_path = self._get_file_path(key)
|
||||
|
||||
entry = MemoryEntry(
|
||||
key=key,
|
||||
value=self._serialize(value),
|
||||
ttl=ttl,
|
||||
)
|
||||
|
||||
# 写入文件
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
json.dump({
|
||||
"key": entry.key,
|
||||
"value": entry.value,
|
||||
"ttl": entry.ttl,
|
||||
"created_at": entry.created_at.isoformat(),
|
||||
}, f, ensure_ascii=False, indent=2)
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
"""从文件读取数据"""
|
||||
async with self._lock:
|
||||
file_path = self._get_file_path(key)
|
||||
|
||||
if not file_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 检查TTL
|
||||
if data.get("ttl"):
|
||||
created_at = datetime.fromisoformat(data["created_at"])
|
||||
from datetime import timedelta
|
||||
if datetime.now() > created_at + timedelta(seconds=data["ttl"]):
|
||||
file_path.unlink() # 删除过期文件
|
||||
return None
|
||||
|
||||
return self._deserialize(data["value"])
|
||||
|
||||
except (json.JSONDecodeError, KeyError, ValueError):
|
||||
return None
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
"""删除数据"""
|
||||
async with self._lock:
|
||||
file_path = self._get_file_path(key)
|
||||
if file_path.exists():
|
||||
file_path.unlink()
|
||||
return True
|
||||
return False
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""检查数据是否存在"""
|
||||
file_path = self._get_file_path(key)
|
||||
if not file_path.exists():
|
||||
return False
|
||||
|
||||
# 检查是否过期
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
if data.get("ttl"):
|
||||
created_at = datetime.fromisoformat(data["created_at"])
|
||||
from datetime import timedelta
|
||||
if datetime.now() > created_at + timedelta(seconds=data["ttl"]):
|
||||
file_path.unlink()
|
||||
return False
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""清空所有数据"""
|
||||
async with self._lock:
|
||||
import shutil
|
||||
if self.storage_dir.exists():
|
||||
shutil.rmtree(self.storage_dir)
|
||||
self.storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
Executable
+144
@@ -0,0 +1,144 @@
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, MemoryEntry
|
||||
from gfa_core.mas.memory.persistent import PersistentMemory
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionData:
|
||||
session_id: str
|
||||
messages: list[dict] = field(default_factory=list)
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
|
||||
|
||||
class SessionMemory(BaseMemory):
|
||||
"""SessionMemory with persistent backend"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_sessions: int = 1000,
|
||||
default_ttl: int = 3600,
|
||||
persistent_storage: Optional[PersistentMemory] = None
|
||||
):
|
||||
self._sessions: OrderedDict[str, MemoryEntry[SessionData]] = OrderedDict()
|
||||
self._max_sessions = max_sessions
|
||||
self._default_ttl = default_ttl
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# 持久化后端
|
||||
self._persistent = persistent_storage
|
||||
if self._persistent:
|
||||
print(f"[SessionMemory] Persistent storage enabled: {self._persistent.storage_dir}")
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
async with self._lock:
|
||||
if len(self._sessions) >= self._max_sessions and key not in self._sessions:
|
||||
self._sessions.popitem(last=False)
|
||||
|
||||
if isinstance(value, SessionData):
|
||||
session_data = value
|
||||
else:
|
||||
session_data = SessionData(session_id=key, state=value)
|
||||
|
||||
entry = MemoryEntry(
|
||||
key=key,
|
||||
value=session_data,
|
||||
ttl=ttl or self._default_ttl,
|
||||
)
|
||||
self._sessions[key] = entry
|
||||
self._sessions.move_to_end(key)
|
||||
|
||||
# 持久化存储
|
||||
if self._persistent:
|
||||
await self._persistent.store(key, session_data, ttl)
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
async with self._lock:
|
||||
# 先查内存
|
||||
entry = self._sessions.get(key)
|
||||
|
||||
# 内存未命中，查持久化
|
||||
if entry is None and self._persistent:
|
||||
persistent_data = await self._persistent.retrieve(key)
|
||||
if persistent_data:
|
||||
self._sessions[key] = MemoryEntry(
|
||||
key=key,
|
||||
value=persistent_data,
|
||||
ttl=self._default_ttl
|
||||
)
|
||||
entry = self._sessions[key]
|
||||
|
||||
if entry is None:
|
||||
return None
|
||||
|
||||
if entry.ttl and datetime.now() > entry.created_at + timedelta(seconds=entry.ttl):
|
||||
del self._sessions[key]
|
||||
if self._persistent:
|
||||
await self._persistent.delete(key)
|
||||
return None
|
||||
|
||||
self._sessions.move_to_end(key)
|
||||
return entry.value
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
async with self._lock:
|
||||
deleted = False
|
||||
if key in self._sessions:
|
||||
del self._sessions[key]
|
||||
deleted = True
|
||||
if self._persistent:
|
||||
persistent_deleted = await self._persistent.delete(key)
|
||||
deleted = deleted or persistent_deleted
|
||||
return deleted
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
async with self._lock:
|
||||
# 先查内存
|
||||
entry = self._sessions.get(key)
|
||||
if entry:
|
||||
if entry.ttl and datetime.now() > entry.created_at + timedelta(seconds=entry.ttl):
|
||||
del self._sessions[key]
|
||||
if self._persistent:
|
||||
await self._persistent.delete(key)
|
||||
return False
|
||||
return True
|
||||
|
||||
# 查持久化
|
||||
if self._persistent:
|
||||
return await self._persistent.exists(key)
|
||||
return False
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._sessions.clear()
|
||||
if self._persistent:
|
||||
await self._persistent.clear()
|
||||
|
||||
async def add_message(self, session_id: str, message: dict) -> None:
|
||||
session = await self.retrieve(session_id)
|
||||
if session is None:
|
||||
session = SessionData(session_id=session_id)
|
||||
|
||||
session.messages.append(message)
|
||||
session.updated_at = datetime.now()
|
||||
await self.store(session_id, session)
|
||||
|
||||
async def get_messages(self, session_id: str) -> list[dict]:
|
||||
session = await self.retrieve(session_id)
|
||||
return session.messages if session else []
|
||||
|
||||
async def update_state(self, session_id: str, state: dict[str, Any]) -> None:
|
||||
session = await self.retrieve(session_id)
|
||||
if session is None:
|
||||
session = SessionData(session_id=session_id)
|
||||
|
||||
session.state.update(state)
|
||||
session.updated_at = datetime.now()
|
||||
await self.store(session_id, session)
|
||||
Reference in New Issue
Block a user