python-library-ai-agent 0.1.0__py3-none-any.whl

ai_agent/memory/snapshot_merge.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+from ai_agent.memory.compression_work import (
+    CompressImportantResult,
+    CompressImportantWork,
+    CompressLongResult,
+    CompressLongWork,
+    CompressionResult,
+    CompressionWork,
+    DateToLongResult,
+    DateToLongWork,
+    ShortToDateResult,
+    ShortToDateWork,
+)
+from ai_agent.memory.config import MemoryConfig
+from ai_agent.memory.models import (
+    DateMemoryDay,
+    ImportantMemoryEntry,
+    MemorySnapshot,
+)
+from ai_agent.memory.worker import MemoryTask, MemoryTaskKind
+
+
+def prepare_work(
+    snapshot: MemorySnapshot,
+    task: MemoryTask,
+    *,
+    config: MemoryConfig,
+) -> CompressionWork | None:
+    if task.kind == MemoryTaskKind.SHORT_TO_DATE:
+        batch_size = int(task.payload.get("batch_size", 1))
+        if not snapshot.short_term:
+            return None
+        take = min(batch_size, len(snapshot.short_term))
+        batch = tuple(
+            m.model_copy(deep=True) for m in snapshot.short_term[:take]
+        )
+        return ShortToDateWork(batch=batch, batch_size=take)
+    if task.kind == MemoryTaskKind.DATE_TO_LONG:
+        day_label = str(task.payload["day"])
+        day = _find_day(snapshot, day_label)
+        if day is None or not day.entries:
+            return None
+        entries = tuple(e.model_copy(deep=True) for e in day.entries)
+        return DateToLongWork(day_label=day_label, entries=entries)
+    if task.kind == MemoryTaskKind.COMPRESS_LONG:
+        if len(snapshot.long_term) <= config.long_term_max_chunks:
+            return None
+        chunks = tuple(c.model_copy(deep=True) for c in snapshot.long_term)
+        return CompressLongWork(chunks=chunks)
+    if task.kind == MemoryTaskKind.COMPRESS_IMPORTANT:
+        if len(snapshot.important) <= config.important_max_entries:
+            return None
+        entries = tuple(e.model_copy(deep=True) for e in snapshot.important)
+        return CompressImportantWork(entries=entries)
+    return None
+
+
+def apply_result(
+    snapshot: MemorySnapshot,
+    task: MemoryTask,
+    work: CompressionWork,
+    result: CompressionResult,
+    *,
+    config: MemoryConfig,
+) -> list[MemoryTask]:
+    if task.kind == MemoryTaskKind.SHORT_TO_DATE:
+        return _apply_short_to_date(snapshot, work, result, config=config)
+    if task.kind == MemoryTaskKind.DATE_TO_LONG:
+        _apply_date_to_long(snapshot, work, result)
+        return []
+    if task.kind == MemoryTaskKind.COMPRESS_LONG:
+        _apply_compress_long(snapshot, result)
+        return []
+    if task.kind == MemoryTaskKind.COMPRESS_IMPORTANT:
+        _apply_compress_important(snapshot, result)
+        return []
+    return []
+
+
+def _apply_short_to_date(
+    snapshot: MemorySnapshot,
+    work: CompressionWork,
+    result: CompressionResult,
+    *,
+    config: MemoryConfig,
+) -> list[MemoryTask]:
+    if not isinstance(work, ShortToDateWork) or not isinstance(
+        result,
+        ShortToDateResult,
+    ):
+        return []
+    remove = min(work.batch_size, len(snapshot.short_term))
+    if remove:
+        snapshot.short_term = snapshot.short_term[remove:]
+    day = _find_or_create_day(snapshot, result.day_label)
+    day.entries.extend(result.entries)
+    now = datetime.now(timezone.utc)
+    for text in result.important_texts:
+        snapshot.important.append(
+            ImportantMemoryEntry(at=now, content=text, source="short_term"),
+        )
+    followups: list[MemoryTask] = []
+    if day is not None and len(day.entries) > config.date_memory_max_entries_per_day:
+        followups.append(
+            MemoryTask(MemoryTaskKind.DATE_TO_LONG, {"day": result.day_label}),
+        )
+    return followups
+
+
+def _apply_date_to_long(
+    snapshot: MemorySnapshot,
+    work: CompressionWork,
+    result: CompressionResult,
+) -> None:
+    if not isinstance(work, DateToLongWork) or not isinstance(
+        result,
+        DateToLongResult,
+    ):
+        return
+    snapshot.long_term.append(result.chunk)
+    snapshot.date_days = [d for d in snapshot.date_days if d.date != work.day_label]
+
+
+def _apply_compress_long(
+    snapshot: MemorySnapshot,
+    result: CompressionResult,
+) -> None:
+    if not isinstance(result, CompressLongResult):
+        return
+    snapshot.long_term = list(result.chunks)
+
+
+def _apply_compress_important(
+    snapshot: MemorySnapshot,
+    result: CompressionResult,
+) -> None:
+    if not isinstance(result, CompressImportantResult):
+        return
+    snapshot.important = list(result.entries)
+
+
+def _find_or_create_day(snapshot: MemorySnapshot, day_label: str) -> DateMemoryDay:
+    existing = _find_day(snapshot, day_label)
+    if existing is not None:
+        return existing
+    day = DateMemoryDay(date=day_label, entries=[])
+    snapshot.date_days.append(day)
+    return day
+
+
+def _find_day(snapshot: MemorySnapshot, day_label: str) -> DateMemoryDay | None:
+    for day in snapshot.date_days:
+        if day.date == day_label:
+            return day
+    return None

