precis-cli 0.1.3__py3-none-any.whl

app/shared/services/ai/agent/memory.py

@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Copyright 2026 Precis Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""@fileoverview Agent 记忆模块
+
+功能概述:
+- 管理 Agent 多轮对话的上下文
+- 按 token 预算截断历史，保留 system prompt 和最近轮次
+- 支持保存和恢复 checkpoint
+
+架构设计:
+- 使用纯字典列表表示消息，兼容 OpenAI messages 格式
+- token 估算委托给 utils.estimate_tokens 统一实现（中文字 1:1、英文单词/数字串/标点各计 1，另加固定开销）
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from app.shared.services.ai.agent.types import AgentTurn
+
+logger = logging.getLogger(__name__)
+
+# 默认最大 token 预算
+DEFAULT_MAX_TOKENS = 120000
+# 单条消息保留的最大字符数
+_MAX_MESSAGE_CHARS = 8000
+
+
+class AgentMemory:
+    """
+    @classdesc Agent 短期记忆
+
+    维护 system prompt + 用户初始任务 + 多轮对话历史。
+    """
+
+    def __init__(self, system_prompt: str, max_tokens: int = DEFAULT_MAX_TOKENS):
+        """
+        @methoddesc 初始化记忆
+
+        参数:
+            system_prompt: 系统提示词
+            max_tokens: 最大 token 预算
+        """
+        self.system_prompt = system_prompt
+        self.max_tokens = max_tokens
+        self._messages: list[dict[str, Any]] = []
+        self._turns: list[AgentTurn] = []
+        self._turn_count: int = 0
+        if system_prompt:
+            self._messages.append({"role": "system", "content": system_prompt})
+
+    def set_task(self, task_message: str) -> None:
+        """设置用户初始任务消息。"""
+        # 确保只有一条用户初始任务，避免重复
+        self._messages = [m for m in self._messages if not (m.get("role") == "user" and m.get("_is_task"))]
+        self._messages.append({"role": "user", "content": task_message, "_is_task": True})
+
+    def add_assistant_message(
+        self,
+        content: str | None,
+        tool_calls: list[dict[str, Any]] | None = None,
+    ) -> None:
+        """添加 assistant 消息。"""
+        msg: dict[str, Any] = {"role": "assistant"}
+        if content:
+            msg["content"] = self._truncate(content)
+        if tool_calls:
+            msg["tool_calls"] = tool_calls
+        self._messages.append(msg)
+
+    def add_tool_result(self, tool_call_id: str, name: str, result: dict[str, Any] | str) -> None:
+        """添加工具结果消息。"""
+        observation = result if isinstance(result, str) else self._dict_to_observation(result)
+        self._messages.append(
+            {
+                "role": "tool",
+                "tool_call_id": tool_call_id,
+                "name": name,
+                "content": self._truncate(observation),
+            }
+        )
+
+    def add_system_reminder(self, content: str) -> None:
+        """注入一条系统级提醒消息（用 user role，带 _is_reminder 标记）。
+
+        为什么不用 role="system"：get_messages 会跳过所有 role=="system" 的非首条消息，
+        导致 reminder 被丢弃。用 user role 可确保它进入 LLM 的上下文。
+
+        用途（P1-2 收敛引导）：
+        - 接近 max_iterations 时提醒 LLM 直接回复、停止调工具
+        - 检测到跨轮重复调用同一工具同参数时提醒 LLM 使用已有结果
+
+        参数:
+            content: 提醒文本
+        """
+        if not content:
+            return
+        self._messages.append({"role": "user", "content": self._truncate(content), "_is_reminder": True})
+
+    def add_turn(self, turn: AgentTurn) -> None:
+        """记录完整 turn。"""
+        self._turns.append(turn)
+        self._turn_count += 1
+
+    def get_messages(self) -> list[dict[str, Any]]:
+        """
+        @methoddesc 获取截断后的消息列表
+
+        保留 system prompt 和最近的用户/assistant/tool 消息，不超出 token 预算。
+        """
+        # 计算 system prompt 开销
+        system_tokens = self._estimate_tokens(self.system_prompt)
+        budget = self.max_tokens - system_tokens
+
+        # 从后往前累加，直到预算用完
+        selected: list[dict[str, Any]] = []
+        total = 0
+        # 跳过 system prompt（已在开头）
+        for msg in reversed(self._messages):
+            if msg.get("role") == "system":
+                continue
+            tokens = self._estimate_message_tokens(msg)
+            if total + tokens > budget and selected:
+                break
+            selected.append(msg)
+            total += tokens
+
+        selected.reverse()
+        result: list[dict[str, Any]] = []
+        if self.system_prompt:
+            result.append({"role": "system", "content": self.system_prompt})
+        result.extend(selected)
+        return result
+
+    def get_turns(self) -> list[AgentTurn]:
+        """获取所有 turn 记录。"""
+        return self._turns
+
+    def create_checkpoint(self) -> dict[str, Any]:
+        """创建当前记忆 checkpoint（含完整 messages，支持跨进程恢复）。"""
+        return {
+            "system_prompt": self.system_prompt,
+            "max_tokens": self.max_tokens,
+            "message_count": len(self._messages),
+            "turn_count": self._turn_count,
+            "messages": [dict(m) for m in self._messages],
+        }
+
+    @classmethod
+    def restore_from_checkpoint(cls, checkpoint: dict[str, Any]) -> AgentMemory:
+        """从 checkpoint 重建 memory 状态。
+
+        参数:
+            checkpoint: create_checkpoint 返回的字典
+
+        返回:
+            重建的 AgentMemory 实例
+        """
+        mem = cls(
+            system_prompt=checkpoint.get("system_prompt", ""),
+            max_tokens=checkpoint.get("max_tokens", DEFAULT_MAX_TOKENS),
+        )
+        mem._messages = [dict(m) for m in checkpoint.get("messages", [])]
+        mem._turn_count = checkpoint.get("turn_count", 0)
+        return mem
+
+    @staticmethod
+    def _truncate(text: str, max_chars: int = _MAX_MESSAGE_CHARS) -> str:
+        """截断超长文本（最终兜底）。"""
+        if len(text) <= max_chars:
+            return text
+        return text[:max_chars] + f"\n... [truncated {len(text) - max_chars} chars]"
+
+    @staticmethod
+    def _summarize_dict(data: dict[str, Any], max_field_chars: int = 1500) -> dict[str, Any]:
+        """对工具结果的 dict 做智能裁剪。
+
+        策略：小字段（< max_field_chars）完整保留；大字段（长列表/长字符串）截断，
+        保留关键摘要信息（如 errors 列表的前几条、success 标志等）。
+        相比直接对 JSON 文本硬截断，避免在字段中间被切断导致 agent 误解。
+        """
+        summarized: dict[str, Any] = {}
+        for key, value in data.items():
+            # 关键小字段始终完整保留
+            if key in {"success", "error", "message", "total_rules", "passed", "issues"}:
+                summarized[key] = value
+                continue
+            if isinstance(value, str):
+                summarized[key] = value if len(value) <= max_field_chars else value[:max_field_chars] + "...[truncated]"
+            elif isinstance(value, list):
+                # 列表：保留前 5 项 + 计数提示
+                if len(value) <= 5:
+                    summarized[key] = value
+                else:
+                    summarized[key] = value[:5] + [f"...[{len(value) - 5} more items]"]
+            elif isinstance(value, dict):
+                # 嵌套 dict 递归裁剪（降阈值避免嵌套膨胀）
+                summarized[key] = AgentMemory._summarize_dict(value, max_field_chars // 2)
+            else:
+                summarized[key] = value
+        return summarized
+
+    @staticmethod
+    def _dict_to_observation(data: dict[str, Any]) -> str:
+        """将字典转为 observation 文本（先智能裁剪大字段，避免 JSON 中途被硬截断）。"""
+        import json
+
+        summarized = AgentMemory._summarize_dict(data)
+        try:
+            return json.dumps(summarized, ensure_ascii=False, indent=2)
+        except (TypeError, ValueError):
+            return str(summarized)
+
+    @staticmethod
+    def _estimate_tokens(text: str) -> int:
+        """粗略估算 token 数。
+
+        委托给 utils.estimate_tokens 统一实现（中文/英文词/数字/标点分类加权），
+        避免项目内两套不一致的 token 估算系数。延迟导入防循环依赖。
+        """
+        if not text:
+            return 0
+        from app.shared.services.ai.utils import estimate_tokens
+
+        return estimate_tokens(text)
+
+    @classmethod
+    def _estimate_message_tokens(cls, msg: dict[str, Any]) -> int:
+        """估算单条消息 token 数。"""
+        content = msg.get("content") or ""
+        tokens = cls._estimate_tokens(content)
+        # tool_calls 也计入
+        tool_calls = msg.get("tool_calls") or []
+        for tc in tool_calls:
+            func = tc.get("function", {})
+            tokens += cls._estimate_tokens(func.get("name", ""))
+            args = func.get("arguments")
+            if isinstance(args, str):
+                tokens += cls._estimate_tokens(args)
+            elif isinstance(args, dict):
+                tokens += cls._estimate_tokens(str(args))
+        return tokens + 4  # 消息格式开销

app/shared/services/ai/agent/planner.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Copyright 2026 Precis Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""@fileoverview Agent 分块规划器
+
+功能概述:
+- 根据数据文件数量、列数自动决定分块策略
+- 输出 chunk 列表，每个 chunk 包含要处理的文件/列范围
+
+架构设计:
+- 简单启发式：文件数多按文件分，列数多按列分
+- 每个 chunk 独立生成配置，最后合并
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class Chunk:
+    """单个分块"""
+
+    chunk_id: str
+    file_paths: list[str] = field(default_factory=list)
+    table_names: list[str] = field(default_factory=list)
+    columns: list[str] = field(default_factory=list)
+    reason: str = ""
+
+
+@dataclass
+class ChunkPlan:
+    """分块计划"""
+
+    chunks: list[Chunk] = field(default_factory=list)
+    strategy: str = ""
+    reason: str = ""
+
+
+def build_chunk_plan(
+    profiling_data: list[dict],
+    file_paths: list[str],
+    max_columns_per_chunk: int = 20,
+    max_files_per_chunk: int = 5,
+) -> ChunkPlan:
+    """
+    @methoddesc 构建分块计划
+
+    参数:
+        profiling_data: 数据画像
+        file_paths: 所有文件路径
+        max_columns_per_chunk: 每个 chunk 最大列数
+        max_files_per_chunk: 每个 chunk 最大文件数
+
+    返回:
+        ChunkPlan 分块计划
+    """
+    total_files = len(profiling_data)
+    total_columns = sum(len(item.get("columns", [])) for item in profiling_data)
+
+    # 决定策略
+    if total_files > max_files_per_chunk:
+        strategy = "by_file"
+        chunks = _chunk_by_file(profiling_data, max_files_per_chunk)
+    elif total_columns > max_columns_per_chunk:
+        strategy = "by_column"
+        chunks = _chunk_by_column(profiling_data, max_columns_per_chunk)
+    else:
+        strategy = "single"
+        chunks = [
+            Chunk(
+                chunk_id="single",
+                file_paths=file_paths,
+                table_names=[item.get("table_name", "") for item in profiling_data],
+                reason="数据量较小，一次性处理",
+            )
+        ]
+
+    reason = f"共 {total_files} 个文件、{total_columns} 列，采用 {strategy} 策略"
+    return ChunkPlan(chunks=chunks, strategy=strategy, reason=reason)
+
+
+def _chunk_by_file(profiling_data: list[dict], max_files_per_chunk: int) -> list[Chunk]:
+    """按文件分块。"""
+    chunks: list[Chunk] = []
+    for i in range(0, len(profiling_data), max_files_per_chunk):
+        group = profiling_data[i : i + max_files_per_chunk]
+        chunk_id = f"files_{i // max_files_per_chunk + 1}"
+        chunks.append(
+            Chunk(
+                chunk_id=chunk_id,
+                file_paths=[item.get("path", "") for item in group],
+                table_names=[item.get("table_name", "") for item in group],
+                reason=f"处理第 {i + 1}-{min(i + max_files_per_chunk, len(profiling_data))} 个文件",
+            )
+        )
+    return chunks
+
+
+def _chunk_by_column(profiling_data: list[dict], max_columns_per_chunk: int) -> list[Chunk]:
+    """按列分块（优先保持同文件列在一起，但单表列数超过限制时拆表）。"""
+    chunks: list[Chunk] = []
+    chunk_index = 1
+
+    for item in profiling_data:
+        table_name = item.get("table_name", "")
+        path = item.get("path", "")
+        columns = item.get("columns", [])
+        col_names = [f"{table_name}.{col.get('name', '')}" for col in columns]
+
+        # 如果单表列数超过限制，拆分该表
+        if len(col_names) > max_columns_per_chunk:
+            for i in range(0, len(col_names), max_columns_per_chunk):
+                chunk = Chunk(
+                    chunk_id=f"cols_{chunk_index}",
+                    file_paths=[path],
+                    table_names=[table_name],
+                    columns=col_names[i : i + max_columns_per_chunk],
+                    reason=f"表 {table_name} 列数过多，拆分为第 {i // max_columns_per_chunk + 1} 组",
+                )
+                chunks.append(chunk)
+                chunk_index += 1
+            continue
+
+        # 否则尝试加入当前 chunk
+        if chunks and len(chunks[-1].columns) + len(col_names) <= max_columns_per_chunk:
+            current = chunks[-1]
+            if path and path not in current.file_paths:
+                current.file_paths.append(path)
+            if table_name and table_name not in current.table_names:
+                current.table_names.append(table_name)
+            current.columns.extend(col_names)
+            current.reason = f"按列分组，共 {len(current.columns)} 列"
+        else:
+            chunk = Chunk(
+                chunk_id=f"cols_{chunk_index}",
+                file_paths=[path],
+                table_names=[table_name],
+                columns=col_names,
+                reason=f"按列分组，共 {len(col_names)} 列",
+            )
+            chunks.append(chunk)
+            chunk_index += 1
+
+    return chunks
+
+
+def plan_to_dict(plan: ChunkPlan) -> dict[str, Any]:
+    """将 ChunkPlan 转为可 JSON 序列化的字典。"""
+    return {
+        "strategy": plan.strategy,
+        "reason": plan.reason,
+        "chunks": [
+            {
+                "chunk_id": c.chunk_id,
+                "file_paths": c.file_paths,
+                "table_names": c.table_names,
+                "columns": c.columns,
+                "reason": c.reason,
+            }
+            for c in plan.chunks
+        ],
+    }
+
+
+# ===== 按源(source)分块（迁移专用，不影响 generation 主路径）=====
+
+
+@dataclass
+class SourceChunk:
+    """单个来源分片"""
+
+    chunk_id: str
+    source_names: list[str] = field(default_factory=list)
+    source_indices: list[int] = field(default_factory=list)  # 指向 parsed_intents 原始列表
+    reason: str = ""
+
+
+@dataclass
+class SourceChunkPlan:
+    """来源分块计划"""
+
+    chunks: list[SourceChunk] = field(default_factory=list)
+    strategy: str = ""  # single | by_source
+    reason: str = ""
+
+
+def _estimate_intent_tokens(intent_text: str) -> int:
+    """保守估算意图文本 token 数。"""
+    return max(1, len(intent_text or "") // 3)
+
+
+def build_source_chunk_plan(
+    parsed_intents: list[dict],
+    max_sources_per_chunk: int = 5,
+    max_tokens_per_chunk: int = 8000,
+) -> SourceChunkPlan:
+    """
+    @methoddesc 构建按源分块计划
+
+    参数:
+        parsed_intents: 解析后的来源意图列表，每个元素含 name/intent
+        max_sources_per_chunk: 每个分片最大来源数
+        max_tokens_per_chunk: 每个分片最大估算 token 数
+
+    返回:
+        SourceChunkPlan 分块计划
+    """
+    n = len(parsed_intents)
+    if n == 0:
+        return SourceChunkPlan(chunks=[], strategy="single", reason="无来源，无需分片")
+
+    total_tokens = sum(_estimate_intent_tokens(it.get("intent", "")) for it in parsed_intents)
+    if n <= max_sources_per_chunk and total_tokens <= max_tokens_per_chunk:
+        return SourceChunkPlan(
+            chunks=[
+                SourceChunk(
+                    chunk_id="source_single",
+                    source_names=[it.get("name") or f"source_{idx + 1}" for idx, it in enumerate(parsed_intents)],
+                    source_indices=list(range(n)),
+                    reason=f"来源数 {n}、估算 {total_tokens} token，单次处理",
+                )
+            ],
+            strategy="single",
+            reason=f"共 {n} 个来源，单次处理",
+        )
+
+    chunks: list[SourceChunk] = []
+    current = SourceChunk(chunk_id="source_chunk_1")
+    current_tokens = 0
+    for idx, it in enumerate(parsed_intents):
+        it_tokens = _estimate_intent_tokens(it.get("intent", ""))
+        over_sources = len(current.source_indices) >= max_sources_per_chunk
+        over_tokens = current_tokens + it_tokens > max_tokens_per_chunk and current.source_indices
+        if over_sources or over_tokens:
+            current.reason = f"分片含 {len(current.source_indices)} 个来源，约 {current_tokens} token"
+            chunks.append(current)
+            current = SourceChunk(chunk_id=f"source_chunk_{len(chunks) + 1}")
+            current_tokens = 0
+        current.source_indices.append(idx)
+        current.source_names.append(it.get("name", f"source_{idx + 1}"))
+        current_tokens += it_tokens
+
+    if current.source_indices:
+        current.reason = f"分片含 {len(current.source_indices)} 个来源，约 {current_tokens} token"
+        chunks.append(current)
+
+    return SourceChunkPlan(
+        chunks=chunks,
+        strategy="by_source",
+        reason=f"共 {n} 个来源、约 {total_tokens} token，按源拆分为 {len(chunks)} 个分片",
+    )
+
+
+def source_plan_to_dict(plan: SourceChunkPlan) -> dict[str, Any]:
+    """将 SourceChunkPlan 转为可 JSON 序列化的字典。"""
+    return {
+        "strategy": plan.strategy,
+        "reason": plan.reason,
+        "chunks": [
+            {
+                "chunk_id": c.chunk_id,
+                "source_names": c.source_names,
+                "source_indices": c.source_indices,
+                "reason": c.reason,
+            }
+            for c in plan.chunks
+        ],
+    }