1"""

2 RAG Pipeline — 检索增强生成管道。

3

4 将文档加载 → 向量化 → 检索 → LLM 生成 串联为一条端到端管道。

5"""

6

7 from __future__ import annotations

8

9 from agentos.rag.store import VectorStore, ChromaStore

10 from agentos.rag.loader import DocumentLoader, Document

11 from agentos.agent.tool_agent import ToolAgent, AgentConfig

12

13

14 class RAGPipeline:

15 """检索增强生成管道。

16

17 使用流程::

18

19 rag = RAGPipeline(agent)

20 rag.ingest_file("docs/report.pdf") # 加载文档到向量库

21 rag.ingest_directory("docs/project/") # 批量加载

22 answer = rag.query("Q3 收入是多少？") # 检索 + 生成

23

24 Args:

25 agent: 用于生成答案的 ToolAgent

26 vector_store: 向量存储（默认 ChromaDB 内存模式）

27 top_k: 检索返回片段数

28 """

29

30 def __init__(

31 self,

32 agent: ToolAgent,

33 vector_store: VectorStore | None = None,

34 top_k: int = 5,

35 ):

36 self._agent = agent

37 self._store = vector_store or ChromaStore()

38 self._top_k = top_k

39 self._loader = DocumentLoader()

40

41 @property

42 def store(self) -> VectorStore:

43 return self._store

44

45 @property

46 def doc_count(self) -> int:

47 return self._store.count()

48

49 def ingest_file(self, path: str) -> int:

50 """加载单个文件到向量库。返回添加的块数。"""

51 docs = self._loader.load_file(path)

52 if not docs:

53 return 0

54 texts = [d.content for d in docs]

55 metadatas = [{"source": d.source, "page": d.page} for d in docs]

56 self._store.add(texts, metadatas)

57 return len(texts)

58

59 def ingest_directory(self, dir_path: str, recursive: bool = True) -> int:

60 """加载目录到向量库。返回添加的块数。"""

61 docs = self._loader.load_directory(dir_path, recursive=recursive)

62 if not docs:

63 return 0

64 texts = [d.content for d in docs]

65 metadatas = [{"source": d.source} for d in docs]

66 self._store.add(texts, metadatas)

67 return len(texts)

68

69 def ingest_texts(self, texts: list[str], metadatas: list[dict] | None = None) -> int:

70 """直接添加文本列表到向量库。"""

71 self._store.add(texts, metadatas)

72 return len(texts)

73

74 def query(self, question: str, top_k: int | None = None) -> str:

75 """检索 + 生成答案。

76

77 Args:

78 question: 用户问题

79 top_k: 检索片段数（覆盖默认值）

80

81 Returns:

82 LLM 生成的答案

83 """

84 k = top_k or self._top_k

85

86 if self._store.count() == 0:

87 return self._agent.run(question).final_answer

88

89 # 检索

90 results = self._store.search(question, top_k=k)

91 if not results:

92 return self._agent.run(question).final_answer

93

94 # 构建上下文

95 context_parts = []

96 sources = set()

97 for i, r in enumerate(results, 1):

98 context_parts.append(f"[{i}] {r.content}")

99 if r.metadata.get("source"):

100 sources.add(r.metadata["source"])

101

102 context = "\n\n".join(context_parts)

103 source_list = "\n".join(f"- {s}" for s in sorted(sources))

104

105 # 构建 RAG prompt

106 rag_task = f"""请基于以下参考资料回答用户的问题。

107 如果参考资料不足以回答问题，请明确说明，不要编造信息。

108

109## 参考资料

110{context}

111

112## 来源文件

113{source_list}

114

115## 用户问题

116{question}

117

118 请用中文回答。回答时引用具体的资料编号（如 [1]）。"""

119

120 return self._agent.run(rag_task).final_answer

121

122 def retrieve(self, question: str, top_k: int | None = None) -> list[str]:

123 """仅检索，不生成。返回匹配的文本片段。"""

124 k = top_k or self._top_k

125 results = self._store.search(question, top_k=k)

126 return [r.content for r in results]