ai_agent/memory/store.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ai_agent.memory.models import (
+    DateMemoryDay,
+    ImportantMemoryEntry,
+    LongTermChunk,
+    MemoryMessage,
+    MemorySnapshot,
+)
+
+
+class MemoryStore:
+    """将会话记忆读写为存储目录下的 JSON 文件。"""
+
+    def __init__(self, root: Path) -> None:
+        self._root = root
+        self._short_term_path = root / "short_term.json"
+        self._date_dir = root / "date"
+        self._long_term_path = root / "long_term.json"
+        self._important_path = root / "important.json"
+
+    @property
+    def root(self) -> Path:
+        return self._root
+
+    def ensure_layout(self) -> None:
+        self._root.mkdir(parents=True, exist_ok=True)
+        self._date_dir.mkdir(parents=True, exist_ok=True)
+
+    def load(self) -> MemorySnapshot:
+        self.ensure_layout()
+        short_term = _read_list(self._short_term_path, MemoryMessage)
+        date_days = self._load_date_days()
+        long_term = _read_list(self._long_term_path, LongTermChunk)
+        important = _read_list(self._important_path, ImportantMemoryEntry)
+        return MemorySnapshot(
+            short_term=short_term,
+            date_days=date_days,
+            long_term=long_term,
+            important=important,
+        )
+
+    def save_short_term(self, messages: list[MemoryMessage]) -> None:
+        self.ensure_layout()
+        _write_list(self._short_term_path, messages)
+
+    def save_date_day(self, day: DateMemoryDay) -> None:
+        self.ensure_layout()
+        path = self._date_dir / f"{day.date}.json"
+        path.write_text(
+            day.model_dump_json(indent=2),
+            encoding="utf-8",
+        )
+
+    def remove_date_day(self, date_label: str) -> None:
+        path = self._date_dir / f"{date_label}.json"
+        if path.is_file():
+            path.unlink()
+
+    def prune_date_files(self, active_dates: set[str]) -> None:
+        if not self._date_dir.is_dir():
+            return
+        for path in self._date_dir.glob("*.json"):
+            if path.stem not in active_dates:
+                path.unlink(missing_ok=True)
+
+    def save_long_term(self, chunks: list[LongTermChunk]) -> None:
+        self.ensure_layout()
+        _write_list(self._long_term_path, chunks)
+
+    def save_important(self, entries: list[ImportantMemoryEntry]) -> None:
+        self.ensure_layout()
+        _write_list(self._important_path, entries)
+
+    def _load_date_days(self) -> list[DateMemoryDay]:
+        if not self._date_dir.is_dir():
+            return []
+        days: list[DateMemoryDay] = []
+        for path in sorted(self._date_dir.glob("*.json")):
+            raw = path.read_text(encoding="utf-8")
+            if not raw.strip():
+                continue
+            days.append(DateMemoryDay.model_validate_json(raw))
+        return days
+
+
+def _read_list(path: Path, model: type) -> list:
+    if not path.is_file():
+        return []
+    raw = path.read_text(encoding="utf-8")
+    if not raw.strip():
+        return []
+    data = json.loads(raw)
+    if not isinstance(data, list):
+        return []
+    return [model.model_validate(item) for item in data]
+
+
+def _write_list(path: Path, items: list) -> None:
+    payload = [item.model_dump(mode="json") for item in items]
+    path.write_text(
+        json.dumps(payload, ensure_ascii=False, indent=2),
+        encoding="utf-8",
+    )

