Spaces:
Sleeping
Sleeping
File size: 2,762 Bytes
9c89d35 d9af1db 9c89d35 d9af1db 9c89d35 d9af1db 9c89d35 d9af1db 9c89d35 d9af1db 9c89d35 d9af1db 9c89d35 d9af1db 9c89d35 d9af1db 9c89d35 dcf4bc4 9c89d35 | 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 | 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())
|