Coverage for agentos/rag/__init__.py: 74%

1""" 

2RAG 模块入口 — 向量存储 / 文档加载 / 检索生成管道 / 混合搜索。 

3 

4v1.5.1: ChromaDB 向量存储 + PDF/DOCX/TXT 文档加载 + 基础 RAG Pipeline。 

5v1.9.0: 混合搜索 + BM25 稀疏检索 + 跨编码器重排 + 引用追踪。 

6""" 

7 

8from agentos.rag.hybrid_search import ( 

9 BM25Retriever, 

10 Citation, 

11 CitationTracker, 

12 CrossEncoderReranker, 

13 DenseRetriever, 

14 FusionMethod, 

15 HybridSearchEngine, 

16 SearchResult, 

17) 

18from agentos.rag.loader import DocumentLoader, load_directory, load_file 

19from agentos.rag.pipeline import RAGPipeline 

20from agentos.rag.store import ChromaStore, VectorStore 

21 

22# 向后兼容别名 

23BaseVectorStore = VectorStore 

24FAISSVectorStore = None # FAISS 已移除，由 ChromaDB 替代 

25ChromaVectorStore = ChromaStore 

26 

27# 配置兼容别名 

28ChunkConfig = dict # chunk_size + chunk_overlap 

29EmbeddingConfig = dict # model_name + device 

30 

31 

32# TextChunker 兼容层 

33class TextChunker: 

34 """文本分块器（向后兼容）。""" 

35 

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

37 self._loader = DocumentLoader(chunk_size=chunk_size, chunk_overlap=chunk_overlap) 

38 

39 def chunk(self, text: str, source: str = "") -> list[str]: 

40 docs = self._loader._chunk(text, source) 

41 return [d.content for d in docs] 

42 

43 def chunk_file(self, path: str) -> list[str]: 

44 docs = self._loader.load_file(path) 

45 return [d.content for d in docs] 

46 

47 

48__all__ = [ 

49 "VectorStore", 

50 "ChromaStore", 

51 "DocumentLoader", 

52 "load_file", 

53 "load_directory", 

54 "RAGPipeline", 

55 # 混合搜索 v1.9.0 

56 "HybridSearchEngine", 

57 "BM25Retriever", 

58 "DenseRetriever", 

59 "CrossEncoderReranker", 

60 "CitationTracker", 

61 "Citation", 

62 "SearchResult", 

63 "FusionMethod", 

64 # 向后兼容 

65 "BaseVectorStore", 

66 "FAISSVectorStore", 

67 "ChromaVectorStore", 

68 "TextChunker", 

69]