ai_agent/memory/worker.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import enum
+import queue
+import threading
+from dataclasses import dataclass
+from datetime import date, datetime, timedelta, timezone
+from typing import Callable
+
+from ai_agent.memory.compression_work import (
+    CompressImportantWork,
+    CompressLongWork,
+    CompressionResult,
+    CompressionWork,
+    CompressImportantResult,
+    CompressLongResult,
+    DateToLongResult,
+    DateToLongWork,
+    ShortToDateResult,
+    ShortToDateWork,
+)
+from ai_agent.memory.compressor import MemoryCompressor
+from ai_agent.memory.config import MemoryConfig
+class MemoryTaskKind(str, enum.Enum):
+    SHORT_TO_DATE = "short_to_date"
+    DATE_TO_LONG = "date_to_long"
+    COMPRESS_LONG = "compress_long"
+    COMPRESS_IMPORTANT = "compress_important"
+
+
+@dataclass
+class MemoryTask:
+    kind: MemoryTaskKind
+    payload: dict
+
+
+class MemoryWorker:
+    """
+    独立线程上的记忆维护：弹出短期、归档日期、压缩长期与重要记忆。
+
+    压缩在锁外执行；提交结果时由调用方短暂持锁合并并发布 Agent 视图。
+    """
+
+    def __init__(
+        self,
+        *,
+        config: MemoryConfig,
+        compressor_factory: Callable[[], MemoryCompressor],
+        prepare: Callable[[MemoryTask], CompressionWork | None],
+        commit: Callable[[MemoryTask, CompressionWork, CompressionResult], None],
+    ) -> None:
+        self._config = config
+        self._compressor_factory = compressor_factory
+        self._prepare = prepare
+        self._commit = commit
+        self._queue: queue.Queue[MemoryTask | None] = queue.Queue()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        if self._thread is not None and self._thread.is_alive():
+            return
+        self._thread = threading.Thread(
+            target=self._thread_main,
+            name="memory-worker",
+            daemon=True,
+        )
+        self._thread.start()
+
+    def stop(self, *, join: bool = True, timeout: float | None = 5.0) -> None:
+        self._queue.put(None)
+        if join and self._thread is not None:
+            self._thread.join(timeout=timeout)
+
+    def enqueue(self, task: MemoryTask) -> None:
+        self._queue.put(task)
+
+    def drain(self, *, timeout: float = 30.0) -> None:
+        """等待队列中已有任务处理完毕（不含后续新任务）。"""
+        done = threading.Event()
+        sentinel = MemoryTask(
+            MemoryTaskKind.COMPRESS_IMPORTANT,
+            {"drain": True, "done": done},
+        )
+        self._queue.put(sentinel)
+        if not done.wait(timeout=timeout):
+            raise TimeoutError("记忆压缩队列等待超时")
+
+    def _thread_main(self) -> None:
+        import asyncio
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(self._run_loop())
+        finally:
+            loop.close()
+
+    async def _run_loop(self) -> None:
+        compressor = self._compressor_factory()
+        while True:
+            task = await _queue_get_async(self._queue)
+            if task is None:
+                return
+            if task.payload.get("drain"):
+                event = task.payload.get("done")
+                if isinstance(event, threading.Event):
+                    event.set()
+                continue
+            work = self._prepare(task)
+            if work is None:
+                continue
+            result = await _process_task(compressor, task, work, self._config)
+            if result is None:
+                continue
+            self._commit(task, work, result)
+
+    @property
+    def config(self) -> MemoryConfig:
+        return self._config
+
+
+async def _process_task(
+    compressor: MemoryCompressor,
+    task: MemoryTask,
+    work: CompressionWork,
+    config: MemoryConfig,
+) -> CompressionResult | None:
+    if task.kind == MemoryTaskKind.SHORT_TO_DATE:
+        return await _process_short_to_date(compressor, work)
+    if task.kind == MemoryTaskKind.DATE_TO_LONG:
+        return await _process_date_to_long(compressor, work)
+    if task.kind == MemoryTaskKind.COMPRESS_LONG:
+        return await _process_compress_long(compressor, work, config)
+    if task.kind == MemoryTaskKind.COMPRESS_IMPORTANT:
+        return await _process_compress_important(compressor, work, config)
+    return None
+
+
+async def _process_short_to_date(
+    compressor: MemoryCompressor,
+    work: CompressionWork,
+) -> ShortToDateResult | None:
+    if not isinstance(work, ShortToDateWork) or not work.batch:
+        return None
+    batch = list(work.batch)
+    entries, important_texts = await compressor.compress_to_date_entries(batch)
+    day_label = _day_label(batch[0].at)
+    return ShortToDateResult(
+        entries=tuple(entries),
+        important_texts=tuple(important_texts),
+        day_label=day_label,
+    )
+
+
+async def _process_date_to_long(
+    compressor: MemoryCompressor,
+    work: CompressionWork,
+) -> DateToLongResult | None:
+    if not isinstance(work, DateToLongWork):
+        return None
+    chunk = await compressor.merge_date_to_long_term(
+        work.day_label,
+        list(work.entries),
+    )
+    return DateToLongResult(chunk=chunk)
+
+
+async def _process_compress_long(
+    compressor: MemoryCompressor,
+    work: CompressionWork,
+    config: MemoryConfig,
+) -> CompressLongResult | None:
+    if not isinstance(work, CompressLongWork):
+        return None
+    if len(work.chunks) <= config.long_term_max_chunks:
+        return None
+    merged = await compressor.merge_long_term_chunks(list(work.chunks))
+    return CompressLongResult(chunks=tuple(merged))
+
+
+async def _process_compress_important(
+    compressor: MemoryCompressor,
+    work: CompressionWork,
+    config: MemoryConfig,
+) -> CompressImportantResult | None:
+    if not isinstance(work, CompressImportantWork):
+        return None
+    if len(work.entries) <= config.important_max_entries:
+        return None
+    merged = await compressor.reconcile_important(list(work.entries))
+    return CompressImportantResult(entries=tuple(merged))
+
+
+async def _queue_get_async(q: queue.Queue) -> MemoryTask | None:
+    import asyncio
+
+    while True:
+        try:
+            return q.get_nowait()
+        except queue.Empty:
+            await asyncio.sleep(0.05)
+
+
+def _day_label(at: datetime) -> str:
+    if at.tzinfo is None:
+        at = at.replace(tzinfo=timezone.utc)
+    return at.astimezone(timezone.utc).date().isoformat()
+
+
+def expire_old_date_days(
+    snapshot,
+    *,
+    config: MemoryConfig,
+    today: date | None = None,
+) -> list[str]:
+    """返回应归档到长期记忆的过期日期标签。"""
+    ref = today or datetime.now(timezone.utc).date()
+    cutoff = ref - timedelta(days=config.date_memory_days - 1)
+    expired: list[str] = []
+    for day in snapshot.date_days:
+        try:
+            day_date = date.fromisoformat(day.date)
+        except ValueError:
+            continue
+        if day_date < cutoff:
+            expired.append(day.date)
+    return expired

