KoRiF commited on
Commit
3d98ed0
·
1 Parent(s): 81917a3

introduce LangGraph-based QA-workflow with CodeAgent answers generation

Browse files
Files changed (2) hide show
  1. answering.py +17 -0
  2. workflow.py +117 -0
answering.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ from smolagents import CodeAgent, InferenceClientModel, FinalAnswerTool, DuckDuckGoSearchTool, WikipediaSearchTool, PythonInterpreterTool
4
+
5
+ from dotenv import load_dotenv
6
+ load_dotenv()
7
+
8
+
9
+ def gen_question_answer(question: str) -> str:
10
+ duck_duck_go_search_tool = DuckDuckGoSearchTool()
11
+ wikipedia_search_tool = WikipediaSearchTool()
12
+ final_answer_tool = FinalAnswerTool()
13
+ python_interpreter_tool = PythonInterpreterTool()
14
+ agent = CodeAgent(tools=[duck_duck_go_search_tool, wikipedia_search_tool, python_interpreter_tool, final_answer_tool], model=InferenceClientModel(), add_base_tools=True, additional_authorized_imports=["pandas"])
15
+
16
+ response = agent.run(question) # Use run() instead of query()
17
+ return str(response)
workflow.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import TypedDict, Annotated, Callable, Optional, Any
2
+ from langgraph.graph import StateGraph, END
3
+ from answering import gen_question_answer
4
+
5
+
6
+ class AgentState(TypedDict):
7
+ question: str
8
+ answer: Annotated[str, lambda x, y: y] # Overwrite with new value
9
+ formatted_answer: Annotated[str, lambda x, y: y]
10
+
11
+ class GAIAAnsweringWorkflow:
12
+ def __init__(
13
+ self,
14
+ qa_function: Optional[Callable[[str], str]] = None,
15
+ formatter: Optional[Callable[[str], str]] = None
16
+ ):
17
+ """
18
+ Initialize the GAIA agent workflow
19
+
20
+ Args:
21
+ qa_function: Core question answering function (gen_question_answer)
22
+ formatter: Answer formatting function (default: GAIA boxed format)
23
+ """
24
+ self.qa_function = gen_question_answer #qa_function or self.default_qa_function
25
+ self.formatter = formatter or self.default_formatter
26
+ self.workflow = self.build_workflow()
27
+
28
+ @staticmethod
29
+ def default_qa_function(question: str) -> str:
30
+ """Placeholder QA function (override with your CodeAgent)"""
31
+ return "42"
32
+
33
+ @staticmethod
34
+ def default_formatter(answer: str) -> str:
35
+ """Default GAIA formatting"""
36
+ return answer #f"\\boxed{{{answer}}}"
37
+
38
+ def build_workflow(self) -> Any:
39
+ """Construct and compile the LangGraph workflow"""
40
+ # Create graph
41
+ workflow = StateGraph(AgentState)
42
+
43
+ # Add nodes
44
+ workflow.add_node("generate_answer", self.generate_answer_node)
45
+ workflow.add_node("format_output", self.format_output_node)
46
+
47
+ # Define edges
48
+ workflow.set_entry_point("generate_answer")
49
+ workflow.add_edge("generate_answer", "format_output")
50
+ workflow.add_edge("format_output", END)
51
+
52
+ return workflow.compile()
53
+
54
+ def generate_answer_node(self, state: AgentState) -> dict:
55
+ """Node that executes the question answering tool"""
56
+ try:
57
+ answer = self.qa_function(state["question"])
58
+ return {"answer": answer}
59
+ except Exception as e:
60
+ print(str(e))
61
+ return {"answer": f"Error: {str(e)}"}
62
+
63
+ def format_output_node(self, state: AgentState) -> dict:
64
+ """Node that formats the answer for GAIA benchmark"""
65
+ try:
66
+ formatted = self.formatter(state["answer"])
67
+ return {"formatted_answer": formatted}
68
+ except Exception as e:
69
+ return {"formatted_answer": f"\\boxed{{\\text{{Formatting error: {str(e)}}}}}"}
70
+
71
+ def __call__(self, question: str) -> str:
72
+ """
73
+ Execute the agent workflow for a given question
74
+
75
+ Args:
76
+ question: Input question string
77
+
78
+ Returns:
79
+ Formatted GAIA answer
80
+ """
81
+ # Initialize state
82
+ initial_state = {
83
+ "question": question,
84
+ "answer": "",
85
+ "formatted_answer": ""
86
+ }
87
+
88
+ # Execute workflow
89
+ result = self.workflow.invoke(initial_state)
90
+ return result["formatted_answer"]
91
+
92
+ # Example usage with custom QA function
93
+ if __name__ == "__main__":
94
+ # Custom QA function (replace with your CodeAgent integration)
95
+ def custom_qa(question: str) -> str:
96
+ if "life" in question:
97
+ return "42"
98
+ elif "prime" in question:
99
+ return "101"
100
+ return "unknown"
101
+
102
+ # Create agent instance
103
+ agent = GAIAAnsweringWorkflow(
104
+ qa_function=gen_question_answer,
105
+ formatter=lambda ans: f"ANSWER: \\boxed{{{ans}}}" # Custom formatting
106
+ )
107
+
108
+ # Test cases
109
+ questions = [
110
+ "What is the answer to life, the universe, and everything?",
111
+ "What is the smallest 3-digit prime number?",
112
+ "Unknown question type?"
113
+ ]
114
+
115
+ for q in questions:
116
+ result = agent(q)
117
+ print(f"Question: {q}\nAnswer: {result}\n")