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 = { 

35 ".pdf", 

36 ".docx", 

37 ".txt", 

38 ".md", 

39 ".markdown", 

40 ".py", 

41 ".json", 

42 ".yaml", 

43 ".yml", 

44 } 

45 

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

47 self.chunk_size = chunk_size 

48 self.chunk_overlap = chunk_overlap 

49 

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

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

52 path = os.path.abspath(path) 

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

54 if suffix not in self.SUPPORTED_SUFFIXES: 

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

56 

57 if suffix == ".pdf": 

58 text = self._read_pdf(path) 

59 elif suffix == ".docx": 

60 text = self._read_docx(path) 

61 else: 

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

63 text = f.read() 

64 

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

66 

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

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

69 docs = [] 

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

71 for fn in sorted(files): 

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

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

74 if suffix in self.SUPPORTED_SUFFIXES: 

75 try: 

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

77 except Exception: 

78 pass 

79 if not recursive: 

80 break 

81 return docs 

82 

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

84 """读取 PDF 文本。""" 

85 try: 

86 import pypdf 

87 

88 reader = pypdf.PdfReader(path) 

89 pages = [] 

90 for page in reader.pages: 

91 text = page.extract_text() 

92 if text: 

93 pages.append(text) 

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

95 except ImportError: 

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

97 

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

99 """读取 DOCX 文本。""" 

100 try: 

101 from docx import Document as DocxDocument 

102 

103 doc = DocxDocument(path) 

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

105 return "\n".join(paragraphs) 

106 except ImportError: 

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

108 

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

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

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

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

113 

114 chunks = [] 

115 start = 0 

116 while start < len(text): 

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

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

119 if chunk: 

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

121 start += self.chunk_size - self.chunk_overlap 

122 return chunks 

123 

124 

125# ── 便捷函数 ────────────────────────────────────────────────── 

126 

127 

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

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

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

131 return loader.load_file(path) 

132 

133 

134def load_directory( 

135 dir_path: str, recursive: bool = True, chunk_size: int = 1000, chunk_overlap: int = 200 

136) -> list[Document]: 

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

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

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