Spaces:
Sleeping
Sleeping
| import os | |
| from langchain_core.messages import HumanMessage | |
| from langchain_core.output_parsers import JsonOutputParser | |
| from langchain_openai import AzureChatOpenAI | |
| from src.prompts import simplifier_prompt | |
| from src.state import AgentState | |
| async def simplifier_node(state: AgentState): | |
| """Processes the agent's state to generate simplified subqueries | |
| from the original query using a language model. | |
| This function takes the current agent state, sends the original query to | |
| a language model via a prompt chain, and parses the output to extract a | |
| list of simpler/atomic subqueries. It then updates the state with the | |
| generated subqueries, resets the current step and results, and appends | |
| the original query as a message. | |
| Args: | |
| state (AgentState): The current state of the agent, expected to | |
| contain the key "original_query". | |
| Returns: | |
| dict: An updated state dictionary containing: | |
| - "subqueries": List of subqueries generated from the | |
| original query. | |
| - "current_step": Reset to 0. | |
| - "results": Reset to an empty list. | |
| - "messages": List containing the original query as a HumanMessage. | |
| - All other original state keys and values. | |
| """ | |
| # Create a message with the original query | |
| messages = [HumanMessage(content=state["original_query"])] | |
| # Create the simplifier chain | |
| chain = ( | |
| simplifier_prompt | | |
| AzureChatOpenAI(deployment_name=os.getenv( | |
| "AZURE_OPENAI_DEPLOYMENT_NAME")) | | |
| JsonOutputParser() | |
| ) | |
| # Invoke the simplifier chain | |
| result = await chain.ainvoke({"query": state["original_query"]}) | |
| return { | |
| **state, | |
| "subqueries": result["subqueries"], | |
| "messages": messages | |
| } | |