scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-09-10 22:41:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

feat(reme_ai): implement summary functionality and enhance vector store operations

Browse source 
- Add summary module with simple summary and personal summary tasks
- Implement new vector store operations for updating and managing memory
- Refactor retrieve module to improve query building and memory rewriting
- Update schema and utils modules for better data handling and PDF processing
This commit is contained in:
jinli.yl 2025-08-26 20:56:30 +08:00
parent cceebc631e
commit 5a3d8ab74d
21 changed files with 932 additions and 602 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
reme_ai/__init__.py
View file

@ -1,4 +1,5 @@
from reme_ai import retrieve
from reme_ai import summary
from reme_ai import vector_store
__version__ = "0.1.0"

15
reme_ai/config/default.yaml
View file

@ -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

1
reme_ai/retrieve/__init__.py
View file

@ -1 +1,2 @@
from reme_ai.retrieve import task
from reme_ai.retrieve import personal

2
reme_ai/retrieve/task/build_query_op.py
View file

@ -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:

2
reme_ai/retrieve/task/rerank_memory_op.py
View file

@ -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"""

5
reme_ai/retrieve/task/rewrite_memory_op.py
View file

@ -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 ""

1
reme_ai/schema/__init__.py
View file

@ -0,0 +1 @@
from flowllm.schema.message import Message, Trajectory, ToolCall

2
reme_ai/summary/__init__.py
View file

@ -0,0 +1,2 @@
from reme_ai.summary import task
from reme_ai.summary import personal

1
reme_ai/summary/task/__init__.py
View file

@ -0,0 +1 @@
from reme_ai.summary.task.simple_summary_op import SimpleSummaryOp

4
reme_ai/summary/task/comparative_extraction_op.py
View file

@ -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()

503
reme_ai/summary/task/pdf_preprocess_op.py
View file

@ -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))

47
reme_ai/summary/task/simple_summary_op.py
View file

@ -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}")
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}")

12
reme_ai/summary/task/simple_summary_prompt.yaml
View file

@ -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": "..."
}
]
```

721
reme_ai/utils/miner_u_pdf_processor.py Normal file
View file

@ -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()

4
reme_ai/utils/op_utils.py
View file

@ -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:

4
reme_ai/vector_store/__init__.py
View file

@ -1 +1,3 @@
from reme_ai.vector_store.recall_vector_store_op import RecallVectorStoreOp
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

3
reme_ai/vector_store/recall_vector_store_op.py
View file

@ -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)

36
reme_ai/vector_store/update_vector_store_op.py
View file

@ -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
}

87
reme_ai/vector_store/vector_store_action_op.py
View file

@ -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)}

58
test/http_client_test.py Normal file
View file

@ -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())

25
test/mcp_client_test.py
View file

@ -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)