82 lines
1.9 KiB
Python
82 lines
1.9 KiB
Python
from tools import TOOL_MAP
|
|
from llm import llm_call
|
|
import re
|
|
from typing import List, Dict
|
|
from models import GraphRAGResponse
|
|
|
|
SYSTEM_PROMPT = """
|
|
You are an educational GraphRAG assistant for schools.
|
|
|
|
You MUST use tools when answering knowledge questions.
|
|
|
|
Available tool:
|
|
- query_knowledge_graph
|
|
|
|
Format:
|
|
Action: query_knowledge_graph
|
|
Action Input: <question>
|
|
|
|
OR:
|
|
|
|
Final Answer: <answer>
|
|
"""
|
|
|
|
def parse_action(text: str):
|
|
action = re.search(r"Action:\s*(\w+)", text)
|
|
input_ = re.search(r"Action Input:\s*(.+)", text)
|
|
|
|
if action and input_:
|
|
return action.group(1), input_.group(1)
|
|
|
|
return None
|
|
|
|
def parse_final(text: str):
|
|
m = re.search(r"Final Answer:\s*(.+)", text, re.DOTALL)
|
|
return m.group(1).strip() if m else None
|
|
|
|
def react_agent(user_message: str, history: List[Dict], G=None, max_steps: int = 5):
|
|
|
|
messages = (
|
|
[{"role": "system", "content": SYSTEM_PROMPT}]
|
|
+ history
|
|
+ [{"role": "user", "content": user_message}]
|
|
)
|
|
|
|
last = ""
|
|
|
|
for _ in range(max_steps):
|
|
|
|
reply = llm_call(messages)
|
|
last = reply
|
|
|
|
final = parse_final(reply)
|
|
if final:
|
|
return GraphRAGResponse(
|
|
question=user_message,
|
|
answer=final,
|
|
evidence=""
|
|
)
|
|
|
|
action = parse_action(reply)
|
|
|
|
if action:
|
|
|
|
tool_name, tool_input = action
|
|
|
|
if tool_name == "query_knowledge_graph":
|
|
result = TOOL_MAP[tool_name](tool_input, G)
|
|
|
|
messages.append({"role": "assistant", "content": reply})
|
|
messages.append({"role": "user", "content": f"Observation: {result}"})
|
|
|
|
return GraphRAGResponse(
|
|
question=user_message,
|
|
answer=result["answer"],
|
|
evidence=result["evidence"]
|
|
)
|
|
|
|
return GraphRAGResponse(
|
|
question=user_message,
|
|
answer=last,
|
|
evidence=""
|
|
) |