diff --git a/gfa_core/mas/memory/__init__.py b/gfa_core/mas/memory/__init__.py new file mode 100755 index 000000000..69ac2f0fb --- /dev/null +++ b/gfa_core/mas/memory/__init__.py @@ -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", +] diff --git a/gfa_core/mas/memory/base.py b/gfa_core/mas/memory/base.py new file mode 100755 index 000000000..823d34f41 --- /dev/null +++ b/gfa_core/mas/memory/base.py @@ -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 diff --git a/gfa_core/mas/memory/case.py b/gfa_core/mas/memory/case.py new file mode 100755 index 000000000..c836afdf3 --- /dev/null +++ b/gfa_core/mas/memory/case.py @@ -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) diff --git a/gfa_core/mas/memory/conversation.py b/gfa_core/mas/memory/conversation.py new file mode 100755 index 000000000..ca0d1fd27 --- /dev/null +++ b/gfa_core/mas/memory/conversation.py @@ -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 diff --git a/gfa_core/mas/memory/knowledge.py b/gfa_core/mas/memory/knowledge.py new file mode 100755 index 000000000..a316f187d --- /dev/null +++ b/gfa_core/mas/memory/knowledge.py @@ -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) diff --git a/gfa_core/mas/memory/manager.py b/gfa_core/mas/memory/manager.py new file mode 100755 index 000000000..ae4f6e9c4 --- /dev/null +++ b/gfa_core/mas/memory/manager.py @@ -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() diff --git a/gfa_core/mas/memory/persistent.py b/gfa_core/mas/memory/persistent.py new file mode 100755 index 000000000..9be9921c1 --- /dev/null +++ b/gfa_core/mas/memory/persistent.py @@ -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) diff --git a/gfa_core/mas/memory/session.py b/gfa_core/mas/memory/session.py new file mode 100755 index 000000000..2961a9035 --- /dev/null +++ b/gfa_core/mas/memory/session.py @@ -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)