1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
| class RAGAgent:
def __init__(self, model: str )
self.client = OpenAI(
api_key = os.getenv("OPENAI_API_KEY"),
base_url = os.getenv("OPENAI_BASE_URL")
)
self.model = model
self.rag = RAGStore(collection_name = "knowledge_base"
self.chunker = SimpleChunker(chunk_size = 1500, overlap = 100)
self.memory = MemorySystem()
self.tool_history : List[Dict] = []
self.max_tool_calls = 5
def add_knowledge(self,text:str,source:str = "知识库路径" )
chunks = self.chunker.chunk(text)
metadatas = [{"source":source,"chunk_index":i}] for i in range(len(chunks))]
self.rag.add_documents(chunks,metadatas = metadatas)
def retrieve(self, query:str ,top_k :int = 3) -> List[dict]:
results = self.rag.query(query,n_results = top_k)
return results
def execute_tool(self,tool_name:str ,arguments: Dict) -> str:
call_signature = json.dumps({"tool":tool_name,"args":arguments},sort_keys = Ture)
recent_calls = [json.dumps(c,sort_keys = Ture) for c in self.tool_history[-3:]]
self.tool_history.append({"tool": tool_name, "args": arguments, "time": time.time()})
if tool_name not in TOOLS_REGISTER:
return f"错误：工具 '{tool_name}' 未注册"
try:
result = base_execute_tool(tool_name,arguments)
return results
def call_llm(self,messages: List[Dict], temperate: fload = 0.1) -> str:
try:
response = self.client.caht.completions.create(
model = self.model,
messages = messages,
temperature = temperature
)
return response.choices[0].message.content.strip()
def run(self, user_input: str) -> Dict[str, Any]:
self.memory.add_memory("user", user_input)
retrieved_docs = self.retrieve(user_input, top_k = 3)
rag_context = self._formate_rag_context(retrieved_docs)
base_message = self._build_base_massages(rag_context)
message = base_massage.copy()
tool_used = []
tool_call_count = 0
while tool_call_count < self.max_tool_calls:
decision = self.call_llm(messages)
tool_name = self._extract_tool_name(decision)
if tool_name is None:
不需要工具直接回答，保存记忆
self.memory.add_message("assistant",answer["answer"])
return {
"answer": answer["answer"],
"sources": answer["sources"],
"tools_used": tools_used,
"retrieved_docs": retrieved_docs
}
tool_call_json = self._generate_tool_call(tool_name, user_input)
parsed = self._extract_json(tool_call_json)
if parsed is None:
messages.append({"role": "assistant", "content": f"TOOL:{tool_name}"})
messages.append({"role": "user", "content": f"工具参数解析失败。请重新输出严格 JSON 格式：{{\"tool\": \"{tool_name}\", \"arguments\": {{...}}}}。不要任何其他文字。"})
actual_tool = parsed.get("tool", tool_name)
arguments = parsed.get("arguments", {})
result = self.execute_tool(actual_tool, arguments)
tools_used.append({"tool": actual_tool, "args": arguments, "result": result[:200]})
messages.append({"role": "assistant", "content": f"TOOL:{actual_tool}"})
messages.append({"role": "user", "content": f"工具【{actual_tool}】返回结果：\n{result}\n\n请基于以上结果，继续回答原始问题：{user_input}"})
tool_call_count += 1
|