Spaces:
Sleeping
Sleeping
| from langgraph.graph import END, StateGraph | |
| from src.nodes import refiner_node, simplifier_node, solver_node | |
| import asyncio | |
| from src.state import AgentState | |
| def should_continue(state: AgentState) -> str: | |
| """Determines whether to continue solving subqueries or | |
| refine the final answer. | |
| Args: | |
| state (AgentState): The current state of the agent, containing | |
| subqueries, results, current step index, and | |
| message history. | |
| Returns: | |
| str: A string indicating the next step: | |
| - "continue_solve": Continue solving subqueries. | |
| - "refine": Refine the final answer. | |
| """ | |
| # Check if we have any subqueries to process | |
| if not state["subqueries"]: | |
| return "refine" | |
| current_step = state.get("current_step", 0) | |
| # Check if we've processed all subqueries | |
| if current_step >= len(state["subqueries"]): | |
| return "refine" | |
| # Continue processing | |
| return "continue_solve" | |
| class AgentGAIA: | |
| """Agent GAIA for multi-step query resolution.""" | |
| def __init__(self): | |
| """Initialize the Agent GAIA with a state graph.""" | |
| # Create the state graph | |
| builder = StateGraph(AgentState) | |
| # Add nodes to the graph | |
| builder.add_node("simplify", simplifier_node) | |
| builder.add_node("solve", solver_node) | |
| builder.add_node("refine", refiner_node) | |
| # Set the initial state | |
| builder.set_entry_point("simplify") | |
| # Define transitions | |
| builder.add_edge("simplify", "solve") | |
| builder.add_conditional_edges( | |
| "solve", | |
| should_continue, | |
| { | |
| "continue_solve": "solve", | |
| "refine": "refine" | |
| } | |
| ) | |
| builder.add_edge("refine", END) | |
| self.graph = builder.compile() | |
| async def __call__(self, query: str) -> str: | |
| """Run a query through the GAIA Agent. | |
| Args: | |
| query (str): The initial query to be processed. | |
| Returns: | |
| str: The final answer generated by the agent. | |
| """ | |
| state = { | |
| "original_query": query, | |
| "subqueries": [], | |
| "current_step": 0, | |
| "results": [], | |
| "final_answer": None, | |
| "messages": [] | |
| } | |
| state = await self.graph.ainvoke(state) | |
| return state["final_answer"] or state["results"][-1] | |
| if __name__ == "__main__": | |
| agent = AgentGAIA() | |
| # Example query | |
| query = ( | |
| "What is the GDP growth difference between France and Germany after " | |
| "their 2020 policy changes?" | |
| ) | |
| async def main(): | |
| answer = await agent(query) | |
| print(answer) | |
| asyncio.run(main()) | |