File size: 3,830 Bytes
63f13b5 0ea9e1b 63f13b5 0ea9e1b a72ec7c d7bbf02 a72ec7c 63f13b5 0ea9e1b 1488b1e 63f13b5 0ea9e1b 63f13b5 a72ec7c ba74371 d7bbf02 9922f58 ba74371 c445d91 d7bbf02 f53d974 a72ec7c 63f13b5 0ea9e1b 63f13b5 a72ec7c 0ea9e1b 63f13b5 a72ec7c ff3a79a 63f13b5 a72ec7c c932e96 63f13b5 a72ec7c ff3a79a a72ec7c ff3a79a a72ec7c ff3a79a a72ec7c ff3a79a a72ec7c ff3a79a 63f13b5 ff3a79a 63f13b5 a72ec7c 63f13b5 a72ec7c 63f13b5 a72ec7c 63f13b5 0ea9e1b | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import os
import pandas as pd
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import StateGraph, MessagesState
from langgraph.prebuilt import ToolNode
from langchain_huggingface import HuggingFaceEndpoint
from tools import TOOLS # your dictionary of tool functions
# --- Load local QA metadata for retriever ---
QA_PATH = "metadata.jsonl"
qa_pairs = pd.read_json(QA_PATH, lines=True)
qa_dict = {
row["Question"].strip(): row["Final answer"].strip()
for _, row in qa_pairs.iterrows()
}
# === LangGraph builder ===
def build_graph():
llm = HuggingFaceEndpoint(
repo_id="mistralai/Mistral-7B-Instruct-v0.3",
task="text-generation",
huggingfacehub_api_token=os.environ["HF_TOKEN"]
)
# Node 1: Retriever
def retriever_node(state: MessagesState):
query = state["messages"][-1].content.strip()
if query in qa_dict:
print("✅ Exact match found in retriever.")
return {"messages": [AIMessage(content=qa_dict[query])]}
print("🔍 No match. Sending to LLM.")
return {"messages": state["messages"]}
# Node 2: Assistant (LLM response parsing)
def assistant_node(state: MessagesState):
query = state["messages"][-1].content.strip()
prompt = (
"You are a helpful assistant for GAIA benchmark.\n"
"If you can answer directly, output ONLY the answer.\n"
"If you need to use a tool, reply in this format:\n"
"use_tool: <tool_name>: <tool_input>\n"
"Never explain anything."
)
full_prompt = f"{prompt}\n\nQuestion: {query}\nAnswer:"
response = llm.invoke(full_prompt).strip()
print(f"🧠 LLM said: {response}")
if response.startswith("use_tool:"):
return {
"messages": state["messages"] + [AIMessage(content=response)],
"tool_call": response # carry tool signal
}
else:
return {"messages": [AIMessage(content=response)]}
# Node 3: Tool execution
def tool_node(state: MessagesState):
try:
tool_signal = state.get("tool_call", "")
_, tool_name, tool_input = tool_signal.split(":", 2)
tool_name = tool_name.strip()
tool_input = tool_input.strip()
tool_fn = TOOLS.get(tool_name)
if not tool_fn:
print(f"❌ Unknown tool: {tool_name}")
return {"messages": [AIMessage(content="Unknown")]}
print(f"🔧 Using tool: {tool_name} with input: {tool_input}")
tool_result = tool_fn(tool_input)
return {"messages": [AIMessage(content=str(tool_result))]}
except Exception as e:
print(f"⚠️ Tool error: {e}")
return {"messages": [AIMessage(content="Unknown")]} # fail-safe
# Build LangGraph
builder = StateGraph(MessagesState)
builder.add_node("retriever", retriever_node)
builder.add_node("assistant", assistant_node)
builder.add_node("tool", tool_node)
builder.set_entry_point("retriever")
builder.add_edge("retriever", "assistant")
builder.add_edge("assistant", "tool")
builder.add_edge("tool", "assistant")
builder.set_finish_point("assistant")
return builder.compile()
# === BasicAgent wrapper ===
class BasicAgent:
def __init__(self):
print("BasicAgent initialized with retriever + LLM + manual tool logic.")
self.graph = build_graph()
def __call__(self, question: str) -> str:
print(f"📥 Question: {question[:100]}")
result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
answer = result["messages"][-1].content.strip()
print(f"📤 Answer: {answer}")
return answer
|