ai_agent/plan/__init__.py

@@ -0,0 +1,15 @@
+from ai_agent.plan.models import Plan, PlanRunResult, PlanStep
+from ai_agent.plan.parse import PlanParseError, parse_plan_text
+from ai_agent.plan.planner import PlanPlanner
+from ai_agent.plan.runner import PlanRunner, PlanStepFailedError
+
+__all__ = [
+    "Plan",
+    "PlanParseError",
+    "parse_plan_text",
+    "PlanPlanner",
+    "PlanRunResult",
+    "PlanRunner",
+    "PlanStep",
+    "PlanStepFailedError",
+]

ai_agent/plan/complete.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from ai_agent.context import ChatMessage, RunContext, RunPhase, RunPhaseKind
+from ai_agent.listener import (
+    AgentListener,
+    notify_output_delta,
+    notify_thinking_delta,
+)
+from ai_agent.llm import LLMClient, StreamKind
+
+
+async def complete_text(
+    llm: LLMClient,
+    *,
+    system_prompt: str,
+    user_content: str,
+    history: list[ChatMessage] | None = None,
+    listeners: Sequence[AgentListener] | None = None,
+    phase: RunPhase | None = None,
+    parse_from_answer_text_only: bool = False,
+) -> str:
+    """
+    无工具单次补全，收集流式文本。
+
+    Args:
+        llm: 与执行 Agent 共用的语言模型客户端
+        system_prompt: 系统提示
+        user_content: 本轮用户内容
+        history: 规划前对话历史（不含本轮 user）
+        parse_from_answer_text_only: 为 True 时返回值仅含 TEXT 流（思考流仍通知 listener，
+            但不参与返回，供规划 JSON 解析，避免思考草稿中的 JSON 覆盖正式回答）
+
+    Returns:
+        模型完整文本
+    """
+    messages = list(history or [])
+    messages.append(ChatMessage(role="user", content=user_content))
+    planning_phase = phase or RunPhase(kind=RunPhaseKind.PLANNING)
+    run = RunContext(
+        system_prompt=system_prompt,
+        messages=messages,
+        phase=planning_phase,
+    )
+    active_listeners = list(listeners or [])
+    all_parts: list[str] = []
+    answer_parts: list[str] = []
+    async for chunk in llm.stream(run, tools=None):
+        if chunk.kind == StreamKind.REASONING and chunk.delta:
+            if not parse_from_answer_text_only:
+                all_parts.append(chunk.delta)
+            await notify_thinking_delta(active_listeners, chunk.delta, run)
+        elif chunk.kind == StreamKind.TEXT and chunk.delta:
+            all_parts.append(chunk.delta)
+            answer_parts.append(chunk.delta)
+            await notify_output_delta(active_listeners, chunk.delta, run)
+    if parse_from_answer_text_only:
+        text = "".join(answer_parts).strip()
+    else:
+        text = "".join(all_parts).strip()
+    if not text and run.output:
+        return run.output.strip()
+    return text

