04. Structured Reasoning & LangGraph

from pydantic import BaseModel, Field
from typing import List
 
class SubQueries(BaseModel):
    queries: List[str] = Field(description="Sub-questions to solve the original query")
 
def decompose_query(query: str) -> List[str]:
    """Decomposes a complex query into sub-queries."""
    completion = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "Break down complex problems into 3-4 sequential sub-questions."},
            {"role": "user", "content": query}
        ],
        response_format=SubQueries
    )
    return completion.choices[0].message.parsed.queries

def solve_decomposed_queries(original_query: str):
    sub_qs = decompose_query(original_query)
    context = ""
 
    for i, q in enumerate(sub_qs):
        print(f"[Step {i+1}] Q: {q}")
        answer = answer_question(q, context)
        print(f"  ✅ A: {answer}")
        context += f"Q: {q}\nA: {answer}\n\n"
 
    # Synthesize final answer
    final_answer = answer_question(f"Answer: '{original_query}'", context)
    return final_answer

# 1. Proposer: Generate candidate continuations
def propose_next_steps(current_state: str, n: int = 3) -> List[str]:
    prompt = f"""Current state: {current_state}
 
Propose {n} possible next steps. Number them 1, 2, 3.
"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}]
    )
    # Parse numbered responses
    lines = response.choices[0].message.content.strip().split("\n")
    return [line.split(". ", 1)[1] for line in lines if ". " in line][:n]
 
# 2. Evaluator: Score candidates
def evaluate_candidates(candidates: List[str], goal: str) -> List[int]:
    scores = []
    for cand in candidates:
        prompt = f"""Goal: {goal}
Candidate: {cand}
 
Rate relevance and quality (1-10). Return number only.
"""
        response = client.chat.completions.create(...)
        scores.append(int(response.choices[0].message.content.strip()))
    return scores
 
# 3. ToT Search Loop
def run_tot(goal: str, steps: int = 3):
    current_state = ""
 
    for i in range(steps):
        candidates = propose_next_steps(current_state)
        scores = evaluate_candidates(candidates, goal)
 
        # Select best candidate
        best_idx = scores.index(max(scores))
        current_state += candidates[best_idx] + " "
 
    return current_state

from langgraph.graph import StateGraph
from typing import TypedDict, List
 
# 1. Define State Schema
class AgentState(TypedDict):
    query: str
    plan: List[str]
    results: List[str]
    done: bool
 
# 2. Create Graph
graph = StateGraph(AgentState)
 
# 3. Add Nodes
graph.add_node("planner", planner_node)
graph.add_node("executor", executor_node)
graph.add_node("checker", checker_node)
 
# 4. Add Edges
graph.add_edge("planner", "executor")
graph.add_edge("executor", "checker")
 
# 5. Add Conditional Edge
graph.add_conditional_edges(
    "checker",
    should_replan,  # Function returning next node name
    {
        "replan": "planner",
        "done": END
    }
)
 
# 6. Compile and Run
app = graph.compile()
result = app.invoke({"query": "...", "plan": [], "results": [], "done": False})

def planner_node(state: AgentState) -> AgentState:
    """Generate execution plan"""
    plan = create_plan(state["query"])
    return {"plan": plan}
 
def executor_node(state: AgentState) -> AgentState:
    """Execute next step in plan"""
    results = execute_step(state["plan"][0])
    remaining_plan = state["plan"][1:]
    return {"plan": remaining_plan, "results": state["results"] + [results]}
 
def checker_node(state: AgentState) -> AgentState:
    """Check if we're done"""
    done = len(state["plan"]) == 0
    return {"done": done}
 
def should_replan(state: AgentState) -> str:
    if state["done"]:
        return "done"
    return "continue"