Initial commit

skills/memory/memory_discovery.py

@@ -0,0 +1,463 @@
+#!/usr/bin/env python3
+"""Memory Discovery and Query Utilities for AI Runtime
+
+- 加载 `.ai-runtime/memory/episodic/index.yml`
+- 提供 SQL 风格 (WHERE / ORDER BY / LIMIT) 的事件查询接口
+- 提供 table/json 两种格式化输出
+
+依赖：PyYAML（项目中已作为核心依赖使用）
+"""
+
+from __future__ import annotations
+
+import datetime as dt
+import json
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional
+
+import yaml
+
+
+@dataclass
+class MemoryEvent:
+    """单条情景记忆事件的索引信息"""
+
+    id: str
+    type: str
+    level: str
+    timestamp: dt.datetime
+    date_bucket: str
+    path: Path
+    title: str = ""
+    tags: List[str] = field(default_factory=list)
+    related: List[str] = field(default_factory=list)
+    meta: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def date(self) -> str:
+        """YYYY-MM-DD 字符串，便于 WHERE 子句使用 date 字段。"""
+
+        return self.timestamp.date().isoformat()
+
+    def to_dict(self) -> Dict[str, Any]:
+        """转换为可 JSON 序列化的字典。"""
+
+        return {
+            "id": self.id,
+            "type": self.type,
+            "level": self.level,
+            "timestamp": self.timestamp.isoformat(),
+            "date": self.date,
+            "date_bucket": self.date_bucket,
+            "path": str(self.path),
+            "title": self.title,
+            "tags": list(self.tags),
+            "related": list(self.related),
+            "meta": dict(self.meta),
+        }
+
+
+class MemoryDiscovery:
+    """Episodic 记忆索引加载与查询"""
+
+    def __init__(self, memory_root: Path) -> None:
+        self.memory_root = memory_root
+        self.episodic_root = memory_root / "episodic"
+        self.index_path = self.episodic_root / "index.yml"
+        self.events: List[MemoryEvent] = []
+        self.refresh()
+
+    # ------------------------------------------------------------------
+    # 加载索引
+    # ------------------------------------------------------------------
+    def refresh(self) -> None:
+        """重新加载索引文件。"""
+
+        self.events = self._load_events()
+
+    def _load_events(self) -> List[MemoryEvent]:
+        """从 episodic 目录扫描 Markdown 事件文件并解析元信息。"""
+
+        if not self.episodic_root.exists():
+            return []
+
+        events: List[MemoryEvent] = []
+        for md_path in self.episodic_root.rglob("*.md"):
+            event = self._parse_event_file(md_path)
+            if event is not None:
+                events.append(event)
+
+        return events
+
+    def _parse_event_file(self, path: Path) -> Optional[MemoryEvent]:
+        """解析单个事件 Markdown 文件。
+
+        协议：
+        - 可选顶部 YAML front matter: `--- ... ---`
+        - 正文中可使用:
+          - `# 标题` 作为事件标题
+          - `## 时间` 下第一行非空文本作为时间
+          - `## 标签` 下第一行非空文本作为标签列表
+        """
+
+        try:
+            text = path.read_text(encoding="utf-8")
+        except Exception:
+            return None
+
+        lines = text.splitlines()
+        front_matter, body_lines = self._parse_front_matter(lines)
+
+        stem = path.stem
+
+        # 基础字段
+        id_value = str(front_matter.get("id") or stem)
+        type_value = str(front_matter.get("type") or "event")
+
+        # level: 优先 front matter，其次目录结构推断
+        level_value = str(front_matter.get("level") or self._infer_level_from_path(path))
+
+        # 标题：优先 front matter.title，其次正文第一个 '# ' 标题
+        title = front_matter.get("title") or self._extract_title_from_body(body_lines) or stem
+
+        # 标签：支持 front matter.tags 或 '## 标签' 段
+        tags = front_matter.get("tags")
+        if isinstance(tags, str):
+            tags = [t.strip() for t in re.split(r"[,\s]+", tags) if t.strip()]
+        elif isinstance(tags, list):
+            tags = [str(t) for t in tags]
+        else:
+            tags = self._extract_tags_from_body(body_lines)
+
+        # 时间：front matter.timestamp/time → 正文 '## 时间' → 文件名/mtime 兜底
+        ts_str = front_matter.get("timestamp") or front_matter.get("time")
+        timestamp: Optional[dt.datetime] = None
+        if isinstance(ts_str, str):
+            timestamp = self._parse_datetime(ts_str)
+        if timestamp is None:
+            body_time = self._extract_time_from_body(body_lines)
+            if body_time:
+                timestamp = self._parse_datetime(body_time)
+        if timestamp is None:
+            timestamp = self._infer_datetime_from_filename_or_mtime(path)
+
+        if timestamp is None:
+            # 无法推断时间的事件对查询意义有限，忽略该文件
+            return None
+
+        # date_bucket: 优先 front matter.date_bucket，其次目录结构 / 时间推断
+        date_bucket = front_matter.get("date_bucket") or self._infer_date_bucket(path, timestamp)
+
+        related = front_matter.get("related") or []
+        if isinstance(related, str):
+            related = [related]
+        elif not isinstance(related, list):
+            related = []
+
+        # meta: 保留所有未被提升为显式字段的 front matter 信息
+        meta: Dict[str, Any] = dict(front_matter)
+        for k in [
+            "id",
+            "type",
+            "level",
+            "title",
+            "tags",
+            "timestamp",
+            "time",
+            "date_bucket",
+            "related",
+        ]:
+            meta.pop(k, None)
+
+        return MemoryEvent(
+            id=id_value,
+            type=type_value,
+            level=level_value,
+            timestamp=timestamp,
+            date_bucket=str(date_bucket),
+            path=path,
+            title=str(title),
+            tags=list(tags or []),
+            related=list(related),
+            meta=meta,
+        )
+
+    def _parse_front_matter(self, lines: List[str]):
+        """解析 YAML front matter，如果不存在则返回空字典和原始行。"""
+
+        if not lines:
+            return {}, []
+
+        if lines[0].strip() != "---":
+            return {}, lines
+
+        for i in range(1, len(lines)):
+            if lines[i].strip() == "---":
+                fm_text = "\n".join(lines[1:i])
+                try:
+                    data = yaml.safe_load(fm_text) or {}
+                except Exception:
+                    data = {}
+                return data, lines[i + 1 :]
+
+        # 未找到结束分隔符，视为无 front matter
+        return {}, lines
+
+    @staticmethod
+    def _extract_title_from_body(body_lines: List[str]) -> Optional[str]:
+        for line in body_lines:
+            s = line.strip()
+            if s.startswith("# "):
+                return s[2:].strip()
+        return None
+
+    @staticmethod
+    def _extract_time_from_body(body_lines: List[str]) -> Optional[str]:
+        for i, line in enumerate(body_lines):
+            if line.strip().startswith("## 时间"):
+                for j in range(i + 1, len(body_lines)):
+                    value = body_lines[j].strip()
+                    if value:
+                        return value
+                break
+        return None
+
+    @staticmethod
+    def _extract_tags_from_body(body_lines: List[str]) -> List[str]:
+        for i, line in enumerate(body_lines):
+            if line.strip().startswith("## 标签"):
+                for j in range(i + 1, len(body_lines)):
+                    raw = body_lines[j].strip()
+                    if not raw:
+                        continue
+                    parts = [p.strip() for p in re.split(r"[,\s]+", raw) if p.strip()]
+                    return parts
+                break
+        return []
+
+    def _infer_level_from_path(self, path: Path) -> str:
+        """根据相对路径推断级别: year/month/day/event。"""
+
+        try:
+            rel = path.relative_to(self.episodic_root)
+        except ValueError:
+            return "event"
+
+        parts = rel.parts
+        if len(parts) >= 3 and parts[0].isdigit() and parts[1].isdigit() and parts[2].isdigit():
+            return "day"
+        if len(parts) >= 2 and parts[0].isdigit() and parts[1].isdigit():
+            return "month"
+        if len(parts) >= 1 and parts[0].isdigit():
+            return "year"
+        return "event"
+
+    def _infer_date_bucket(self, path: Path, ts: dt.datetime) -> str:
+        """推断 date_bucket，例如 "2025/11/14"。"""
+
+        try:
+            rel = path.relative_to(self.episodic_root)
+            parts = rel.parts
+            if len(parts) >= 3 and parts[0].isdigit() and parts[1].isdigit() and parts[2].isdigit():
+                return f"{parts[0]}/{parts[1]}/{parts[2]}"
+        except ValueError:
+            pass
+
+        return ts.date().isoformat()
+
+    def _infer_datetime_from_filename_or_mtime(self, path: Path) -> Optional[dt.datetime]:
+        """从文件名 (YYYYMMDD-HHMM) 或 mtime 推断时间。"""
+
+        m = re.match(r"(\d{8})-(\d{4})", path.stem)
+        if m:
+            date_str, hm = m.groups()
+            try:
+                return dt.datetime.strptime(date_str + hm, "%Y%m%d%H%M")
+            except Exception:
+                pass
+
+        m2 = re.match(r"(\d{4})(\d{2})(\d{2})", path.stem)
+        if m2:
+            y, mth, d = m2.groups()
+            try:
+                return dt.datetime.strptime(f"{y}{mth}{d}", "%Y%m%d")
+            except Exception:
+                pass
+
+        try:
+            return dt.datetime.fromtimestamp(path.stat().st_mtime)
+        except Exception:
+            return None
+
+    # ------------------------------------------------------------------
+    # SQL 风格查询接口
+    # ------------------------------------------------------------------
+    def query(
+        self,
+        where: Optional[str] = None,
+        order_by: Optional[str] = None,
+        limit: Optional[int] = None,
+        offset: int = 0,
+    ) -> List[MemoryEvent]:
+        """基于 SQL 风格参数查询事件列表。"""
+
+        events: List[MemoryEvent] = list(self.events)
+
+        if where:
+            events = list(self._apply_where(events, where))
+
+        if order_by:
+            events = self._apply_order_by(events, order_by)
+
+        if offset:
+            events = events[offset:]
+
+        if limit is not None:
+            events = events[:limit]
+
+        return events
+
+    def _apply_where(
+        self, events: Iterable[MemoryEvent], where: str
+    ) -> Iterable[MemoryEvent]:
+        """简易 WHERE 解析，仅支持 AND，运算符子集。
+
+        支持的形式：
+        - field = 'value' / != / >= / <=
+        - tags CONTAINS 'tag'
+        - 通过 AND 连接多个条件（不支持 OR / 括号）
+        """
+
+        conditions = [part.strip() for part in re.split(r"\s+AND\s+", where, flags=re.I) if part.strip()]
+
+        def match(event: MemoryEvent) -> bool:
+            for cond in conditions:
+                if not self._eval_condition(event, cond):
+                    return False
+            return True
+
+        return (e for e in events if match(e))
+
+    def _eval_condition(self, event: MemoryEvent, cond: str) -> bool:
+        # tags CONTAINS 'tag'
+        if re.search(r"\bCONTAINS\b", cond, flags=re.I):
+            left, right = re.split(r"\bCONTAINS\b", cond, maxsplit=1, flags=re.I)
+            field = left.strip()
+            value = self._strip_quotes(right.strip())
+            if field.lower() != "tags":
+                return False
+            return value in (event.tags or [])
+
+        # field op value
+        m = re.match(r"^(\w+)\s*(=|!=|>=|<=)\s*(.+)$", cond)
+        if not m:
+            return False
+
+        field, op, raw_value = m.groups()
+        field = field.strip().lower()
+        value = self._strip_quotes(raw_value.strip())
+
+        # 取事件属性或 meta 字段
+        lhs: Any
+        if field == "id":
+            lhs = event.id
+        elif field == "type":
+            lhs = event.type
+        elif field == "level":
+            lhs = event.level
+        elif field == "title":
+            lhs = event.title
+        elif field == "date":
+            lhs = event.date
+        elif field == "timestamp":
+            lhs = event.timestamp
+        else:
+            if field in event.meta:
+                lhs = event.meta[field]
+            else:
+                # 未知字段直接返回 False，避免误匹配
+                return False
+
+        # 时间 / 日期字段支持 >= <=
+        if isinstance(lhs, dt.datetime):
+            rhs = self._parse_datetime(value)
+            if rhs is None:
+                return False
+        else:
+            rhs = value
+
+        try:
+            if op == "=":
+                return lhs == rhs
+            if op == "!=":
+                return lhs != rhs
+            if op == ">=":
+                return lhs >= rhs
+            if op == "<=":
+                return lhs <= rhs
+        except TypeError:
+            return False
+
+        return False
+
+    @staticmethod
+    def _strip_quotes(text: str) -> str:
+        if (text.startswith("'") and text.endswith("'")) or (
+            text.startswith('"') and text.endswith('"')
+        ):
+            return text[1:-1]
+        return text
+
+    @staticmethod
+    def _parse_datetime(value: str) -> Optional[dt.datetime]:
+        # 支持 "YYYY-MM-DD" 或 ISO8601 字符串
+        try:
+            if len(value) == 10:
+                return dt.datetime.fromisoformat(value + "T00:00:00")
+            return dt.datetime.fromisoformat(value)
+        except Exception:
+            return None
+
+    # ------------------------------------------------------------------
+    # 格式化输出
+    # ------------------------------------------------------------------
+    def format_events(
+        self,
+        events: List[MemoryEvent],
+        select: Optional[List[str]] = None,
+        format_type: str = "table",
+    ) -> str:
+        if select is None or not select:
+            select = ["id", "timestamp", "title"]
+
+        rows = []
+        for ev in events:
+            d = ev.to_dict()
+            rows.append({field: d.get(field) for field in select})
+
+        if format_type == "json":
+            return json.dumps(rows, ensure_ascii=False, indent=2)
+
+        # table 格式
+        return self._format_table(rows, select)
+
+    @staticmethod
+    def _format_table(rows: List[Dict[str, Any]], headers: List[str]) -> str:
+        if not rows:
+            return "(no events)"
+
+        # 计算列宽
+        widths: Dict[str, int] = {}
+        for h in headers:
+            widths[h] = max(len(h), *(len(str(row.get(h, ""))) for row in rows))
+
+        def fmt_row(row: Dict[str, Any]) -> str:
+            return "  ".join(str(row.get(h, "")).ljust(widths[h]) for h in headers)
+
+        header_line = "  ".join(h.ljust(widths[h]) for h in headers)
+        sep_line = "  ".join("-" * widths[h] for h in headers)
+        data_lines = [fmt_row(r) for r in rows]
+
+        return "\n".join([header_line, sep_line, *data_lines])