diff --git a/reme_ai/__init__.py b/reme_ai/__init__.py index a72cea87..39d39f5c 100644 --- a/reme_ai/__init__.py +++ b/reme_ai/__init__.py @@ -1,4 +1,5 @@ from reme_ai import retrieve +from reme_ai import summary from reme_ai import vector_store __version__ = "0.1.0" diff --git a/reme_ai/config/default.yaml b/reme_ai/config/default.yaml index 592f919e..26eafc6e 100644 --- a/reme_ai/config/default.yaml +++ b/reme_ai/config/default.yaml @@ -16,7 +16,7 @@ http: limit_concurrency: 64 flow: - retrieve_task_memory: + retrieve_task_memory_simple: flow_content: build_query_op >> recall_vector_store_op >> merge_memory_op description: "Retrieve the most relevant top_k memory experience from historical memory based on the query to help solve tasks better now" input_schema: @@ -24,10 +24,15 @@ flow: type: "str" description: "current query" required: true - top_k: - type: "int" - description: "top_k similar memories" - required: true + + summary_task_memory_simple: + flow_content: simple_summary_op >> update_vector_store_op + description: "Summarize trajectories or messages into memories" + input_schema: + trajectories: + type: "list" + description: "A list of conversation trajectory information, including message content and score. This field does not need to be filled in, the system will complete it automatically." + required: false # task_summarizer: simple_summary_op->update_vector_store_op # vector_store: vector_store_action_op diff --git a/reme_ai/retrieve/__init__.py b/reme_ai/retrieve/__init__.py index c5e9fa3a..53a36c1f 100644 --- a/reme_ai/retrieve/__init__.py +++ b/reme_ai/retrieve/__init__.py @@ -1 +1,2 @@ from reme_ai.retrieve import task +from reme_ai.retrieve import personal diff --git a/reme_ai/retrieve/task/build_query_op.py b/reme_ai/retrieve/task/build_query_op.py index ca15d2ea..e043a2eb 100644 --- a/reme_ai/retrieve/task/build_query_op.py +++ b/reme_ai/retrieve/task/build_query_op.py @@ -5,7 +5,7 @@ from loguru import logger @C.register_op() class BuildQueryOp(BaseLLMOp): - current_path: str = __file__ + file_path: str = __file__ def execute(self): if "query" in self.context: diff --git a/reme_ai/retrieve/task/rerank_memory_op.py b/reme_ai/retrieve/task/rerank_memory_op.py index 9b617fd1..45923232 100644 --- a/reme_ai/retrieve/task/rerank_memory_op.py +++ b/reme_ai/retrieve/task/rerank_memory_op.py @@ -15,7 +15,7 @@ class RerankMemoryOp(BaseLLMOp): """ Rerank and filter recalled experiences using LLM and score-based filtering """ - current_path: str = __file__ + file_path: str = __file__ def execute(self): """Execute rerank operation""" diff --git a/reme_ai/retrieve/task/rewrite_memory_op.py b/reme_ai/retrieve/task/rewrite_memory_op.py index b9c7cc9a..13e19363 100644 --- a/reme_ai/retrieve/task/rewrite_memory_op.py +++ b/reme_ai/retrieve/task/rewrite_memory_op.py @@ -15,7 +15,7 @@ class RewriteMemoryOp(BaseLLMOp): """ Generate and rewrite context messages from reranked experiences """ - current_path: str = __file__ + file_path: str = __file__ def execute(self): """Execute rewrite operation""" @@ -37,8 +37,7 @@ class RewriteMemoryOp(BaseLLMOp): # Store results in context self.context.response.answer = rewritten_memory - def _generate_context_message(self, query: str, messages: List[Message], memories: List[BaseMemory], - ) -> str: + def _generate_context_message(self, query: str, messages: List[Message], memories: List[BaseMemory]) -> str: """Generate context message from retrieved memories""" if not memories: return "" diff --git a/reme_ai/schema/__init__.py b/reme_ai/schema/__init__.py index e69de29b..7089e89b 100644 --- a/reme_ai/schema/__init__.py +++ b/reme_ai/schema/__init__.py @@ -0,0 +1 @@ +from flowllm.schema.message import Message, Trajectory, ToolCall \ No newline at end of file diff --git a/reme_ai/summary/__init__.py b/reme_ai/summary/__init__.py index e69de29b..1f47c200 100644 --- a/reme_ai/summary/__init__.py +++ b/reme_ai/summary/__init__.py @@ -0,0 +1,2 @@ +from reme_ai.summary import task +from reme_ai.summary import personal \ No newline at end of file diff --git a/reme_ai/summary/task/__init__.py b/reme_ai/summary/task/__init__.py index e69de29b..c6b143af 100644 --- a/reme_ai/summary/task/__init__.py +++ b/reme_ai/summary/task/__init__.py @@ -0,0 +1 @@ +from reme_ai.summary.task.simple_summary_op import SimpleSummaryOp \ No newline at end of file diff --git a/reme_ai/summary/task/comparative_extraction_op.py b/reme_ai/summary/task/comparative_extraction_op.py index 9f3f417d..0a7309c8 100644 --- a/reme_ai/summary/task/comparative_extraction_op.py +++ b/reme_ai/summary/task/comparative_extraction_op.py @@ -5,8 +5,8 @@ from loguru import logger from flowllm import C, BaseLLMOp from reme_ai.schema.memory import BaseMemory, TaskMemory -from reme_ai.schema.message import Message, Trajectory -from reme_ai.utils.memory_utils import merge_messages_content, parse_json_experience_response, get_trajectory_context +from reme_ai.schema import Message, Trajectory +from reme_ai.utils.op_utils import merge_messages_content, parse_json_experience_response, get_trajectory_context @C.register_op() diff --git a/reme_ai/summary/task/pdf_preprocess_op.py b/reme_ai/summary/task/pdf_preprocess_op.py deleted file mode 100644 index 9bf10014..00000000 --- a/reme_ai/summary/task/pdf_preprocess_op.py +++ /dev/null @@ -1,503 +0,0 @@ -#!/usr/bin/env python3 -""" -MinerU PDF 处理器 -返回 Markdown 内容和结构化的 content_list -""" - -import os -import json -import subprocess -import tempfile -import logging -from pathlib import Path -from typing import Dict, List, Any, Tuple, Optional, Union -import platform -import re - - -class MinerUPDFProcessor: - """ - 基于 MinerU 的 PDF 处理器 - 仿照 RAGAnything 的处理逻辑，但独立使用 MinerU - """ - - def __init__(self, log_level: str = "INFO"): - """ - 初始化处理器 - - Args: - log_level: 日志级别 ("DEBUG", "INFO", "WARNING", "ERROR") - """ - # 设置日志 - logging.basicConfig( - level=getattr(logging, log_level.upper()), - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' - ) - self.logger = logging.getLogger(__name__) - - # 检查 MinerU 安装 - if not self.check_mineru_installation(): - raise RuntimeError( - "MinerU 未正确安装。请使用以下命令安装：\n" - "pip install -U 'mineru[core]' 或 uv pip install -U 'mineru[core]'" - ) - - def check_mineru_installation(self) -> bool: - """检查 MinerU 是否正确安装""" - try: - subprocess_kwargs = { - "capture_output": True, - "text": True, - "check": True, - "encoding": "utf-8", - "errors": "ignore", - } - - # Windows 下隐藏控制台窗口 - if platform.system() == "Windows": - subprocess_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW - - result = subprocess.run(["mineru", "--version"], **subprocess_kwargs) - self.logger.debug(f"MinerU 版本: {result.stdout.strip()}") - return True - except (subprocess.CalledProcessError, FileNotFoundError): - return False - - def _run_mineru_command( - self, - input_path: Union[str, Path], - output_dir: Union[str, Path], - method: str = "auto", - lang: Optional[str] = None, - backend: str = "pipeline", - start_page: Optional[int] = None, - end_page: Optional[int] = None, - formula: bool = True, - table: bool = True, - device: Optional[str] = None, - source: str = "modelscope", - vlm_url: Optional[str] = None, - ) -> None: - """ - 运行 MinerU 命令行工具 - - Args: - input_path: 输入文件路径 - output_dir: 输出目录路径 - method: 解析方法 (auto, txt, ocr) - lang: 文档语言，用于 OCR 优化 - backend: 解析后端 - start_page: 起始页码 (0-based) - end_page: 结束页码 (0-based) - formula: 启用公式解析 - table: 启用表格解析 - device: 推理设备 - source: 模型来源 - vlm_url: VLM 服务器 URL（当 backend 为 vlm-sglang-client 时需要） - """ - cmd = [ - "mineru", - "-p", str(input_path), - "-o", str(output_dir), - "-m", method, - # "-b", backend, - # "--source", source, - ] - - # 添加可选参数 - if lang: - cmd.extend(["-l", lang]) - if start_page is not None: - cmd.extend(["-s", str(start_page)]) - if end_page is not None: - cmd.extend(["-e", str(end_page)]) - if not formula: - cmd.extend(["-f", "false"]) - if not table: - cmd.extend(["-t", "false"]) - if device: - cmd.extend(["-d", device]) - if vlm_url: - cmd.extend(["-u", vlm_url]) - - try: - subprocess_kwargs = { - "capture_output": True, - "text": True, - "check": True, - "encoding": "utf-8", - "errors": "ignore", - } - - # Windows 下隐藏控制台窗口 - if platform.system() == "Windows": - subprocess_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW - - self.logger.info(f"执行 MinerU 命令: {' '.join(cmd)}") - result = subprocess.run(cmd, **subprocess_kwargs) - - self.logger.info("MinerU 命令执行成功") - if result.stdout: - self.logger.debug(f"MinerU 输出: {result.stdout}") - - except subprocess.CalledProcessError as e: - self.logger.error(f"MinerU 命令执行错误: {e}") - if e.stderr: - self.logger.error(f"错误详情: {e.stderr}") - raise - except FileNotFoundError: - raise RuntimeError( - "mineru 命令未找到。请确保 MinerU 2.0 已正确安装：\n" - "pip install -U 'mineru[core]' 或 uv pip install -U 'mineru[core]'" - ) - - def _read_output_files( - self, - output_dir: Path, - file_stem: str, - method: str = "auto" - ) -> Tuple[List[Dict[str, Any]], str]: - """ - 读取 MinerU 生成的输出文件 - - Args: - output_dir: 输出目录 - file_stem: 文件名（不含扩展名） - method: 解析方法 - - Returns: - Tuple[List[Dict[str, Any]], str]: (content_list, markdown_content) - """ - # 查找生成的文件 - md_file = output_dir / f"{file_stem}.md" - json_file = output_dir / f"{file_stem}_content_list.json" - images_base_dir = output_dir - - # 检查子目录结构 - file_stem_subdir = output_dir / file_stem - if file_stem_subdir.exists(): - md_file = file_stem_subdir / method / f"{file_stem}.md" - json_file = file_stem_subdir / method / f"{file_stem}_content_list.json" - images_base_dir = file_stem_subdir / method - - # 读取 Markdown 内容 - md_content = "" - if md_file.exists(): - try: - with open(md_file, "r", encoding="utf-8") as f: - md_content = f.read() - self.logger.info(f"成功读取 Markdown 文件: {md_file}") - except Exception as e: - self.logger.warning(f"无法读取 Markdown 文件 {md_file}: {e}") - else: - self.logger.warning(f"Markdown 文件不存在: {md_file}") - - # 读取 JSON 内容列表 - content_list = [] - if json_file.exists(): - try: - with open(json_file, "r", encoding="utf-8") as f: - content_list = json.load(f) - - # 修复相对路径为绝对路径 - self.logger.info(f"修复图片路径，基础目录: {images_base_dir}") - for item in content_list: - if isinstance(item, dict): - for field_name in ["img_path", "table_img_path", "equation_img_path"]: - if field_name in item and item[field_name]: - img_path = item[field_name] - if not os.path.isabs(img_path): - absolute_img_path = (images_base_dir / img_path).resolve() - item[field_name] = str(absolute_img_path) - self.logger.debug(f"更新 {field_name}: {img_path} -> {item[field_name]}") - - self.logger.info(f"成功读取 JSON 文件: {json_file}, 包含 {len(content_list)} 个内容块") - - except Exception as e: - self.logger.warning(f"无法读取 JSON 文件 {json_file}: {e}") - else: - self.logger.warning(f"JSON 文件不存在: {json_file}") - - return content_list, md_content - - def process_pdf( - self, - pdf_path: Union[str, Path], - output_dir: Optional[Union[str, Path]] = None, - method: str = "auto", - lang: Optional[str] = None, - backend: str = "pipeline", - **kwargs - ) -> Tuple[List[Dict[str, Any]], str]: - """ - 处理 PDF 文件 - - Args: - pdf_path: PDF 文件路径 - output_dir: 输出目录路径（可选，默认在 PDF 文件同目录下创建） - method: 解析方法 ("auto", "txt", "ocr") - lang: 文档语言，用于 OCR 优化 (如 "ch", "en", "ja") - backend: 解析后端 ("pipeline", "vlm-transformers", "vlm-sglang-engine", "vlm-sglang-client") - **kwargs: 其他 MinerU 参数 - - Returns: - Tuple[List[Dict[str, Any]], str]: (content_list, markdown_content) - - Raises: - FileNotFoundError: PDF 文件不存在 - RuntimeError: MinerU 处理失败 - """ - # 转换为 Path 对象 - pdf_path = Path(pdf_path) - if not pdf_path.exists(): - raise FileNotFoundError(f"PDF 文件不存在: {pdf_path}") - - if not pdf_path.suffix.lower() == '.pdf': - raise ValueError(f"文件不是 PDF 格式: {pdf_path}") - - name_without_suffix = pdf_path.stem - - # 准备输出目录 - if output_dir: - base_output_dir = Path(output_dir) - else: - base_output_dir = pdf_path.parent / "mineru_output" - - base_output_dir.mkdir(parents=True, exist_ok=True) - - try: - # 运行 MinerU 命令 - self.logger.info(f"开始处理 PDF 文件: {pdf_path}") - - self._run_mineru_command( - input_path=pdf_path, - output_dir=base_output_dir, - method=method, - lang=lang, - backend=backend, - **kwargs - ) - - # 读取生成的输出文件 - backend_method = method - if backend.startswith("vlm-"): - backend_method = "vlm" - - content_list, markdown_content = self._read_output_files( - base_output_dir, name_without_suffix, method=backend_method - ) - - # 统计处理结果 - content_stats = {} - for item in content_list: - if isinstance(item, dict): - content_type = item.get("type", "unknown") - content_stats[content_type] = content_stats.get(content_type, 0) + 1 - - self.logger.info(f"PDF 处理完成! 提取了 {len(content_list)} 个内容块") - self.logger.info("内容类型统计:") - for content_type, count in content_stats.items(): - self.logger.info(f" - {content_type}: {count}") - - return content_list, markdown_content - - except Exception as e: - self.logger.error(f"处理 PDF 文件时出错: {str(e)}") - raise - - def save_results( - self, - content_list: List[Dict[str, Any]], - markdown_content: str, - output_path: Union[str, Path], - save_markdown: bool = True, - save_json: bool = True, - indent: int = 2 - ) -> Dict[str, Path]: - """ - 保存处理结果到文件 - - Args: - content_list: 内容列表 - markdown_content: Markdown 内容 - output_path: 输出路径（不含扩展名） - save_markdown: 是否保存 Markdown 文件 - save_json: 是否保存 JSON 文件 - indent: JSON 文件缩进 - - Returns: - Dict[str, Path]: 保存的文件路径字典 - """ - output_path = Path(output_path) - saved_files = {} - - try: - # 确保输出目录存在 - output_path.parent.mkdir(parents=True, exist_ok=True) - - # 保存 Markdown 文件 - if save_markdown and markdown_content: - md_path = output_path.with_suffix('.md') - with open(md_path, 'w', encoding='utf-8') as f: - f.write(markdown_content) - saved_files['markdown'] = md_path - self.logger.info(f"Markdown 文件已保存: {md_path}") - - # 保存 JSON 文件 - if save_json and content_list: - json_path = output_path.with_suffix('.json') - with open(json_path, 'w', encoding='utf-8') as f: - json.dump(content_list, f, indent=indent, ensure_ascii=False) - saved_files['json'] = json_path - self.logger.info(f"JSON 文件已保存: {json_path}") - - return saved_files - - except Exception as e: - self.logger.error(f"保存文件时出错: {e}") - raise - - -def chunk_pdf_content(content_list: List[Dict[str, Any]], max_length: int = 4000) -> List[str]: - """ - 将MinerU解析的content_list分割成指定长度的文本块 - - Args: - content_list: MinerU解析的内容列表 - max_length: 每个chunk的最大字符长度 - - Returns: - List[str]: 分块后的文本列表，每个文本都带有chunk标记 - """ - - def extract_text(item): - """提取单个item的文本""" - if item.get("type") == "text": - text = item.get("text", "").strip() - if not text: - return "" - # 如果是标题，添加#标记 - level = item.get("text_level", 0) - if level > 0: - return f"{'#' * min(level, 6)} {text}" - return text - - elif item.get("type") == "table": - parts = [] - if item.get("table_caption"): - parts.append("表格: " + " | ".join(item["table_caption"])) - if item.get("table_body"): - # 简单清理HTML标签 - table_text = re.sub(r'<[^>]+>', ' | ', item["table_body"]) - table_text = re.sub(r'\s+', ' ', table_text).strip() - parts.append(table_text) - return "\n".join(parts) if parts else "" - - elif item.get("type") == "image": - if item.get("image_caption"): - return "图片: " + " | ".join(item["image_caption"]) - return "" - - return "" - - # 提取所有文本 - all_text = "" - for item in content_list: - text = extract_text(item) - if text.strip(): - all_text += text + "\n" - - if not all_text.strip(): - return [] - - # 分割成chunks - chunks = [] - current_chunk = "" - - for line in all_text.split('\n'): - if len(current_chunk) + len(line) + 1 > max_length and current_chunk: - chunks.append(current_chunk.strip()) - current_chunk = line - else: - current_chunk += line + "\n" if current_chunk else line - - # 添加最后一个chunk - if current_chunk.strip(): - chunks.append(current_chunk.strip()) - - # 添加标记 - total_chunks = len(chunks) - marked_chunks = [] - for i, chunk in enumerate(chunks): - header = f"=== CHUNK {i + 1}/{total_chunks} ({len(chunk)}字符) ===\n" - marked_chunks.append(header + chunk) - - return marked_chunks - - -def main(): - """主函数 - 直接处理 PDF 文件""" - - # 配置输入和输出路径 - input_pdf_path = "/Users/dengjiaji/Downloads/bandaoti_research.pdf" # 修改为你的 PDF 文件路径 - output_dir = "/Users/dengjiaji/Downloads/bandaoti_research" # 修改为你的输出目录 - - try: - # 创建处理器 - processor = MinerUPDFProcessor(log_level="INFO") - - # 处理 PDF 文件（使用默认参数） - print(f"开始处理 PDF: {input_pdf_path}") - content_list, markdown_content = processor.process_pdf( - pdf_path=input_pdf_path, - output_dir=output_dir, - method="auto", # 自动选择最佳解析方法 - backend="pipeline" # 使用默认后端 - ) - - # 保存结果到输出目录 - # output_base = Path(output_dir) / Path(input_pdf_path).stem - # saved_files = processor.save_results( - # content_list=content_list, - # markdown_content=markdown_content, - # output_path=output_base - # ) - - # 显示结果 - print(f"\n✅ 处理完成!") - print(f"📄 提取内容块数量: {len(content_list)}") - print(f"📝 Markdown 内容长度: {len(markdown_content)} 字符") - print(f"\n💾 保存的文件:") - # for file_type, file_path in saved_files.items(): - # print(f" {file_type}: {file_path}") - - # 显示内容类型统计 - content_stats = {} - for item in content_list: - if isinstance(item, dict): - content_type = item.get("type", "unknown") - content_stats[content_type] = content_stats.get(content_type, 0) + 1 - - print(f"\n📊 内容类型统计:") - for content_type, count in content_stats.items(): - print(f" {content_type}: {count}") - - return content_list, markdown_content - - except FileNotFoundError as e: - print(f"❌ 文件未找到: {e}") - print("请检查 input_pdf_path 是否正确") - return None, None - except Exception as e: - print(f"❌ 处理错误: {e}") - return None, None - - -if __name__ == "__main__": - # main() - with open("/Users/dengjiaji/Downloads/bandaoti_research/bandaoti_research/auto/bandaoti_research_content_list.json", 'r', encoding='utf-8') as f: - content_list = json.load(f) - - # 生成chunks - chunks = chunk_pdf_content(content_list, max_length=10000) - print("len(chunks)", len(chunks)) diff --git a/reme_ai/summary/task/simple_summary_op.py b/reme_ai/summary/task/simple_summary_op.py index 0050cd18..bc07463a 100644 --- a/reme_ai/summary/task/simple_summary_op.py +++ b/reme_ai/summary/task/simple_summary_op.py @@ -1,24 +1,24 @@ import json from typing import List +from flowllm import C, BaseLLMOp from loguru import logger -from flowllm import C, BaseLLMOp +from reme_ai.schema import Message, Trajectory from reme_ai.schema.memory import BaseMemory, TaskMemory -from reme_ai.schema.message import Message, Trajectory -from reme_ai.utils.memory_utils import merge_messages_content +from reme_ai.utils.op_utils import merge_messages_content @C.register_op() class SimpleSummaryOp(BaseLLMOp): - current_path: str = __file__ + file_path: str = __file__ def summary_trajectory(self, trajectory: Trajectory) -> List[BaseMemory]: execution_process = merge_messages_content(trajectory.messages) success_score_threshold: float = self.op_params.get("success_score_threshold", 0.9) logger.info(f"success_score_threshold={success_score_threshold}") - execution_result = "success" if trajectory.score > success_score_threshold else "fail" + execution_result = "success" if trajectory.score >= success_score_threshold else "fail" summary_prompt = self.prompt_format(prompt_name="summary_prompt", execution_process=execution_process, execution_result=execution_result, @@ -26,22 +26,24 @@ class SimpleSummaryOp(BaseLLMOp): def parse_content(message: Message): content = message.content - experience_list = [] + memory_list = [] try: - content = content.split("```")[1].strip() + if "```" in content: + content = content.split("```")[1].strip() + if content.startswith("json"): content = content.strip("json") for exp_dict in json.loads(content): when_to_use = exp_dict.get("when_to_use", "").strip() - experience = exp_dict.get("experience", "").strip() - if when_to_use and experience: - experience_list.append(TaskMemory(workspace_id=self.context.get("workspace_id", ""), - when_to_use=when_to_use, - content=experience, - author=getattr(self.llm, 'model_name', 'system'))) + memory = exp_dict.get("memory", "").strip() + if when_to_use and memory: + memory_list.append(TaskMemory(workspace_id=self.context.get("workspace_id", ""), + when_to_use=when_to_use, + content=memory, + author=getattr(self.llm, "model_name", "system"))) - return experience_list + return memory_list except Exception as e: logger.exception(f"parse content failed!\n{content}") @@ -50,13 +52,16 @@ class SimpleSummaryOp(BaseLLMOp): return self.llm.chat(messages=[Message(content=summary_prompt)], callback_fn=parse_content) def execute(self): - trajectories: List[Trajectory] = self.context.get("trajectories", []) + trajectories: list = self.context.get("trajectories", []) + trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories] - experience_list = [] + memory_list: List[BaseMemory] = [] for trajectory in trajectories: - experiences = self.summary_trajectory(trajectory) - experience_list.extend(experiences) + memories = self.summary_trajectory(trajectory) + if memories: + memory_list.extend(memories) - self.context.summary_experiences = experience_list - for e in experience_list: - logger.info(f"add experience when_to_use={e.when_to_use}\ncontent={e.content}") \ No newline at end of file + self.context.response.answer = json.dumps([x.model_dump() for x in memory_list]) + self.context.memory_list = memory_list + for memory in memory_list: + logger.info(f"add memory: when_to_use={memory.when_to_use}\ncontent={memory.content}") diff --git a/reme_ai/summary/task/simple_summary_prompt.yaml b/reme_ai/summary/task/simple_summary_prompt.yaml index bb20a51e..563cbc78 100644 --- a/reme_ai/summary/task/simple_summary_prompt.yaml +++ b/reme_ai/summary/task/simple_summary_prompt.yaml @@ -7,11 +7,11 @@ summary_prompt: | # Task Reflect on what went well and what did not go well in the execution path based on the user's question. - The summarized experiences should be generalizable, providing assistance in solving similar types of problems in the future. + The summarized memory should be generalizable, providing assistance in solving similar types of problems in the future. They can include good suggestions or highlight areas to avoid and potential pitfalls. - The type of experience can be plain text or a piece of code that solves a specific problem. - If there are no experiences to summarize, output an empty list []. - For each experience, first state the scenario when it is applicable (when to use), then provide the experience itself, with a maximum of two experiences summarized. + The type of memory can be plain text or a piece of code that solves a specific problem. + If there are no memory to summarize, output an empty list []. + For each memory, first state the scenario when it is applicable (when to use), then provide the memory itself, with a maximum of two memories summarized. # Output Format {summary_example} @@ -21,11 +21,11 @@ summary_example: | [ { "when_to_use": "...", - "experience": "..." + "memory": "..." }, { "when_to_use": "...", - "experience": "..." + "memory": "..." } ] ``` \ No newline at end of file diff --git a/reme_ai/utils/miner_u_pdf_processor.py b/reme_ai/utils/miner_u_pdf_processor.py new file mode 100644 index 00000000..8b8317e2 --- /dev/null +++ b/reme_ai/utils/miner_u_pdf_processor.py @@ -0,0 +1,721 @@ +#!/usr/bin/env python3 +""" +MinerU PDF Processor + +A comprehensive PDF processing utility that leverages MinerU for extracting structured content +from PDF documents. Returns both Markdown content and structured content lists for further processing. + +This processor provides a high-level interface to MinerU's command-line tools, handling +file I/O, error management, and result parsing automatically. +""" + +import json +import logging +import os +import platform +import re +import subprocess +from pathlib import Path +from typing import Dict, List, Any, Tuple, Optional, Union + + +class MinerUPDFProcessor: + """ + MinerU-based PDF Processing Engine + + A robust PDF processor that wraps MinerU functionality to extract structured content + from PDF documents. Inspired by RAGAnything's processing logic but operates independently + with MinerU as the core engine. + + Features: + - Automatic MinerU installation validation + - Multiple parsing methods (auto, txt, ocr) + - Language-specific OCR optimization + - Structured content extraction with metadata + - Image path resolution and management + - Comprehensive error handling and logging + + Example: + processor = MinerUPDFProcessor(log_level="INFO") + content_list, markdown = processor.process_pdf("document.pdf") + """ + + def __init__(self, log_level: str = "INFO"): + """ + Initialize the PDF processor with logging configuration. + + Args: + log_level (str): Logging level for the processor. + Options: "DEBUG", "INFO", "WARNING", "ERROR" + + Raises: + RuntimeError: If MinerU is not properly installed or accessible + """ + # Configure logging system + logging.basicConfig( + level=getattr(logging, log_level.upper()), + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + ) + self.logger = logging.getLogger(__name__) + + # Validate MinerU installation before proceeding + if not self.check_mineru_installation(): + raise RuntimeError( + "MinerU is not properly installed. Please install using:\n" + "pip install -U 'mineru[core]' or uv pip install -U 'mineru[core]'" + ) + + @classmethod + def create_with_defaults(cls, log_level: str = "INFO") -> "MinerUPDFProcessor": + """ + Create a MinerUPDFProcessor instance with default settings. + + Convenience method for quick instantiation with standard configuration. + + Args: + log_level (str): Logging level (default: "INFO") + + Returns: + MinerUPDFProcessor: Configured processor instance + """ + return cls(log_level=log_level) + + def check_mineru_installation(self) -> bool: + """ + Verify that MinerU is properly installed and accessible. + + Attempts to run the MinerU command-line tool to check its availability + and version information. + + Returns: + bool: True if MinerU is properly installed, False otherwise + """ + try: + # Configure subprocess parameters for cross-platform compatibility + subprocess_kwargs = { + "capture_output": True, + "text": True, + "check": True, + "encoding": "utf-8", + "errors": "ignore", + } + + # Hide console window on Windows systems + if platform.system() == "Windows": + subprocess_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW + + # Execute version check command + result = subprocess.run(["mineru", "--version"], **subprocess_kwargs) + self.logger.debug(f"MinerU version detected: {result.stdout.strip()}") + return True + except (subprocess.CalledProcessError, FileNotFoundError): + return False + + def _run_mineru_command( + self, + input_path: Union[str, Path], + output_dir: Union[str, Path], + method: str = "auto", + lang: Optional[str] = None, + backend: str = "pipeline", + start_page: Optional[int] = None, + end_page: Optional[int] = None, + formula: bool = True, + table: bool = True, + device: Optional[str] = None, + source: str = "modelscope", + vlm_url: Optional[str] = None, + ) -> None: + """ + Execute MinerU command-line tool with specified parameters. + + This method constructs and executes the MinerU command with all provided + options, handling cross-platform subprocess execution and error management. + + Args: + input_path (Union[str, Path]): Path to the input PDF file + output_dir (Union[str, Path]): Directory path for output files + method (str): Parsing method - "auto", "txt", or "ocr" + lang (Optional[str]): Document language for OCR optimization (e.g., "en", "ch", "ja") + backend (str): Processing backend to use + start_page (Optional[int]): Starting page number (0-based indexing) + end_page (Optional[int]): Ending page number (0-based indexing) + formula (bool): Enable mathematical formula parsing + table (bool): Enable table structure parsing + device (Optional[str]): Computing device for inference (e.g., "cuda", "cpu") + source (str): Model source repository + vlm_url (Optional[str]): VLM server URL (required for vlm-sglang-client backend) + + Raises: + subprocess.CalledProcessError: If MinerU command execution fails + FileNotFoundError: If MinerU executable is not found + RuntimeError: If MinerU is not properly installed + """ + # Build base command with required parameters + cmd = [ + "mineru", + "-p", str(input_path), + "-o", str(output_dir), + "-m", method, + # Note: backend and source parameters are commented out as they may not be + # available in all MinerU versions or configurations + # "-b", backend, + # "--source", source, + ] + + # Add optional parameters if specified + if lang: + cmd.extend(["-l", lang]) + if start_page is not None: + cmd.extend(["-s", str(start_page)]) + if end_page is not None: + cmd.extend(["-e", str(end_page)]) + if not formula: + cmd.extend(["-f", "false"]) + if not table: + cmd.extend(["-t", "false"]) + if device: + cmd.extend(["-d", device]) + if vlm_url: + cmd.extend(["-u", vlm_url]) + + try: + # Configure subprocess execution parameters + subprocess_kwargs = { + "capture_output": True, + "text": True, + "check": True, + "encoding": "utf-8", + "errors": "ignore", + } + + # Hide console window on Windows systems + if platform.system() == "Windows": + subprocess_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW + + self.logger.info(f"Executing MinerU command: {' '.join(cmd)}") + result = subprocess.run(cmd, **subprocess_kwargs) + + self.logger.info("MinerU command executed successfully") + if result.stdout: + self.logger.debug(f"MinerU output: {result.stdout}") + + except subprocess.CalledProcessError as e: + self.logger.error(f"MinerU command execution failed: {e}") + if e.stderr: + self.logger.error(f"Error details: {e.stderr}") + raise + except FileNotFoundError: + raise RuntimeError( + "MinerU command not found. Please ensure MinerU 2.0 is properly installed:\n" + "pip install -U 'mineru[core]' or uv pip install -U 'mineru[core]'" + ) + + def _read_output_files( + self, + output_dir: Path, + file_stem: str, + method: str = "auto" + ) -> Tuple[List[Dict[str, Any]], str]: + """ + Read and parse MinerU-generated output files. + + This method locates and reads the Markdown and JSON files generated by MinerU, + handling different directory structures and resolving image paths to absolute paths. + + Args: + output_dir (Path): Directory containing the MinerU output files + file_stem (str): Base filename without extension + method (str): Parsing method used ("auto", "txt", "ocr", "vlm") + + Returns: + Tuple[List[Dict[str, Any]], str]: A tuple containing: + - content_list: Structured content list with metadata + - markdown_content: Raw markdown text content + """ + # Locate generated output files - handle both flat and nested directory structures + md_file = output_dir / f"{file_stem}.md" + json_file = output_dir / f"{file_stem}_content_list.json" + images_base_dir = output_dir + + # Check for nested subdirectory structure (common with newer MinerU versions) + file_stem_subdir = output_dir / file_stem + if file_stem_subdir.exists(): + md_file = file_stem_subdir / method / f"{file_stem}.md" + json_file = file_stem_subdir / method / f"{file_stem}_content_list.json" + images_base_dir = file_stem_subdir / method + + # Read Markdown content + md_content = "" + if md_file.exists(): + try: + with open(md_file, "r", encoding="utf-8") as f: + md_content = f.read() + self.logger.info(f"Successfully read Markdown file: {md_file}") + except Exception as e: + self.logger.warning(f"Failed to read Markdown file {md_file}: {e}") + else: + self.logger.warning(f"Markdown file not found: {md_file}") + + # Read structured content list from JSON + content_list = [] + if json_file.exists(): + try: + with open(json_file, "r", encoding="utf-8") as f: + content_list = json.load(f) + + # Convert relative image paths to absolute paths for proper access + self.logger.info(f"Resolving image paths relative to: {images_base_dir}") + for item in content_list: + if isinstance(item, dict): + # Process various image path fields that may be present + for field_name in ["img_path", "table_img_path", "equation_img_path"]: + if field_name in item and item[field_name]: + img_path = item[field_name] + if not os.path.isabs(img_path): + absolute_img_path = (images_base_dir / img_path).resolve() + item[field_name] = str(absolute_img_path) + self.logger.debug(f"Updated {field_name}: {img_path} -> {item[field_name]}") + + self.logger.info(f"Successfully read JSON file: {json_file}, containing {len(content_list)} content blocks") + + except Exception as e: + self.logger.warning(f"Failed to read JSON file {json_file}: {e}") + else: + self.logger.warning(f"JSON file not found: {json_file}") + + return content_list, md_content + + def process_pdf( + self, + pdf_path: Union[str, Path], + output_dir: Optional[Union[str, Path]] = None, + method: str = "auto", + lang: Optional[str] = None, + backend: str = "pipeline", + **kwargs + ) -> Tuple[List[Dict[str, Any]], str]: + """ + Process a PDF file and extract structured content using MinerU. + + This is the main entry point for PDF processing. It validates input, + executes MinerU processing, and returns both structured content and markdown. + + Args: + pdf_path (Union[str, Path]): Path to the input PDF file + output_dir (Optional[Union[str, Path]]): Output directory path. + If None, creates 'mineru_output' in PDF's directory + method (str): Parsing method - "auto" (recommended), "txt", or "ocr" + lang (Optional[str]): Document language for OCR optimization + (e.g., "ch" for Chinese, "en" for English, "ja" for Japanese) + backend (str): Processing backend - "pipeline", "vlm-transformers", + "vlm-sglang-engine", or "vlm-sglang-client" + **kwargs: Additional MinerU parameters (start_page, end_page, formula, table, etc.) + + Returns: + Tuple[List[Dict[str, Any]], str]: A tuple containing: + - content_list: Structured list of content blocks with metadata + - markdown_content: Complete document in Markdown format + + Raises: + FileNotFoundError: If the specified PDF file does not exist + ValueError: If the file is not a valid PDF format + RuntimeError: If MinerU processing fails or encounters errors + """ + # Convert to Path object and validate input + pdf_path = Path(pdf_path) + if not pdf_path.exists(): + raise FileNotFoundError(f"PDF file does not exist: {pdf_path}") + + if not pdf_path.suffix.lower() == '.pdf': + raise ValueError(f"File is not a PDF format: {pdf_path}") + + name_without_suffix = pdf_path.stem + + # Prepare output directory + if output_dir: + base_output_dir = Path(output_dir) + else: + base_output_dir = pdf_path.parent / "mineru_output" + + base_output_dir.mkdir(parents=True, exist_ok=True) + + try: + # Execute MinerU processing + self.logger.info(f"Starting PDF processing: {pdf_path}") + + self._run_mineru_command( + input_path=pdf_path, + output_dir=base_output_dir, + method=method, + lang=lang, + backend=backend, + **kwargs + ) + + # Read generated output files + backend_method = method + if backend.startswith("vlm-"): + backend_method = "vlm" + + content_list, markdown_content = self._read_output_files( + base_output_dir, name_without_suffix, method=backend_method + ) + + # Generate processing statistics + content_stats = {} + for item in content_list: + if isinstance(item, dict): + content_type = item.get("type", "unknown") + content_stats[content_type] = content_stats.get(content_type, 0) + 1 + + self.logger.info(f"PDF processing completed! Extracted {len(content_list)} content blocks") + self.logger.info("Content type statistics:") + for content_type, count in content_stats.items(): + self.logger.info(f" - {content_type}: {count}") + + return content_list, markdown_content + + except Exception as e: + self.logger.error(f"Error occurred during PDF processing: {str(e)}") + raise + + def save_results( + self, + content_list: List[Dict[str, Any]], + markdown_content: str, + output_path: Union[str, Path], + save_markdown: bool = True, + save_json: bool = True, + indent: int = 2 + ) -> Dict[str, Path]: + """ + Save processing results to files. + + Saves the extracted content in both JSON (structured) and Markdown (text) formats + for different use cases and downstream processing needs. + + Args: + content_list (List[Dict[str, Any]]): Structured content list with metadata + markdown_content (str): Complete document in Markdown format + output_path (Union[str, Path]): Output file path (without extension) + save_markdown (bool): Whether to save Markdown file + save_json (bool): Whether to save JSON file with structured content + indent (int): JSON file indentation for readability + + Returns: + Dict[str, Path]: Dictionary mapping file types to their saved paths + Keys: 'markdown', 'json' (if respective files were saved) + + Raises: + Exception: If file writing operations fail + """ + output_path = Path(output_path) + saved_files = {} + + try: + # Ensure output directory exists + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Save Markdown file + if save_markdown and markdown_content: + md_path = output_path.with_suffix('.md') + with open(md_path, 'w', encoding='utf-8') as f: + f.write(markdown_content) + saved_files['markdown'] = md_path + self.logger.info(f"Markdown file saved: {md_path}") + + # Save JSON file with structured content + if save_json and content_list: + json_path = output_path.with_suffix('.json') + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(content_list, f, indent=indent, ensure_ascii=False) + saved_files['json'] = json_path + self.logger.info(f"JSON file saved: {json_path}") + + return saved_files + + except Exception as e: + self.logger.error(f"Error occurred while saving files: {e}") + raise + + def get_content_statistics(self, content_list: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Generate detailed statistics about the processed content. + + Analyzes the content list to provide insights into document structure, + content types, and processing results. + + Args: + content_list (List[Dict[str, Any]]): Structured content list from MinerU + + Returns: + Dict[str, Any]: Dictionary containing various statistics: + - total_blocks: Total number of content blocks + - content_types: Count of each content type + - text_stats: Text-specific statistics (characters, words, etc.) + - image_count: Number of images found + - table_count: Number of tables found + """ + stats = { + "total_blocks": len(content_list), + "content_types": {}, + "text_stats": {"total_characters": 0, "total_words": 0, "title_levels": {}}, + "image_count": 0, + "table_count": 0, + "has_formulas": False + } + + for item in content_list: + if not isinstance(item, dict): + continue + + content_type = item.get("type", "unknown") + stats["content_types"][content_type] = stats["content_types"].get(content_type, 0) + 1 + + if content_type == "text": + text = item.get("text", "") + stats["text_stats"]["total_characters"] += len(text) + stats["text_stats"]["total_words"] += len(text.split()) + + level = item.get("text_level", 0) + if level > 0: + stats["text_stats"]["title_levels"][level] = stats["text_stats"]["title_levels"].get(level, 0) + 1 + + elif content_type == "image": + stats["image_count"] += 1 + + elif content_type == "table": + stats["table_count"] += 1 + + elif content_type == "formula": + stats["has_formulas"] = True + + return stats + + def validate_output_quality(self, content_list: List[Dict[str, Any]], markdown_content: str) -> Dict[str, Any]: + """ + Validate the quality and completeness of the processing output. + + Performs various checks to ensure the processed content meets quality standards + and provides warnings or suggestions for improvement. + + Args: + content_list (List[Dict[str, Any]]): Structured content list + markdown_content (str): Markdown content string + + Returns: + Dict[str, Any]: Validation results containing: + - is_valid: Overall validation status + - warnings: List of warning messages + - suggestions: List of improvement suggestions + - quality_score: Numeric quality score (0-100) + """ + validation = { + "is_valid": True, + "warnings": [], + "suggestions": [], + "quality_score": 100 + } + + # Check if content was extracted + if not content_list and not markdown_content.strip(): + validation["is_valid"] = False + validation["warnings"].append("No content was extracted from the PDF") + validation["quality_score"] = 0 + return validation + + # Check content diversity + stats = self.get_content_statistics(content_list) + if stats["total_blocks"] < 5: + validation["warnings"].append("Very few content blocks extracted - document may be complex or image-heavy") + validation["quality_score"] -= 20 + + # Check text content ratio + text_blocks = stats["content_types"].get("text", 0) + if text_blocks == 0: + validation["warnings"].append("No text blocks found - consider using OCR method for image-based PDFs") + validation["quality_score"] -= 30 + elif text_blocks / stats["total_blocks"] < 0.3: + validation["suggestions"].append("Low text content ratio - document may benefit from OCR processing") + validation["quality_score"] -= 10 + + # Check for images without processing + if stats["image_count"] > 0 and stats["content_types"].get("text", 0) == 0: + validation["suggestions"].append("Images detected but no text extracted - consider using VLM backend for image analysis") + + # Check markdown length vs content blocks + if len(markdown_content.strip()) < 100 and stats["total_blocks"] > 10: + validation["warnings"].append("Markdown content seems unusually short for the number of content blocks") + validation["quality_score"] -= 15 + + return validation + + +def chunk_pdf_content(content_list: List[Dict[str, Any]], max_length: int = 4000) -> List[str]: + """ + Split MinerU-parsed content list into text chunks of specified length. + + This utility function converts structured content from MinerU into manageable + text chunks suitable for downstream processing like embedding generation or + language model input. + + Args: + content_list (List[Dict[str, Any]]): MinerU-parsed structured content list + max_length (int): Maximum character length per chunk (default: 4000) + + Returns: + List[str]: List of text chunks, each prefixed with chunk metadata + including chunk number, total chunks, and character count + """ + + def extract_text(item: Dict[str, Any]) -> str: + """ + Extract text content from a single content item. + + Handles different content types (text, table, image) and formats them + appropriately for text-based processing. + + Args: + item (Dict[str, Any]): Single content item from MinerU output + + Returns: + str: Extracted and formatted text content + """ + if item.get("type") == "text": + text = item.get("text", "").strip() + if not text: + return "" + # Add markdown header formatting for titles + level = item.get("text_level", 0) + if level > 0: + return f"{'#' * min(level, 6)} {text}" + return text + + elif item.get("type") == "table": + parts = [] + if item.get("table_caption"): + parts.append("Table: " + " | ".join(item["table_caption"])) + if item.get("table_body"): + # Simple HTML tag cleanup and formatting + table_text = re.sub(r'<[^>]+>', ' | ', item["table_body"]) + table_text = re.sub(r'\s+', ' ', table_text).strip() + parts.append(table_text) + return "\n".join(parts) if parts else "" + + elif item.get("type") == "image": + if item.get("image_caption"): + return "Image: " + " | ".join(item["image_caption"]) + return "" + + return "" + + # Extract all text content from the structured list + all_text = "" + for item in content_list: + text = extract_text(item) + if text.strip(): + all_text += text + "\n" + + if not all_text.strip(): + return [] + + # Split into chunks based on max_length + chunks = [] + current_chunk = "" + + for line in all_text.split('\n'): + # Check if adding this line would exceed max_length + if len(current_chunk) + len(line) + 1 > max_length and current_chunk: + chunks.append(current_chunk.strip()) + current_chunk = line + else: + current_chunk += line + "\n" if current_chunk else line + + # Add the final chunk if it contains content + if current_chunk.strip(): + chunks.append(current_chunk.strip()) + + # Add chunk metadata headers + total_chunks = len(chunks) + marked_chunks = [] + for i, chunk in enumerate(chunks): + header = f"=== CHUNK {i + 1}/{total_chunks} ({len(chunk)} characters) ===\n" + marked_chunks.append(header + chunk) + + return marked_chunks + + +# Example usage and demonstration +if __name__ == "__main__": + """ + Example usage of the MinerUPDFProcessor class. + + This example demonstrates the basic workflow for processing a PDF file + and working with the extracted content. + """ + import sys + + # Example usage + def example_usage(): + """Demonstrate basic PDF processing workflow.""" + try: + # Initialize processor + processor = MinerUPDFProcessor.create_with_defaults(log_level="INFO") + + # Example PDF path (replace with actual PDF file) + pdf_path = "example_document.pdf" + + if not Path(pdf_path).exists(): + print(f"Example PDF file not found: {pdf_path}") + print("Please provide a valid PDF file path to test the processor.") + return + + # Process PDF with different methods + print("Processing PDF with auto method...") + content_list, markdown_content = processor.process_pdf( + pdf_path=pdf_path, + method="auto", + lang="en" # Specify language for better OCR results + ) + + # Generate statistics + stats = processor.get_content_statistics(content_list) + print(f"Processing Statistics:") + print(f" Total blocks: {stats['total_blocks']}") + print(f" Content types: {stats['content_types']}") + print(f" Text characters: {stats['text_stats']['total_characters']}") + print(f" Text words: {stats['text_stats']['total_words']}") + + # Validate output quality + validation = processor.validate_output_quality(content_list, markdown_content) + print(f"Quality Score: {validation['quality_score']}/100") + if validation['warnings']: + print("Warnings:", validation['warnings']) + if validation['suggestions']: + print("Suggestions:", validation['suggestions']) + + # Save results + output_path = Path(pdf_path).stem + "_processed" + saved_files = processor.save_results( + content_list=content_list, + markdown_content=markdown_content, + output_path=output_path + ) + print(f"Results saved to: {saved_files}") + + # Create text chunks for downstream processing + chunks = chunk_pdf_content(content_list, max_length=2000) + print(f"Created {len(chunks)} text chunks") + + # Display first chunk as example + if chunks: + print("First chunk preview:") + print(chunks[0][:200] + "..." if len(chunks[0]) > 200 else chunks[0]) + + except Exception as e: + print(f"Error during processing: {e}") + sys.exit(1) + + # Run example if script is executed directly + example_usage() diff --git a/reme_ai/utils/op_utils.py b/reme_ai/utils/op_utils.py index 6a43c78a..50bfc61d 100644 --- a/reme_ai/utils/op_utils.py +++ b/reme_ai/utils/op_utils.py @@ -3,9 +3,13 @@ import re from typing import List from flowllm.schema.message import Message, Trajectory +from flowllm.utils.llm_utils import merge_messages_content as merge_messages_content_flowllm from loguru import logger +def merge_messages_content(messages: List[Message | dict]) -> str: + return merge_messages_content_flowllm(messages) + def parse_json_experience_response(response: str) -> List[dict]: """Parse JSON formatted experience response""" try: diff --git a/reme_ai/vector_store/__init__.py b/reme_ai/vector_store/__init__.py index d676d89a..a7bef7a7 100644 --- a/reme_ai/vector_store/__init__.py +++ b/reme_ai/vector_store/__init__.py @@ -1 +1,3 @@ -from reme_ai.vector_store.recall_vector_store_op import RecallVectorStoreOp \ No newline at end of file +from reme_ai.vector_store.recall_vector_store_op import RecallVectorStoreOp +from reme_ai.vector_store.update_vector_store_op import UpdateVectorStoreOp +from reme_ai.vector_store.vector_store_action_op import VectorStoreActionOp diff --git a/reme_ai/vector_store/recall_vector_store_op.py b/reme_ai/vector_store/recall_vector_store_op.py index 53cb8163..d0e89e0b 100644 --- a/reme_ai/vector_store/recall_vector_store_op.py +++ b/reme_ai/vector_store/recall_vector_store_op.py @@ -12,10 +12,11 @@ class RecallVectorStoreOp(BaseLLMOp): def execute(self): recall_key: str = self.op_params.get("recall_key", "query") + top_k: int = self.op_params.get("top_k", 3) + query: str = self.context[recall_key] assert query, "query should be not empty!" - top_k: int = self.context.top_k workspace_id: str = self.context.workspace_id nodes: List[VectorNode] = self.vector_store.search(query=query, workspace_id=workspace_id, top_k=top_k) diff --git a/reme_ai/vector_store/update_vector_store_op.py b/reme_ai/vector_store/update_vector_store_op.py index dae56a2a..bc1d0074 100644 --- a/reme_ai/vector_store/update_vector_store_op.py +++ b/reme_ai/vector_store/update_vector_store_op.py @@ -1,28 +1,32 @@ import json from typing import List +from flowllm import C, BaseLLMOp +from flowllm.schema.vector_node import VectorNode from loguru import logger -from experiencemaker.op import OP_REGISTRY -from experiencemaker.op.base_op import BaseOp -from experiencemaker.schema.experience import BaseExperience -from experiencemaker.schema.request import BaseRequest -from experiencemaker.schema.vector_node import VectorNode +from reme_ai.schema.memory import BaseMemory -@OP_REGISTRY.register() -class UpdateVectorStoreOp(BaseOp): +@C.register_op() +class UpdateVectorStoreOp(BaseLLMOp): def execute(self): - request: BaseRequest = self.context.request + workspace_id: str = self.context.workspace_id - experience_ids: List[str] | None = self.context.response.deleted_experience_ids - if experience_ids: - self.vector_store.delete(node_ids=experience_ids, workspace_id=request.workspace_id) - logger.info(f"delete experience_ids={json.dumps(experience_ids, indent=2)}") + deleted_memory_ids: List[str] = self.context.get("deleted_memory_ids", []) + if deleted_memory_ids: + self.vector_store.delete(node_ids=deleted_memory_ids, workspace_id=workspace_id) + logger.info(f"delete memory_ids={json.dumps(deleted_memory_ids, indent=2)}") - insert_experience_list: List[BaseExperience] | None = self.context.response.experience_list - if insert_experience_list: - insert_nodes: List[VectorNode] = [x.to_vector_node() for x in insert_experience_list] - self.vector_store.insert(nodes=insert_nodes, workspace_id=request.workspace_id) + insert_memory_list: List[BaseMemory] | None = self.context.get("memory_list", []) + if insert_memory_list: + insert_nodes: List[VectorNode] = [x.to_vector_node() for x in insert_memory_list] + self.vector_store.insert(nodes=insert_nodes, workspace_id=workspace_id) logger.info(f"insert insert_node.size={len(insert_nodes)}") + + # Store results in context + self.context.update_result = { + "deleted_count": len(deleted_memory_ids) if deleted_memory_ids else 0, + "inserted_count": len(insert_memory_list) if insert_memory_list else 0 + } diff --git a/reme_ai/vector_store/vector_store_action_op.py b/reme_ai/vector_store/vector_store_action_op.py index fe0d897b..975f1add 100644 --- a/reme_ai/vector_store/vector_store_action_op.py +++ b/reme_ai/vector_store/vector_store_action_op.py @@ -1,58 +1,67 @@ -from experiencemaker.op import OP_REGISTRY -from experiencemaker.op.base_op import BaseOp -from experiencemaker.schema.experience import vector_node_to_experience, dict_to_experience, BaseExperience -from experiencemaker.schema.request import VectorStoreRequest -from experiencemaker.schema.response import VectorStoreResponse -from experiencemaker.schema.vector_node import VectorNode +from flowllm import C, BaseLLMOp +from flowllm.schema.vector_node import VectorNode + +from reme_ai.schema.memory import vector_node_to_memory, dict_to_experience, BaseMemory -@OP_REGISTRY.register() -class VectorStoreActionOp(BaseOp): +@C.register_op() +class VectorStoreActionOp(BaseLLMOp): def execute(self): - request: VectorStoreRequest = self.context.request - response: VectorStoreResponse = self.context.response + workspace_id: str = self.context.workspace_id + action: str = self.context.action - if request.action == "copy": - result = self.vector_store.copy_workspace(src_workspace_id=request.src_workspace_id, - dest_workspace_id=request.workspace_id) + if action == "copy": + src_workspace_id: str = self.context.src_workspace_id + result = self.vector_store.copy_workspace(src_workspace_id=src_workspace_id, + dest_workspace_id=workspace_id) - elif request.action == "delete": - result = self.vector_store.delete_workspace(workspace_id=request.workspace_id) + elif action == "delete": + result = self.vector_store.delete_workspace(workspace_id=workspace_id) - elif request.action == "delete_ids": - result = self.vector_store.delete(workspace_id=request.workspace_id,node_ids = request.experience_ids) + elif action == "delete_ids": + memory_ids: list = self.context.memory_ids + result = self.vector_store.delete(workspace_id=workspace_id, node_ids=memory_ids) - elif request.action == "dump": - def node_to_experience(node: VectorNode) -> dict: - return vector_node_to_experience(node).model_dump() + elif action == "dump": + path: str = self.context.path + def node_to_memory(node: VectorNode) -> dict: + return vector_node_to_memory(node).model_dump() - result = self.vector_store.dump_workspace(workspace_id=request.workspace_id, - path=request.path, - callback_fn=node_to_experience) + result = self.vector_store.dump_workspace(workspace_id=workspace_id, + path=path, + callback_fn=node_to_memory) - elif request.action == "load": - def experience_dict_to_node(experience_dict: dict) -> VectorNode: - experience: BaseExperience = dict_to_experience(experience_dict=experience_dict) - return experience.to_vector_node() + elif action == "load": + path: str = self.context.path + def memory_dict_to_node(memory_dict: dict) -> VectorNode: + memory: BaseMemory = dict_to_experience(memory_dict=memory_dict) + return memory.to_vector_node() - result = self.vector_store.load_workspace(workspace_id=request.workspace_id, - path=request.path, - callback_fn=experience_dict_to_node) + result = self.vector_store.load_workspace(workspace_id=workspace_id, + path=path, + callback_fn=memory_dict_to_node) - elif request.action == "update_freq": - result = self.vector_store.update_freq(workspace_id=request.workspace_id, node_ids = request.experience_ids) + elif action == "update_freq": + memory_ids: list = self.context.memory_ids + result = self.vector_store.update_freq(workspace_id=workspace_id, node_ids=memory_ids) - elif request.action == "update_utility": - result = self.vector_store.update_utility(workspace_id=request.workspace_id, node_ids = request.experience_ids) + elif action == "update_utility": + memory_ids: list = self.context.memory_ids + result = self.vector_store.update_utility(workspace_id=workspace_id, node_ids=memory_ids) - elif request.action == "utility_based_delete": - result = self.vector_store.utility_based_delete(workspace_id=request.workspace_id, freq_threshold = request.freq_threshold, utility_threshold = request.utility_threshold) + elif action == "utility_based_delete": + freq_threshold: float = self.context.freq_threshold + utility_threshold: float = self.context.utility_threshold + result = self.vector_store.utility_based_delete(workspace_id=workspace_id, + freq_threshold=freq_threshold, + utility_threshold=utility_threshold) else: - raise ValueError(f"invalid action={request.action}") + raise ValueError(f"invalid action={action}") + # Store results in context if isinstance(result, dict): - response.metadata.update(result) + self.context.action_result = result else: - response.metadata["result"] = str(result) + self.context.action_result = {"result": str(result)} diff --git a/test/http_client_test.py b/test/http_client_test.py new file mode 100644 index 00000000..eb211597 --- /dev/null +++ b/test/http_client_test.py @@ -0,0 +1,58 @@ +import asyncio + +import aiohttp + + +async def main(): + base_url = "http://0.0.0.0:8002" + + async with aiohttp.ClientSession() as session: + # 获取工具列表 + print("获取工具列表...") + async with session.get(f"{base_url}/list") as response: + if response.status == 200: + tools = await response.json() + print("可用工具:") + for tool in tools: + print(tool) + else: + print(f"获取工具列表失败: {response.status}") + return + + workspace_id = "default1" + + trajectories = [ + { + "task_id": "t1", + "messages": [ + {"role": "user", "content": "搜索可以使用websearch工具"} + ], + "score": 0.9, + } + ] + + async with session.post( + f"{base_url}/summary_task_memory_simple", + json={ + "trajectories": trajectories, + "workspace_id": workspace_id, + }, + headers={"Content-Type": "application/json"} + ) as response: + result = await response.json() + print(result) + + async with session.post( + f"{base_url}/retrieve_task_memory_simple", + json={ + "query": "茅台怎么样？", + "workspace_id": workspace_id, + }, + headers={"Content-Type": "application/json"} + ) as response: + result = await response.json() + print(result) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/test/mcp_client_test.py b/test/mcp_client_test.py index 6a0de445..aae8f6fb 100644 --- a/test/mcp_client_test.py +++ b/test/mcp_client_test.py @@ -8,11 +8,30 @@ async def main(): for tool in tools: print(tool.model_dump_json()) - result: CallToolResult = await client.call_tool("retrieve_task_memory", + workspace_id = "default" + + result: CallToolResult = await client.call_tool("retrieve_task_memory_simple", arguments={ "query": "茅台怎么样？", - "workspace_id": "default", - "top_k": 1, + "workspace_id": workspace_id, + }) + print(result.content) + + trajectories = [ + { + "task_id": "t1", + "messages": [ + {"role": "user", "content": "今天天气不错"} + ], + "score": 0.9, + } + ] + + + result: CallToolResult = await client.call_tool("summary_task_memory_simple", + arguments={ + "trajectories": trajectories, + "workspace_id": workspace_id, }) print(result.content)