Coverage for agentos/rag/loader.py: 20%

1""" 

2文档加载器 — PDF / DOCX / TXT / Markdown 解析与分块。 

3 

4零外部 HTTP 依赖，纯本地解析。 

5""" 

6 

7from __future__ import annotations 

8 

9import os 

10from pathlib import Path 

11 

12 

13class Document: 

14 """文档片段。""" 

15 

16 def __init__(self, content: str, source: str = "", page: int = 0, metadata: dict | None = None): 

17 self.content = content 

18 self.source = source 

19 self.page = page 

20 self.metadata = metadata or {} 

21 

22 def __repr__(self): 

23 return f"Document(source={self.source!r}, chars={len(self.content)})" 

24 

25 

26class DocumentLoader: 

27 """文档加载器 — 支持多种格式的文档解析与智能分块。 

28 

29 Args: 

30 chunk_size: 分块大小（字符数） 

31 chunk_overlap: 块间重叠字符数 

32 """ 

33 

34 SUPPORTED_SUFFIXES = {".pdf", ".docx", ".txt", ".md", ".markdown", ".py", ".json", ".yaml", ".yml"} 

35 

36 def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200): 

37 self.chunk_size = chunk_size 

38 self.chunk_overlap = chunk_overlap 

39 

40 def load_file(self, path: str) -> list[Document]: 

41 """加载单个文件，自动识别格式。""" 

42 path = os.path.abspath(path) 

43 suffix = Path(path).suffix.lower() 

44 if suffix not in self.SUPPORTED_SUFFIXES: 

45 raise ValueError(f"不支持的文件格式: {suffix}。支持: {self.SUPPORTED_SUFFIXES}") 

46 

47 if suffix == ".pdf": 

48 text = self._read_pdf(path) 

49 elif suffix == ".docx": 

50 text = self._read_docx(path) 

51 else: 

52 with open(path, "r", encoding="utf-8", errors="replace") as f: 

53 text = f.read() 

54 

55 return self._chunk(text, source=path) 

56 

57 def load_directory(self, dir_path: str, recursive: bool = True) -> list[Document]: 

58 """加载目录下所有支持的文件。""" 

59 docs = [] 

60 for root, _, files in os.walk(dir_path): 

61 for fn in sorted(files): 

62 fp = os.path.join(root, fn) 

63 suffix = Path(fp).suffix.lower() 

64 if suffix in self.SUPPORTED_SUFFIXES: 

65 try: 

66 docs.extend(self.load_file(fp)) 

67 except Exception: 

68 pass 

69 if not recursive: 

70 break 

71 return docs 

72 

73 def _read_pdf(self, path: str) -> str: 

74 """读取 PDF 文本。""" 

75 try: 

76 import pypdf 

77 reader = pypdf.PdfReader(path) 

78 pages = [] 

79 for page in reader.pages: 

80 text = page.extract_text() 

81 if text: 

82 pages.append(text) 

83 return "\n\n".join(pages) 

84 except ImportError: 

85 raise ImportError("pypdf 未安装。运行: pip install pypdf") 

86 

87 def _read_docx(self, path: str) -> str: 

88 """读取 DOCX 文本。""" 

89 try: 

90 from docx import Document as DocxDocument 

91 doc = DocxDocument(path) 

92 paragraphs = [p.text for p in doc.paragraphs if p.text.strip()] 

93 return "\n".join(paragraphs) 

94 except ImportError: 

95 raise ImportError("python-docx 未安装。运行: pip install python-docx") 

96 

97 def _chunk(self, text: str, source: str = "") -> list[Document]: 

98 """固定大小+重叠分块。""" 

99 if len(text) <= self.chunk_size: 

100 return [Document(content=text.strip(), source=source)] 

101 

102 chunks = [] 

103 start = 0 

104 while start < len(text): 

105 end = min(start + self.chunk_size, len(text)) 

106 chunk = text[start:end].strip() 

107 if chunk: 

108 chunks.append(Document(content=chunk, source=source)) 

109 start += self.chunk_size - self.chunk_overlap 

110 return chunks 

111 

112 

113# ── 便捷函数 ────────────────────────────────────────────────── 

114 

115def load_file(path: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> list[Document]: 

116 """便捷函数：加载单个文件。""" 

117 loader = DocumentLoader(chunk_size=chunk_size, chunk_overlap=chunk_overlap) 

118 return loader.load_file(path) 

119 

120 

121def load_directory(dir_path: str, recursive: bool = True, chunk_size: int = 1000, chunk_overlap: int = 200) -> list[Document]: 

122 """便捷函数：加载目录。""" 

123 loader = DocumentLoader(chunk_size=chunk_size, chunk_overlap=chunk_overlap) 

124 return loader.load_directory(dir_path, recursive=recursive)