ai_agent/plan/delivery.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import re
+
+from ai_agent.plan.models import PlanStep
+
+_SKILL_REF_PATTERN = re.compile(r"\b[\w-]+/[\w-]+(?:/[\w-]+)*\b")
+_ENABLE_TOOL_NAMES = frozenset({"enable_skill", "skill__enable_skill"})
+_DELIVERY_SKILL_ALIASES: tuple[tuple[str, str], ...] = (
+    ("skills/chat-search-answer", "skills/chat-search-answer"),
+    ("chat-search-answer", "skills/chat-search-answer"),
+    ("聊天搜索回答", "skills/chat-search-answer"),
+)
+
+_PLAN_DELIVERY_PRELOAD_NOTE = (
+    "改写技能已载入本步系统上下文，无需再调用 enable_skill，"
+    "直接按技能正文原则交付终稿。"
+)
+
+
+def delivery_skill_refs_for_step(step: PlanStep) -> tuple[str, ...]:
+    """
+    从计划步 objective / title 解析终稿改写 skill 引用。
+
+    匹配 ``{root_key}/{skill_id}`` 形式（如 ``skills/chat-search-answer``）。
+    """
+    text = f"{step.title}\n{step.objective}"
+    refs: list[str] = []
+    for match in _SKILL_REF_PATTERN.finditer(text):
+        candidate = match.group(0)
+        if candidate not in refs and "/" in candidate:
+            refs.append(candidate)
+    for needle, ref in _DELIVERY_SKILL_ALIASES:
+        if needle in text and ref not in refs:
+            refs.append(ref)
+    return tuple(refs)
+
+
+def plan_delivery_preload_note() -> str:
+    """终稿步已预载 skill 时拼入用户消息的说明。"""
+    return _PLAN_DELIVERY_PRELOAD_NOTE

ai_agent/plan/display.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from ai_agent.plan.models import Plan, PlanStep
+
+
+def format_plan_for_terminal(plan: Plan) -> str:
+    """
+    将计划格式化为终端可读的固定版式文本。
+
+    Args:
+        plan: 规划结果
+
+    Returns:
+        多行字符串，不含首尾空行
+    """
+    lines: list[str] = ["--- plan ---"]
+    if plan.summary and plan.summary.strip():
+        lines.append(f"摘要: {plan.summary.strip()}")
+    total = len(plan.steps)
+    lines.append(f"步骤（共 {total} 步）:")
+    for index, step in enumerate(plan.steps, start=1):
+        lines.extend(_format_step_lines(index, step))
+    return "\n".join(lines)
+
+
+def _format_step_lines(index: int, step: PlanStep) -> list[str]:
+    optional_tag = "（可选）" if step.optional else ""
+    head = f"  {index}. {step.id} · {step.title}{optional_tag}"
+    body: list[str] = [head]
+    objective = step.objective.strip()
+    if objective:
+        for line in objective.splitlines():
+            body.append(f"     目标: {line}")
+    tools_line = _format_tools_line(step)
+    if tools_line:
+        body.append(f"     {tools_line}")
+    return body
+
+
+def _format_tools_line(step: PlanStep) -> str:
+    parts: list[str] = []
+    if step.hint_tools:
+        parts.append(f"建议工具: {', '.join(step.hint_tools)}")
+    if step.required_tool:
+        parts.append(f"须用工具: {step.required_tool}")
+    return "；".join(parts)